feat: websocket calls multiple agents and returns messages

Author: zhayujie

agentmesh/api/websocket_api.py

@@ -1,9 +1,10 @@
 import uuid
 import json
+import threading
 from fastapi import APIRouter, WebSocket, WebSocketDisconnect
 from fastapi.responses import JSONResponse
 
-from ..service.websocket_service import websocket_manager, task_processor
+from ..service.websocket_service import websocket_manager, task_processor, thread_manager
 from ..common.models import UserInputMessage
 
 # Create router
@@ -21,19 +22,66 @@ async def websocket_endpoint(websocket: WebSocket):
 
     try:
         # Accept the WebSocket connection
-        await websocket_manager.connect(websocket, connection_id)
+        await websocket.accept()
+        
+        # Create a simple connection wrapper for FastAPI WebSocket
+        class FastAPIWebSocketWrapper:
+            def __init__(self, websocket: WebSocket):
+                self.websocket = websocket
+                self.connection_id = connection_id
+            
+            def send(self, message: str):
+                """Send message using FastAPI WebSocket"""
+                import asyncio
+                try:
+                    # Create a new event loop for this thread
+                    loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(loop)
+                    loop.run_until_complete(self.websocket.send_text(message))
+                except Exception as e:
+                    print(f"Error sending message: {e}")
+                finally:
+                    try:
+                        loop.close()
+                    except:
+                        pass
+        
+        # Create wrapper and register with manager
+        ws_wrapper = FastAPIWebSocketWrapper(websocket)
+        websocket_manager.connect(ws_wrapper, connection_id)
         print(f"WebSocket client connected: {connection_id}")
 
         # Handle messages from client
         while True:
             try:
+                # Check if shutdown is requested
+                if thread_manager.shutdown_event.is_set():
+                    print(f"Shutdown requested, closing connection {connection_id}")
+                    break
+                
                 # Receive message from client
                 data = await websocket.receive_text()
+                
+                # Process message
                 message_data = json.loads(data)
 
                 # Handle different message types
                 if message_data.get("event") == "user_input":
-                    await handle_user_input(connection_id, message_data)
+                    # Start task execution in a separate thread with proper tracking
+                    task_thread = threading.Thread(
+                        target=handle_user_input,
+                        args=(connection_id, message_data),
+                        name=f"task-{connection_id}",
+                        daemon=False  # Changed to False for better control
+                    )
+                    
+                    # Add thread to manager for tracking
+                    thread_manager.add_thread(task_thread)
+                    task_thread.start()
+                    
+                    # Clean up completed threads
+                    thread_manager.remove_thread(task_thread)
+                    
                 else:
                     print(f"Unknown event type: {message_data.get('event')}")
 
@@ -55,27 +103,32 @@ async def websocket_endpoint(websocket: WebSocket):
         print(f"WebSocket connection cleaned up: {connection_id}")
 
 
-async def handle_user_input(connection_id: str, message_data: dict):
+def handle_user_input(connection_id: str, message_data: dict):
     """Handle user input message and start task execution"""
     try:
+        # Check if shutdown is requested
+        if thread_manager.shutdown_event.is_set():
+            print(f"Shutdown requested, skipping task for connection {connection_id}")
+            return
+        
         # Extract user input data
         user_input = message_data.get("data", {})
         text = user_input.get("text", "")
+        team_name = user_input.get("team", "general_team")  # Default to general_team
 
         if not text:
             print("Empty user input received")
             return
 
-        print(f"Processing user input: {text[:50]}...")
+        print(f"Processing user input: {text[:50]}... with team: {team_name}")
 
         # Process user input and create task
-        task_id = await task_processor.process_user_input(connection_id, user_input)
+        task_id = task_processor.process_user_input(connection_id, user_input)
 
-        # Start task execution in background
-        import asyncio
-        asyncio.create_task(task_processor.execute_task(task_id, text))
+        # Execute task synchronously (this will stream results via WebSocket)
+        task_processor.execute_task(task_id, text, team_name)
 
-        print(f"Task {task_id} started for connection {connection_id}")
+        print(f"Task {task_id} completed for connection {connection_id} with team {team_name}")
 
     except Exception as e:
         print(f"Error handling user input: {e}")
@@ -87,4 +140,4 @@ async def handle_user_input(connection_id: str, message_data: dict):
                 "msg": f"Failed to process task: {str(e)}"
             }
         }
-        await websocket_manager.send_message(connection_id, error_message) 
+        websocket_manager.send_message(connection_id, error_message) 

agentmesh/protocol/agent.py

