update the discord main bot to allow it to be packaged with Docker

This commit is contained in:
Josh Knapp 2025-01-02 19:41:08 -08:00
parent 45249174ba
commit 8144971707

View File

@ -1,9 +1,10 @@
import os import os
import discord import discord
from discord.ext import commands from discord.ext import commands
import openai
from openai import OpenAI from openai import OpenAI
import base64
import requests
from io import BytesIO
from collections import deque from collections import deque
from dotenv import load_dotenv from dotenv import load_dotenv
@ -35,28 +36,55 @@ bot = commands.Bot(command_prefix='!', intents=intents)
# Message history cache # Message history cache
channel_history = {} channel_history = {}
async def download_image(url):
response = requests.get(url)
if response.status_code == 200:
image_data = BytesIO(response.content)
base64_image = base64.b64encode(image_data.read()).decode('utf-8')
return base64_image
return None
async def get_chat_history(channel, limit=100): async def get_chat_history(channel, limit=100):
messages = [] messages = []
async for message in channel.history(limit=limit): async for message in channel.history(limit=limit):
messages.append(f"{message.author.name}: {message.content}") content = f"{message.author.name}: {message.content}"
# Handle attachments (images)
for attachment in message.attachments:
if any(attachment.filename.lower().endswith(ext) for ext in ['.png', '.jpg', '.jpeg', '.gif', '.webp']):
content += f" [Image: {attachment.url}]"
messages.append(content)
return "\n".join(reversed(messages)) return "\n".join(reversed(messages))
async def get_ai_response(context, user_message): async def get_ai_response(context, user_message, image_urls=None):
formatted_prompt = f"##CONTEXT##\n{context}\n##ENDCONTEXT##\n\n{user_message}" messages = [{"role": "user", "content": []}]
# Add text content
text_content = f"##CONTEXT##\n{context}\n##ENDCONTEXT##\n\n{user_message}"
messages[0]["content"].append({"type": "text", "text": text_content})
# Add image content if present
if image_urls:
for url in image_urls:
base64_image = await download_image(url)
if base64_image:
messages[0]["content"].append({
"type": "image_url",
"image_url": {
"url": f"data:image/jpeg;base64,{base64_image}"
}
})
try: try:
response = client.chat.completions.create(model=MODEL_NAME, response = client.chat.completions.create(
messages=[ model=MODEL_NAME,
{"role": "user", "content": formatted_prompt} messages=messages
]) )
return response.choices[0].message.content return response.choices[0].message.content
except Exception as e: except Exception as e:
return f"Error: {str(e)}" return f"Error: {str(e)}"
@bot.event
async def on_ready():
print(f'{bot.user} has connected to Discord!')
@bot.event @bot.event
async def on_message(message): async def on_message(message):
# Ignore messages from the bot itself # Ignore messages from the bot itself
@ -81,14 +109,25 @@ async def on_message(message):
# Remove bot mention from the message # Remove bot mention from the message
user_message = message.content.replace(f'<@{bot.user.id}>', '').strip() user_message = message.content.replace(f'<@{bot.user.id}>', '').strip()
# Collect image URLs from the message
image_urls = []
for attachment in message.attachments:
if any(attachment.filename.lower().endswith(ext) for ext in ['.png', '.jpg', '.jpeg', '.gif', '.webp']):
image_urls.append(attachment.url)
# Get AI response # Get AI response
response = await get_ai_response(history, user_message) response = await get_ai_response(history, user_message, image_urls)
# Send response # Send response
await message.reply(response) await message.reply(response)
await bot.process_commands(message) await bot.process_commands(message)
@bot.event
async def on_ready():
print(f'{bot.user} has connected to Discord!')
def main(): def main():
if not all([DISCORD_TOKEN, OPENAI_API_KEY, OPENWEBUI_API_BASE, MODEL_NAME]): if not all([DISCORD_TOKEN, OPENAI_API_KEY, OPENWEBUI_API_BASE, MODEL_NAME]):
print("Error: Missing required environment variables") print("Error: Missing required environment variables")