Integrate OpenWebUI bot with AI service

ai_service/bot_manager.py

@@ -0,0 +1,170 @@
+"""
+Bot manager for the AI service.
+
+This module provides functionality to manage the OpenWebUI bot.
+"""
+
+import os
+import sys
+import asyncio
+import logging
+from typing import Optional
+
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+# Add the openwebui-bot directory to the Python path
+# Use a path relative to the current file's directory to ensure it works in any environment
+current_dir = os.path.dirname(os.path.abspath(__file__))
+root_dir = os.path.dirname(current_dir)
+openwebui_bot_dir = os.path.join(root_dir, 'openwebui-bot')
+sys.path.insert(0, openwebui_bot_dir)
+
+# Log the path for debugging
+logger.info(f"Adding OpenWebUI bot directory to Python path: {openwebui_bot_dir}")
+
+# Import the bot modules
+try:
+    from env import WEBUI_URL, TOKEN
+    from examples.custom_ai import main as custom_bot_main
+except ImportError as e:
+    print(f"Error importing OpenWebUI bot modules: {str(e)}")
+    print("Make sure the openwebui-bot directory exists and contains the required files.")
+    WEBUI_URL = None
+    TOKEN = None
+    custom_bot_main = None
+
+# Global variable to store the bot task
+bot_task = None
+
+async def start_bot(
+    openwebui_url: str,
+    api_key: str,
+    model_id: str,
+    system_prompt: str = None,
+    temperature: float = None,
+    max_tokens: int = None,
+    top_p: float = None,
+    triggers: list = None,
+    respond_to_all: bool = None
+) -> bool:
+    """
+    Start the OpenWebUI bot.
+
+    Args:
+        openwebui_url: URL of the OpenWebUI instance.
+        api_key: API key for authentication.
+        model_id: ID of the model to use.
+        system_prompt: System prompt for the bot.
+        temperature: Temperature for response generation.
+        max_tokens: Maximum number of tokens to generate.
+        top_p: Top-p sampling parameter.
+        triggers: List of trigger words that will make the bot respond.
+        respond_to_all: Whether to respond to all messages or only to mentions.
+
+    Returns:
+        True if the bot was started successfully, False otherwise.
+    """
+    global bot_task
+
+    # Check if the bot is already running
+    if bot_task is not None and not bot_task.done():
+        logger.warning("Bot is already running!")
+        return True
+
+    # Check if the bot modules were imported successfully
+    if custom_bot_main is None:
+        logger.error("OpenWebUI bot modules not found. Bot cannot be started.")
+        return False
+
+    # Set environment variables for the bot
+    os.environ["WEBUI_URL"] = openwebui_url
+    os.environ["TOKEN"] = api_key
+    os.environ["MODEL_ID"] = model_id
+
+    # Set optional parameters if provided
+    if system_prompt:
+        os.environ["SYSTEM_PROMPT"] = system_prompt
+    if temperature is not None:
+        os.environ["TEMPERATURE"] = str(temperature)
+    if max_tokens is not None:
+        os.environ["MAX_TOKENS"] = str(max_tokens)
+    if top_p is not None:
+        os.environ["TOP_P"] = str(top_p)
+    if triggers:
+        os.environ["TRIGGERS"] = ",".join(triggers)
+    if respond_to_all is not None:
+        os.environ["RESPOND_TO_ALL"] = str(respond_to_all).lower()
+
+    try:
+        # Start the bot in a background task
+        logger.info("Starting OpenWebUI bot...")
+        logger.info(f"OpenWebUI URL: {openwebui_url}")
+        logger.info(f"Model: {model_id}")
+
+        if system_prompt:
+            logger.info(f"System prompt: {system_prompt[:50]}...")
+        if temperature is not None:
+            logger.info(f"Temperature: {temperature}")
+        if max_tokens is not None:
+            logger.info(f"Max tokens: {max_tokens}")
+        if top_p is not None:
+            logger.info(f"Top-p: {top_p}")
+        if triggers:
+            logger.info(f"Triggers: {triggers}")
+        if respond_to_all is not None:
+            logger.info(f"Respond to all: {respond_to_all}")
+
+        # Create a task for the bot
+        bot_task = asyncio.create_task(custom_bot_main())
+
+        logger.info("Bot started successfully!")
+        return True
+    except Exception as e:
+        logger.error(f"Error starting bot: {str(e)}")
+        return False
+
+async def stop_bot() -> bool:
+    """
+    Stop the OpenWebUI bot.
+
+    Returns:
+        True if the bot was stopped successfully, False otherwise.
+    """
+    global bot_task
+
+    # Check if the bot is running
+    if bot_task is None or bot_task.done():
+        logger.warning("Bot is not running!")
+        return True
+
+    try:
+        # Cancel the bot task
+        logger.info("Stopping OpenWebUI bot...")
+        bot_task.cancel()
+
+        try:
+            # Wait for the task to be cancelled
+            await bot_task
+        except asyncio.CancelledError:
+            pass
+
+        logger.info("Bot stopped successfully!")
+        return True
+    except Exception as e:
+        logger.error(f"Error stopping bot: {str(e)}")
+        return False
+
+def is_bot_running() -> bool:
+    """
+    Check if the bot is running.
+
+    Returns:
+        True if the bot is running, False otherwise.
+    """
+    global bot_task
+    return bot_task is not None and not bot_task.done()