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
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()