Skip to content

Instantly share code, notes, and snippets.

@DuskyElf
Last active September 20, 2024 03:57
Show Gist options
  • Save DuskyElf/1157f3147dd8fc43e135c0b0541e9e27 to your computer and use it in GitHub Desktop.
Save DuskyElf/1157f3147dd8fc43e135c0b0541e9e27 to your computer and use it in GitHub Desktop.
ollama bot integration into discord bot
import os
import discord
import dotenv
import requests
from dataclasses import dataclass
from typing import Optional, Protocol
dotenv.load_dotenv()
intents = discord.Intents.default()
intents.message_content = True
client = discord.Client(intents=intents)
command_tree = discord.app_commands.CommandTree(client)
class ChatBot(Protocol):
def generate_response(self, prompt: str) -> str:
...
class DolphinAlex:
REQUEST_TEMPLATE = {
"model": "dolphin-alex",
"stream": False,
"prompt": "",
}
def __init__(self, ollama_url: str):
self.url = ollama_url
def generate_response(self, prompt: str) -> str:
ollama_request = DolphinAlex.REQUEST_TEMPLATE.copy()
ollama_request["prompt"] = prompt
ollama_response = requests.post(self.url, json=ollama_request)
chat_result = ollama_response.json()["response"].strip()
return chat_result
@dataclass
class BotSettings:
reply_to_all: bool = False
# replace this with any ChatBot
chat_bot: ChatBot = DolphinAlex("http://localhost:11434/api/generate")
@client.event
async def on_ready():
await command_tree.sync()
await client.change_presence(activity=discord.Activity(name="with God, DuskyElf <3", type=1))
print(f'We have logged in as {client.user}')
@client.event
async def on_message(message: discord.Message):
if message.author == client.user:
return
# print(BotSettings.reply_to_all)
if BotSettings.reply_to_all:
async with message.channel.typing():
response = BotSettings.chat_bot.generate_response(message.content)
await message.reply(response)
@command_tree.command(name="chat", description="Chat with the bot")
async def chat(interaction: discord.Interaction, prompt: str):
await interaction.response.defer()
response = BotSettings.chat_bot.generate_response(prompt)
response = f"prompt: {prompt}\n\n{response}"
if len(response) > 2000:
embed = discord.Embed(title="Response", description=response, color=0xf1c40f)
await interaction.followup.send(embed=embed)
else:
await interaction.followup.send(response)
@command_tree.command(name="public", description="Bot will respond to all messages")
async def public(interaction: discord.Interaction):
if BotSettings.reply_to_all:
await interaction.response.send_message("Bot is already in Public mode")
return
BotSettings.reply_to_all = True
await interaction.response.send_message("Bot will respond to all messages")
@command_tree.command(name="private", description="Bot will not respond to all messages")
async def private(interaction: discord.Interaction):
if not BotSettings.reply_to_all:
await interaction.response.send_message("Bot is already in Private mode")
return
BotSettings.reply_to_all = False
await interaction.response.send_message("Bot will not respond to all messages")
if (token := os.getenv("TOKEN")):
client.run(token)
else:
print("Provide me the discord token")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment