From 4df065b5bc5223bc340932dd7c551061948fae1f Mon Sep 17 00:00:00 2001 From: Amish Suchak Date: Tue, 31 Dec 2024 18:05:14 -0600 Subject: [PATCH] feat: integrate OpenRouter API and improve message handling - Added OpenRouter API integration for AI responses - Added debug logging for message handling - Improved AI suffix detection logic --- bot/env.py | 2 + bot/main.py | 132 +++++++++++++++++++++++++++++++++++++++++++--------- 2 files changed, 111 insertions(+), 23 deletions(-) diff --git a/bot/env.py b/bot/env.py index 4cc234d..13fd85f 100644 --- a/bot/env.py +++ b/bot/env.py @@ -10,3 +10,5 @@ except ImportError: WEBUI_URL = os.getenv("WEBUI_URL", "http://localhost:8080") TOKEN = os.getenv("TOKEN", "") +OPENROUTER_API_KEY = "" +OPENROUTER_API_URL = "https://openrouter.ai/api/v1/chat/completions" diff --git a/bot/main.py b/bot/main.py index ed1d5e3..6c86c07 100644 --- a/bot/main.py +++ b/bot/main.py @@ -1,8 +1,51 @@ import asyncio import socketio -from env import WEBUI_URL, TOKEN +import aiohttp +import json +from env import WEBUI_URL, TOKEN, OPENROUTER_API_KEY, OPENROUTER_API_URL from utils import send_message, send_typing +# Create a single session for all API calls +session = None + +async def get_session(): + global session + if session is None: + session = aiohttp.ClientSession() + return session + +async def cleanup(): + global session + if session: + await session.close() + session = None + +async def call_openrouter(message: str) -> str: + headers = { + "Authorization": f"Bearer {OPENROUTER_API_KEY}", + "HTTP-Referer": "http://localhost:8080", + "Content-Type": "application/json" + } + + payload = { + "model": "mistralai/mistral-7b-instruct", + "messages": [{"role": "user", "content": message}] + } + + try: + session = await get_session() + async with session.post(OPENROUTER_API_URL, headers=headers, json=payload) as response: + if response.status == 200: + data = await response.json() + return data["choices"][0]["message"]["content"] + else: + error_text = await response.text() + print(f"OpenRouter API error: {error_text}") + return f"Error: Unable to get response from AI (Status {response.status})" + except Exception as e: + print(f"Error calling OpenRouter: {e}") + return "Error: Unable to connect to AI service" + # Create an asynchronous Socket.IO client instance sio = socketio.AsyncClient(logger=False, engineio_logger=False) @@ -19,18 +62,57 @@ async def disconnect(): # Define a function to handle channel events -def events(user_id): - @sio.on("channel-events") - async def channel_events(data): - if data["user"]["id"] == user_id: - # Ignore events from the bot itself - return +# Keep track of processed message IDs to avoid duplicates +processed_messages = set() +@sio.on("channel-events") +async def on_channel_event(data): + try: + print(f"Received event: {data}") + + if not isinstance(data, dict) or "data" not in data: + return + if data["data"]["type"] == "message": - print(f'{data["user"]["name"]}: {data["data"]["data"]["content"]}') - await send_typing(sio, data["channel_id"]) - await asyncio.sleep(1) # Simulate a delay - await send_message(data["channel_id"], "Pong!") + message_id = data["data"]["data"].get("id") + if message_id in processed_messages: + print(f"Skipping already processed message: {message_id}") + return + + # Get the message content and channel ID + message_content = data["data"]["data"]["content"] + channel_id = data["channel_id"] + + # Debug logging + print(f"Received message: {message_content}") + print(f"From user ID: {data['user']['id']}") + print(f"Bot ID: {sio.user_id}") + print(f"Ends with AI: {message_content.strip().upper().endswith('AI')}") + + # Only respond to messages ending with AI (case-insensitive) + if (message_content and channel_id and + "user" in data and + data["user"]["id"] != sio.user_id and + message_content.strip().upper().endswith("AI")): + + print(f'Processing AI request from {data["user"]["name"]}: {message_content}') + print(f'Responding in channel: {channel_id}') + + # Remove the "AI" suffix and any trailing whitespace + query = message_content.strip()[:-2].strip() + if query: # Only proceed if there's a message after removing "AI" + await send_typing(sio, channel_id) + try: + response = await call_openrouter(query) + await send_message(channel_id, response) + except Exception as e: + print(f"Error calling OpenRouter API: {e}") + await send_message(channel_id, "Sorry, I'm unable to get a response from the AI service right now. Please try again later.") + + if message_id: + processed_messages.add(message_id) + except Exception as e: + print(f"Error handling event: {e}") # Define an async function for the main workflow @@ -41,21 +123,25 @@ async def main(): WEBUI_URL, socketio_path="/ws/socket.io", transports=["websocket"] ) print("Connection established!") + + # Authenticate with the server + print("Authenticating with server...") + response = await sio.call("user-join", {"auth": {"token": TOKEN}}) + sio.user_id = response["id"] # Store user ID for later use + print(f"Authentication successful. User ID: {sio.user_id}") + + # Wait indefinitely to keep the connection open + await sio.wait() except Exception as e: - print(f"Failed to connect: {e}") + print(f"Failed to connect/authenticate: {e}") return - # Callback function for user-join - async def join_callback(data): - events(data["id"]) # Attach the event handlers dynamically - - # Authenticate with the server - await sio.emit("user-join", {"auth": {"token": TOKEN}}, callback=join_callback) - - # Wait indefinitely to keep the connection open - await sio.wait() - # Actually run the async `main` function using `asyncio` if __name__ == "__main__": - asyncio.run(main()) + try: + asyncio.run(main()) + finally: + # Ensure we clean up the session + if session: + asyncio.run(cleanup())