@@ -38,6 +38,7 @@ def __init__(self, name: str, system_prompt: str, description: str, model: LLMMo
         self.max_steps = max_steps  # max ReAct steps, None means no limit
         self.conversation_history = []
         self.action_history = []
+        self.captured_actions = []  # Initialize captured actions list
         self.ext_data = ""
         self.output_mode = output_mode
         if tools:
@@ -140,10 +141,6 @@ def step(self):
         final_answer = None
         current_step = 0
 
-        # Initialize captured actions list (if it doesn't exist)
-        if not hasattr(self, 'captured_actions'):
-            self.captured_actions = []
-
         # Initialize final answer (if it doesn't exist)
         if not hasattr(self, 'final_answer'):
             self.final_answer = ""
@@ -277,7 +274,25 @@ def step(self):
                 tool: BaseTool = self._find_tool(parsed["action"])
                 observation = ""
                 if tool:
+                    # Record start time for execution timing
+                    start_time = time.time()
+
                     tool_result = tool.execute_tool(parsed.get("action_input", {}))
+
+                    # Calculate execution time
+                    execution_time = time.time() - start_time
+
+                    # Capture tool use for tracking
+                    self.capture_tool_use(
+                        tool_name=parsed["action"],
+                        input_params=parsed.get("action_input", {}),
+                        thought=parsed['thought'],
+                        output=tool_result.result,
+                        status=tool_result.status,
+                        error_message=tool_result.result if tool_result.status == "error" else None,
+                        execution_time=execution_time,
+                    )
+
                     # Update conversation history
                     parsed["Observation"] = {
                         "status": tool_result.status,
@@ -326,9 +341,25 @@ def _execute_post_process_tools(self):
             # Set tool context
             tool.context = self
 
+            # Record start time for execution timing
+            start_time = time.time()
+
             # Execute tool (with empty parameters, tool will extract needed info from context)
             result = tool.execute({})
 
+            # Calculate execution time
+            execution_time = time.time() - start_time
+
+            # Capture tool use for tracking
+            self.capture_tool_use(
+                tool_name=tool.name,
+                input_params={},  # Post-process tools typically don't take parameters
+                output=result.result,
+                status=result.status,
+                error_message=str(result.result) if result.status == "error" else None,
+                execution_time=execution_time
+            )
+
             # Log result
             if result.status == "success":
                 # Print tool execution result in the desired format
@@ -424,10 +455,12 @@ def _fetch_agents_outputs(self) -> str:
                 f"member name: {agent_output.agent_name}\noutput content: {agent_output.output}\n\n")
         return "\n".join(agent_outputs_list)
 
-    def capture_tool_use(self, tool_name, input_params, output, status, error_message=None, execution_time=0.0):
+    def capture_tool_use(self, tool_name, input_params, output, status, thought=None, error_message=None,
+                         execution_time=0.0):
         """
         Capture a tool use action.
         
+        :param thought: thought content
         :param tool_name: Name of the tool used
         :param input_params: Parameters passed to the tool
         :param output: Output from the tool
@@ -448,52 +481,10 @@ def capture_tool_use(self, tool_name, input_params, output, status, error_messag
             agent_id=self.id if hasattr(self, 'id') else str(id(self)),
             agent_name=self.name,
             action_type=AgentActionType.TOOL_USE,
-            tool_result=tool_result
+            tool_result=tool_result,
+            thought=thought
         )
 
-        if not hasattr(self, 'captured_actions'):
-            self.captured_actions = []
-
-        self.captured_actions.append(action)
-
-        return action
-
-    def capture_thinking(self, thought_content):
-        """
-        Capture a thinking action.
-        
-        :param thought_content: Content of the thought
-        """
-        action = AgentAction(
-            agent_id=self.id if hasattr(self, 'id') else str(id(self)),
-            agent_name=self.name,
-            action_type=AgentActionType.THINKING,
-            content=thought_content
-        )
-
-        if not hasattr(self, 'captured_actions'):
-            self.captured_actions = []
-
-        self.captured_actions.append(action)
-
-        return action
-
-    def capture_final_answer(self, answer_content):
-        """
-        Capture a final answer action.
-        
-        :param answer_content: Content of the final answer
-        """
-        action = AgentAction(
-            agent_id=self.id if hasattr(self, 'id') else str(id(self)),
-            agent_name=self.name,
-            action_type=AgentActionType.FINAL_ANSWER,
-            content=answer_content
-        )
-
-        if not hasattr(self, 'captured_actions'):
-            self.captured_actions = []
-
         self.captured_actions.append(action)
 
         return action

agentmesh/protocol/result.py

@@ -55,6 +55,7 @@ class AgentAction:
     id: str = field(default_factory=lambda: str(uuid.uuid4()))
     content: str = ""
     tool_result: Optional[ToolResult] = None
+    thought: Optional[str] = None
     timestamp: float = field(default_factory=time.time)
 
 
@@ -113,13 +114,14 @@ def to_dict(self) -> Dict[str, Any]:
                     "agent_name": action.agent_name,
                     "action_type": action.action_type.value,
                     "content": action.content,
+                    "thought": action.thought,
                     "tool_result": {
                         "tool_name": action.tool_result.tool_name,
                         "input_params": action.tool_result.input_params,
                         "output": action.tool_result.output,
                         "status": action.tool_result.status,
                         "error_message": action.tool_result.error_message,
-                        "execution_time": action.tool_result.execution_time
+                        "execution_time": action.tool_result.execution_time,
                     } if action.tool_result else None,
                     "timestamp": action.timestamp
                 }