From 2e6136caece098aeb37bd3454fa9c35fbdfea32d Mon Sep 17 00:00:00 2001 From: Slipstream Date: Wed, 21 May 2025 18:20:33 -0600 Subject: [PATCH] refactor: Improve custom bot management by adding import error handling and updating bot initialization logic --- api_service/dashboard_api_endpoints.py | 29 ++++++++- main.py | 89 ++++++++++++++------------ 2 files changed, 76 insertions(+), 42 deletions(-) diff --git a/api_service/dashboard_api_endpoints.py b/api_service/dashboard_api_endpoints.py index e31d894..1f480e7 100644 --- a/api_service/dashboard_api_endpoints.py +++ b/api_service/dashboard_api_endpoints.py @@ -36,7 +36,11 @@ import os parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..')) if parent_dir not in sys.path: sys.path.append(parent_dir) -import custom_bot_manager +try: + import custom_bot_manager +except ImportError: + print("Warning: Could not import custom_bot_manager. Custom bot functionality will be disabled.") + custom_bot_manager = None # Set up logging log = logging.getLogger(__name__) @@ -1113,6 +1117,15 @@ async def get_custom_bot_status( ): """Get the status of the user's custom bot.""" try: + # Check if custom bot manager is available + if not custom_bot_manager: + return CustomBotStatus( + exists=False, + status="not_available", + error="Custom bot functionality is not available", + is_running=False + ) + user_id = _user.get('user_id') if not user_id: raise HTTPException( @@ -1136,6 +1149,13 @@ async def start_custom_bot( ): """Start the user's custom bot.""" try: + # Check if custom bot manager is available + if not custom_bot_manager: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="Custom bot functionality is not available" + ) + user_id = _user.get('user_id') if not user_id: raise HTTPException( @@ -1220,6 +1240,13 @@ async def stop_custom_bot( ): """Stop the user's custom bot.""" try: + # Check if custom bot manager is available + if not custom_bot_manager: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="Custom bot functionality is not available" + ) + user_id = _user.get('user_id') if not user_id: raise HTTPException( diff --git a/main.py b/main.py index abd12c3..1183910 100644 --- a/main.py +++ b/main.py @@ -256,53 +256,55 @@ async def on_ready(): # Start custom bots for users who have enabled them try: log.info("Starting custom bots for users...") - if bot.pg_pool: - async with bot.pg_pool.acquire() as conn: - # Get all users with custom bots enabled - users = await conn.fetch(""" - SELECT user_id, settings FROM user_settings - WHERE settings->>'custom_bot_enabled' = 'true' - AND settings->>'custom_bot_token' IS NOT NULL - """) - log.info(f"Found {len(users)} users with custom bots enabled") + # Import the API database to access user settings + try: + from api_service.api_server import db + if not db: + log.warning("API database not initialized, cannot start custom bots") + return + except ImportError: + log.warning("Could not import API database, cannot start custom bots") + return - for user in users: - user_id = user['user_id'] - settings = user['settings'] + # Get all user settings from the API database + user_settings_dict = db.user_settings + enabled_bots = 0 - # Extract bot settings - token = settings.get('custom_bot_token') - prefix = settings.get('custom_bot_prefix', '!') - status_type = settings.get('custom_bot_status_type', 'listening') - status_text = settings.get('custom_bot_status_text', '!help') + for user_id, settings in user_settings_dict.items(): + # Check if custom bot is enabled and has a token + if settings.custom_bot_enabled and settings.custom_bot_token: + enabled_bots += 1 + token = settings.custom_bot_token + prefix = settings.custom_bot_prefix + status_type = settings.custom_bot_status_type + status_text = settings.custom_bot_status_text - if token: - log.info(f"Creating and starting custom bot for user {user_id}") - # Create the bot - success, message = await custom_bot_manager.create_custom_bot( - user_id=user_id, - token=token, - prefix=prefix, - status_type=status_type, - status_text=status_text - ) + log.info(f"Creating and starting custom bot for user {user_id}") + # Create the bot + success, message = await custom_bot_manager.create_custom_bot( + user_id=user_id, + token=token, + prefix=prefix, + status_type=status_type, + status_text=status_text + ) - if success: - # Start the bot - success, message = custom_bot_manager.run_custom_bot_in_thread( - user_id=user_id, - token=token - ) + if success: + # Start the bot + success, message = custom_bot_manager.run_custom_bot_in_thread( + user_id=user_id, + token=token + ) - if success: - log.info(f"Successfully started custom bot for user {user_id}") - else: - log.error(f"Failed to start custom bot for user {user_id}: {message}") - else: - log.error(f"Failed to create custom bot for user {user_id}: {message}") - else: - log.warning("Bot Postgres pool not initialized, cannot start custom bots") + if success: + log.info(f"Successfully started custom bot for user {user_id}") + else: + log.error(f"Failed to start custom bot for user {user_id}: {message}") + else: + log.error(f"Failed to create custom bot for user {user_id}: {message}") + + log.info(f"Found {enabled_bots} users with custom bots enabled") except Exception as e: log.exception(f"Error starting custom bots: {e}") @@ -656,14 +658,19 @@ async def main(args): # Pass parsed args log.info("Stopping all custom bots...") # Get all running custom bots bot_statuses = custom_bot_manager.get_all_custom_bot_statuses() + stopped_count = 0 + for user_id, status in bot_statuses.items(): if status.get('is_running', False): log.info(f"Stopping custom bot for user {user_id}") success, message = custom_bot_manager.stop_custom_bot(user_id) if success: + stopped_count += 1 log.info(f"Successfully stopped custom bot for user {user_id}") else: log.error(f"Failed to stop custom bot for user {user_id}: {message}") + + log.info(f"Stopped {stopped_count} custom bots") except Exception as e: log.exception(f"Error stopping custom bots: {e}")