ds_zagres_ai/ai_service/bot_manager.py

"""
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})...")

                    # Define a monitoring task to restart the bot if it fails
                    async def monitor_bot_task():
                        global bot_task
                        try:
                            # Start the main bot task
                            main_task = asyncio.create_task(custom_bot_main())

                            # Wait for the task to complete or fail
                            try:
                                await main_task
                                logger.warning("Bot task completed unexpectedly")
                            except asyncio.CancelledError:
                                logger.info("Bot task was cancelled")
                                raise
                            except Exception as e:
                                logger.error(f"Bot task failed with error: {str(e)}")
                                logger.error(traceback.format_exc())

                                # Try to restart the bot
                                logger.info("Attempting to restart the bot...")
                                await asyncio.sleep(5)  # Wait a bit before restarting
                                restart_task = asyncio.create_task(custom_bot_main())
                                bot_task = restart_task  # Update the global task reference
                                logger.info("Bot restarted successfully")
                        except Exception as e:
                            logger.error(f"Error in monitor task: {str(e)}")
                            logger.error(traceback.format_exc())

                    # Start the monitoring task
                    bot_task = asyncio.create_task(monitor_bot_task())
                    logger.info("Bot started successfully with monitoring!")
                    return True
                except Exception as e:
                    last_error = e
                    logger.error(f"Error starting bot (attempt {attempt}/{max_retries}): {str(e)}")
                    logger.error(traceback.format_exc())
                    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()