OpenWebUI-Discordbot/scripts/discordbot.py

import os
import discord
from discord.ext import commands
from openai import OpenAI
import base64
import requests
from io import BytesIO
from collections import deque
from dotenv import load_dotenv
import json
import datetime
import aiohttp
from typing import Dict, Any, List

# Load environment variables
load_dotenv()

# Get environment variables
DISCORD_TOKEN = os.getenv('DISCORD_TOKEN')
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
OPENWEBUI_API_BASE = os.getenv('OPENWEBUI_API_BASE')
MODEL_NAME = os.getenv('MODEL_NAME')

# Configure OpenAI client to point to OpenWebUI
client = OpenAI(
    api_key=os.getenv('OPENAI_API_KEY'),
    base_url=os.getenv('OPENWEBUI_API_BASE')  # e.g., "http://localhost:8080/v1"
)

# Configure OpenAI
# TODO: The 'openai.api_base' option isn't read in the client API. You will need to pass it when you instantiate the client, e.g. 'OpenAI(base_url=OPENWEBUI_API_BASE)'
# openai.api_base = OPENWEBUI_API_BASE

# Initialize Discord bot
intents = discord.Intents.default()
intents.message_content = True
intents.messages = True
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):
        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_available_tools() -> List[Dict[str, Any]]:
    """Fetch available tools from OpenWebUI API."""
    try:
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {OPENAI_API_KEY}"
        }
        
        async with aiohttp.ClientSession() as session:
            async with session.get(
                f"{OPENWEBUI_API_BASE}/v1/tools/list",
                headers=headers
            ) as response:
                if response.status == 200:
                    tools = await response.json()
                    return tools
                else:
                    print(f"Error fetching tools: {await response.text()}")
                    return []
    except Exception as e:
        print(f"Error fetching tools: {str(e)}")
        return []

async def get_ai_response(context, user_message, image_urls=None):
    # Fetch available tools
    tools = await get_available_tools()
    tools_json = json.dumps(tools, indent=2)
    
    system_message = f"\"\"\"Previous conversation context:{context}\nAvailable Tools: {tools_json}\nReturn an empty string if no tools match the query." + """If a function tool matches, construct and return a JSON object in the format {"name": "functionName", "parameters": {"requiredFunctionParamKey": "requiredFunctionParamValue\"}} using the appropriate tool and its parameters. Only return the object and limit the response to the JSON object without additional text."""

    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": [] if image_urls else user_message}
    ]
    
    # Handle messages with images differently
    if image_urls:
        content_parts = [{"type": "text", "text": user_message}]
        
        for url in image_urls:
            base64_image = await download_image(url)
            if base64_image:
                content_parts.append({
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/jpeg;base64,{base64_image}"
                    }
                })
        messages[1]["content"] = content_parts

    try:
        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_message(message):
    # Ignore messages from the bot itself
    if message.author == bot.user:
        return

    should_respond = False

    # Check if bot was mentioned
    if bot.user in message.mentions:
        should_respond = True

    # Check if message is a DM
    if isinstance(message.channel, discord.DMChannel):
        should_respond = True

    if should_respond:
        async with message.channel.typing():
            # Get chat history
            history = await get_chat_history(message.channel)

            # 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, 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")
        return

    bot.run(DISCORD_TOKEN)

if __name__ == "__main__":
    main()