refactor: Improve custom bot management by adding import error handling and updating bot initialization logic

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(

main.py

@@ -256,28 +256,30 @@ 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(
@@ -301,8 +303,8 @@ async def on_ready():
                         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")
+
+        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}")