diff --git a/scripts/discordbot.py b/scripts/discordbot.py index 9edf967..cf24c9e 100644 --- a/scripts/discordbot.py +++ b/scripts/discordbot.py @@ -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")