""" Bot manager for the AI service. This module provides functionality to manage the OpenWebUI bot. """ import os import sys import asyncio import logging import traceback 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 custom_bot_main = None # Check if the openwebui-bot directory exists if not os.path.exists(openwebui_bot_dir): logger.error(f"OpenWebUI bot directory not found: {openwebui_bot_dir}") else: logger.info(f"OpenWebUI bot directory found: {openwebui_bot_dir}") # Add examples directory to path examples_dir = os.path.join(openwebui_bot_dir, 'examples') if os.path.exists(examples_dir): sys.path.insert(0, examples_dir) logger.info(f"Added examples directory to path: {examples_dir}") # Check for required files env_path = os.path.join(openwebui_bot_dir, 'env.py') custom_ai_path = os.path.join(openwebui_bot_dir, 'examples', 'custom_ai.py') if not os.path.exists(env_path): logger.error(f"env.py not found at {env_path}") else: logger.info(f"env.py found at {env_path}") if not os.path.exists(custom_ai_path): logger.error(f"custom_ai.py not found at {custom_ai_path}") else: logger.info(f"custom_ai.py found at {custom_ai_path}") # Try to import the modules try: # Use a different approach to import the modules sys.path.insert(0, openwebui_bot_dir) # Import the custom_ai module try: from examples.custom_ai import main as custom_bot_main logger.info("Successfully imported custom_ai module") except ImportError as e: logger.error(f"Error importing custom_ai module: {str(e)}") # Try a different approach try: import importlib.util spec = importlib.util.spec_from_file_location("custom_ai", custom_ai_path) custom_ai = importlib.util.module_from_spec(spec) spec.loader.exec_module(custom_ai) custom_bot_main = custom_ai.main logger.info("Successfully imported custom_ai module using importlib") except Exception as e: logger.error(f"Error importing custom_ai module using importlib: {str(e)}") except Exception as e: logger.error(f"Error importing OpenWebUI bot modules: {str(e)}") logger.error(f"Traceback: {traceback.format_exc()}") # Log the current state if custom_bot_main is None: logger.error("Failed to import custom_bot_main function") else: logger.info("Successfully imported custom_bot_main function") # 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. This function creates a .env file for the bot and starts it in a background task. The bot connects to OpenWebUI via WebSocket and listens for messages in channels. When a message mentions the bot (using trigger words like @ai), the bot processes the message and sends a response back to the channel. 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 # Create .env file for the bot try: env_file_path = os.path.join(openwebui_bot_dir, '.env') with open(env_file_path, 'w') as f: f.write(f"# OpenWebUI configuration\n") f.write(f"WEBUI_URL={openwebui_url}\n") f.write(f"TOKEN={api_key}\n\n") f.write(f"# Model configuration\n") f.write(f"MODEL_ID={model_id}\n\n") f.write(f"# Bot behavior configuration\n") f.write(f'SYSTEM_PROMPT="{system_prompt or "You are a helpful AI assistant."}"\n') f.write(f"TEMPERATURE={temperature or 0.7}\n") f.write(f"MAX_TOKENS={max_tokens or 2048}\n") f.write(f"TOP_P={top_p or 0.9}\n") f.write(f"TRIGGERS={','.join(triggers) if triggers else '@ai,@bot,@assistant,@chatbot'}\n") f.write(f"RESPOND_TO_ALL={'true' if respond_to_all else 'false'}\n") logger.info(f"Created .env file at {env_file_path}") except Exception as e: logger.error(f"Error creating .env file: {str(e)}") # Continue anyway, as we'll set environment variables too # 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"API Key: {api_key[:4]}...{api_key[-4:] if len(api_key) > 8 else ''}") 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 try: # Try to reload the env module to pick up new environment variables try: import importlib if 'env' in sys.modules: importlib.reload(sys.modules['env']) logger.info("Reloaded env module") except Exception as e: logger.warning(f"Failed to reload env module: {str(e)}") # Start the bot with a retry mechanism max_retries = 3 retry_delay = 5 # seconds last_error = None for attempt in range(1, max_retries + 1): try: logger.info(f"Starting bot (attempt {attempt}/{max_retries})...") bot_task = asyncio.create_task(custom_bot_main()) logger.info("Bot started successfully!") return True except Exception as e: last_error = e logger.error(f"Error starting bot (attempt {attempt}/{max_retries}): {str(e)}") if attempt < max_retries: logger.info(f"Retrying in {retry_delay} seconds...") await asyncio.sleep(retry_delay) else: logger.error("Maximum retry attempts reached") # If we get here, all attempts failed logger.error(f"Failed to start bot after {max_retries} attempts") logger.error(f"Last error: {str(last_error)}") logger.error(f"Traceback: {traceback.format_exc()}") return False except Exception as e: logger.error(f"Error creating bot task: {str(e)}") logger.error(f"Traceback: {traceback.format_exc()}") return False except Exception as e: logger.error(f"Error starting bot: {str(e)}") logger.error(f"Traceback: {traceback.format_exc()}") 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()