diff --git a/src/nalgonda/agency_config_lock_manager.py b/src/nalgonda/agency_config_lock_manager.py index a469083b..7efa4487 100644 --- a/src/nalgonda/agency_config_lock_manager.py +++ b/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""" diff --git a/src/nalgonda/data/default_config.json b/src/nalgonda/data/default_config.json index 6c41effb..f38a48b5 100644 --- a/src/nalgonda/data/default_config.json +++ b/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": [] }, diff --git a/src/nalgonda/main.py b/src/nalgonda/main.py index 484b2439..2b9aae4b 100644 --- a/src/nalgonda/main.py +++ b/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