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 discord
from discord.ext import commands
import openai
from openai import OpenAI
import base64
import requests
from io import BytesIO
from collections import deque
from dotenv import load_dotenv
@ -35,28 +36,55 @@ bot = commands.Bot(command_prefix='!', intents=intents)
# Message history cache
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):
messages = []
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))
async def get_ai_response(context, user_message):
formatted_prompt = f"##CONTEXT##\n{context}\n##ENDCONTEXT##\n\n{user_message}"
async def get_ai_response(context, user_message, image_urls=None):
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:
response = client.chat.completions.create(model=MODEL_NAME,
messages=[
{"role": "user", "content": formatted_prompt}
])
response = client.chat.completions.create(
model=MODEL_NAME,
messages=messages
)
return response.choices[0].message.content
except Exception as e:
return f"Error: {str(e)}"
@bot.event
async def on_ready():
print(f'{bot.user} has connected to Discord!')
@bot.event
async def on_message(message):
# Ignore messages from the bot itself
@ -81,14 +109,25 @@ async def on_message(message):
# Remove bot mention from the message
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
response = await get_ai_response(history, user_message)
response = await get_ai_response(history, user_message, image_urls)
# Send response
await message.reply(response)
await bot.process_commands(message)
@bot.event
async def on_ready():
print(f'{bot.user} has connected to Discord!')
def main():
if not all([DISCORD_TOKEN, OPENAI_API_KEY, OPENWEBUI_API_BASE, MODEL_NAME]):
print("Error: Missing required environment variables")