Iyeoluwa Akinrinola
285 lines
10 KiB
Python
285 lines
10 KiB
Python
"""
|
|
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()
|