Update agents/agents_nodes.py

agents/agents_nodes.py

@@ -86,83 +86,46 @@ F_MAPPING = {
 
 def format_output(state: AgentState):
     try:
-        # Validate messages list length
+        # 1. Validate minimum message count
         if len(state["messages"]) < 2:
-            return {"output": {"error": "Insufficient messages to process"}}
-
-        # Safely access list indices
+            return {"output": {"error": f"Only {len(state['messages'])} messages, need 2 for processing"}}
+        
+        # 2. Safe element access
         last_msg = state["messages"][-1]
         second_last_msg = state["messages"][-2]
 
-        # Verify message types
+        # 3. Verify message types
         if not isinstance(last_msg, ToolMessage):
-            return {"output": {"error": "Last message must be ToolMessage"}}
+            return {"output": {"error": f"Last message is {type(last_msg).__name__}, expected ToolMessage"}}
         if not isinstance(second_last_msg, AIMessage):
-            return {"output": {"error": "Second-last message must be AIMessage"}}
+            return {"output": {"error": f"Second-last message is {type(second_last_msg).__name__}, expected AIMessage"}}
 
-        # Parse tool result
+        # 4. Parse tool result
         try:
             tool_result = json.loads(last_msg.content)
         except json.JSONDecodeError:
             return {"output": {"error": "ToolMessage content is invalid JSON"}}
 
-        # Validate tool call structure
-        if not second_last_msg.tool_calls:
-            return {"output": {"error": "AIMessage has no tool calls"}}
+        # 5. Validate tool call structure
+        if not hasattr(second_last_msg, 'tool_calls') or not second_last_msg.tool_calls:
+            return {"output": {"error": "AIMessage contains no tool calls"}}
         
-        first_tool_call = second_last_msg.tool_calls
-        args = first_tool_call.get("args", {})
+        # 6. Handle single vs multiple tool calls
+        tool_call = second_last_msg.tool_calls if isinstance(second_last_msg.tool_calls, list) else second_last_msg.tool_calls
+        args = tool_call.get("args", {})
         
-        # Process factor type
-        factor_type = args.get("F")
+        # 7. Process factor type
+        factor_type = args.get("F", "").strip()
         if not factor_type or factor_type not in F_MAPPING:
-            return {"output": {"error": f"Invalid factor type: {factor_type}"}}
+            return {"output": {"error": f"Invalid/missing factor type: '{factor_type}'"}}
         
         result_key = F_MAPPING[factor_type]
         if result_key not in tool_result:
-            return {"output": {"error": f"Result key {result_key} not found"}}
+            return {"output": {"error": f"Result key '{result_key}' not found in tool output"}}
             
+        # 8. Format final output
         value = tool_result[result_key]
         return {"output": {result_key: round(float(value), 2)}}
         
     except Exception as e:
-        return {"output": {"error": f"Formatting failed: {str(e)}"}}
-
-
-
-# def format_output(state: AgentState):
-#     try:
-#         # The last message should be the ToolMessage (from the tool node)
-#         if not state["messages"] or not isinstance(state["messages"][-1], ToolMessage):
-#             return {"output": {"error": "No tool result found in the last message"}}
-        
-#         tool_message = state["messages"][-1]
-#         # Parse the content of the tool message as JSON
-#         tool_result = json.loads(tool_message.content)
-        
-#         # The second last message should be the AIMessage with the tool call
-#         if len(state["messages"]) < 2 or not isinstance(state["messages"][-2], AIMessage):
-#             return {"output": {"error": "No AI message (with tool call) found before the tool message"}}
-        
-#         ai_message = state["messages"][-2]
-#         if not ai_message.tool_calls:
-#             return {"output": {"error": "The AI message does not contain tool calls"}}
-        
-#         # We take the first tool call (since we forced one tool)
-#         tool_call = ai_message.tool_calls
-#         args = tool_call["args"]
-        
-#         # Get the factor type from the args
-#         factor_type = args["F"]
-#         if factor_type not in F_MAPPING:
-#             return {"output": {"error": f"Unrecognized factor type: {factor_type}"}}
-        
-#         result_key = F_MAPPING[factor_type]
-#         if result_key not in tool_result:
-#             return {"output": {"error": f"Expected key {result_key} not found in tool result"}}
-        
-#         value = tool_result[result_key]
-#         return {"output": {result_key: round(float(value), 2)}}
-        
-#     except (KeyError, TypeError, json.JSONDecodeError, IndexError) as e:
-#         return {"output": {"error": f"Result formatting failed: {str(e)}"}}
+        return {"output": {"error": f"Unhandled exception: {str(e)}"}}