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,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:
+        # Get all user settings from the API database
-                    user_id = user['user_id']
+        user_settings_dict = db.user_settings
-                    settings = user['settings']
+        enabled_bots = 0
 
-                    # Extract bot settings
+        for user_id, settings in user_settings_dict.items():
-                    token = settings.get('custom_bot_token')
+            # Check if custom bot is enabled and has a token
-                    prefix = settings.get('custom_bot_prefix', '!')
+            if settings.custom_bot_enabled and settings.custom_bot_token:
-                    status_type = settings.get('custom_bot_status_type', 'listening')
+                enabled_bots += 1
-                    status_text = settings.get('custom_bot_status_text', '!help')
+                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}")
-                        log.info(f"Creating and starting custom bot for user {user_id}")
+                # Create the bot
-                        # Create the bot
+                success, message = await custom_bot_manager.create_custom_bot(
-                        success, message = await custom_bot_manager.create_custom_bot(
+                    user_id=user_id,
-                            user_id=user_id,
+                    token=token,
-                            token=token,
+                    prefix=prefix,
-                            prefix=prefix,
+                    status_type=status_type,
-                            status_type=status_type,
+                    status_text=status_text
-                            status_text=status_text
+                )
-                        )
 
-                        if success:
+                if success:
-                            # Start the bot
+                    # Start the bot
-                            success, message = custom_bot_manager.run_custom_bot_in_thread(
+                    success, message = custom_bot_manager.run_custom_bot_in_thread(
-                                user_id=user_id,
+                        user_id=user_id,
-                                token=token
+                        token=token
-                            )
+                    )
 
-                            if success:
+                    if success:
-                                log.info(f"Successfully started custom bot for user {user_id}")
+                        log.info(f"Successfully started custom bot for user {user_id}")
-                            else:
+                    else:
-                                log.error(f"Failed to start custom bot for user {user_id}: {message}")
+                        log.error(f"Failed to start custom bot for user {user_id}: {message}")
-                        else:
+                else:
-                            log.error(f"Failed to create custom bot for user {user_id}: {message}")
+                    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}")