Update agents/agents_nodes.py

agents/agents_nodes.py

@@ -86,37 +86,83 @@ F_MAPPING = {
 
 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"}}
+        # Validate messages list length
+        if len(state["messages"]) < 2:
+            return {"output": {"error": "Insufficient messages to process"}}
+
+        # Safely access list indices
+        last_msg = state["messages"][-1]
+        second_last_msg = state["messages"][-2]
+
+        # Verify message types
+        if not isinstance(last_msg, ToolMessage):
+            return {"output": {"error": "Last message must be ToolMessage"}}
+        if not isinstance(second_last_msg, AIMessage):
+            return {"output": {"error": "Second-last message must be AIMessage"}}
+
+        # 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"}}
         
-        tool_message = state["messages"][-1]
-        # Parse the content of the tool message as JSON
-        tool_result = json.loads(tool_message.content)
+        first_tool_call = second_last_msg.tool_calls
+        args = first_tool_call.get("args", {})
         
-        # 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}"}}
+        # Process factor type
+        factor_type = args.get("F")
+        if not factor_type or factor_type not in F_MAPPING:
+            return {"output": {"error": f"Invalid 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"}}
-        
+            return {"output": {"error": f"Result key {result_key} not found"}}
+            
         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)}"}}
+    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)}"}}