Fix blocking send message

src/nalgonda/agency_config_lock_manager.py

@@ -1,6 +1,7 @@
 import threading
 from collections import defaultdict
 
+
 class AgencyConfigLockManager:
     """Lock manager for agency config files"""
 

src/nalgonda/data/default_config.json

@@ -5,7 +5,7 @@
       "id": null,
       "role": "CEO",
       "description": "Responsible for client communication, task planning and management.",
-      "instructions": "# Instructions for CEO Agent\n\n- Ensure that proposal is send to the user before proceeding with task execution.\n- Delegate tasks to appropriate agents, ensuring they align with their expertise and capabilities.\n- Clearly define the objectives and expected outcomes for each task.\n- Provide necessary context and background information for successful task completion.\n- Maintain ongoing communication with agents until complete task execution.\n- Review completed tasks to ensure they meet the set objectives.\n- Report the results back to the user.",
+      "instructions": "# Instructions for CEO Agent\n\n- Send the proposal to the user before beginning task execution.\n- Assign tasks to agents based on their expertise and capabilities.\n- Clearly outline the goals and expected outcomes for each task.\n- Provide essential context and background for successful task completion.\n- Keep in constant communication with agents throughout task execution.\n- Review completed tasks to ensure they meet the objectives.\n- Report the outcomes to the user.\n- Pass on any user feedback to the agents. Note: All conversations with agents are private. Information must be relayed directly by you, as cross-referencing or referencing 'above' is not possible in these separate, private conversations.",
       "files_folder": null,
       "tools": []
     },

src/nalgonda/main.py

@@ -3,6 +3,7 @@
 import uuid
 
 from agency_manager import AgencyManager
+from agency_swarm import Agency
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from nalgonda.constants import DATA_DIR
 
@@ -69,24 +70,41 @@ async def websocket_endpoint(websocket: WebSocket, agency_id: str):
 
     try:
         while True:
-            user_message = await websocket.receive_text()
+            try:
+                user_message = await websocket.receive_text()
 
-            if not user_message.strip():
-                await ws_manager.send_message("message not provided", websocket)
-                await ws_manager.disconnect(websocket)
-                await websocket.close(code=1003)
-                return
+                if not user_message.strip():
+                    await ws_manager.send_message("message not provided", websocket)
+                    continue
 
-            gen = await asyncio.to_thread(agency.get_completion, message=user_message, yield_messages=True)
-            for response in gen:
-                response_text = response.get_formatted_content()
-                await ws_manager.send_message(response_text, websocket)
+                await process_message(user_message, agency, websocket)
+
+            except Exception as e:
+                logger.exception(e)
+                await ws_manager.send_message(f"Error: {e}\nPlease try again.", websocket)
+                continue
 
     except WebSocketDisconnect:
         await ws_manager.disconnect(websocket)
         logger.info(f"WebSocket disconnected for agency_id: {agency_id}")
 
 
+async def process_message(user_message: str, agency: Agency, websocket: WebSocket):
+    """Process the user message and send the response to the websocket."""
+    gen = agency.get_completion(message=user_message, yield_messages=True)
+
+    async for response in async_gen(gen):
+        response_text = response.get_formatted_content()
+        await ws_manager.send_message(response_text, websocket)
+
+
+async def async_gen(gen):
+    """Asynchronous wrapper for a synchronous generator."""
+    for value in gen:
+        # Offload the blocking operation to a separate thread
+        yield await asyncio.to_thread(lambda v=value: v)
+
+
 if __name__ == "__main__":
     import uvicorn