Update agents/agents_nodes.py

agents/agents_nodes.py

@@ -26,54 +26,19 @@ text_generator = pipeline(
 
 
 llm = HuggingFacePipeline(pipeline=text_generator)
-# llm_instantiated = llm.bind_tools(    
-#     [time_value_tool],
-#     tool_choice={"type": "function", "function": {"name": "time_value_tool"}}
-# )
 
-# llm_instantiated = llm | RunnableLambda(
-#     lambda x: x.bind_tools(
-#         [time_value_tool],
-#         tool_choice={"type": "function", "function": {"name": "time_value_tool"}}
-#     )
-# )
-
-# llm_instantiated = llm | RunnableLambda(
-#     lambda x: x.bind(
-#         [time_value_tool],
-#         tool_choice={"type": "function", "function": {"name": "time_value_tool"}}
-#     )
-# )
-
-# Replace this
-
-
-# With this
-llm_instantiated = llm.bind(
-    tools=[time_value_tool],
-    tool_choice="time_value_tool"
+llm_instantiated = llm.bind(    
+    [time_value_tool],
+    tool_choice={"type": "function", "function": {"name": "time_value_tool"}}
 )
 
-
-
-# def agent_node(state: AgentState):
-#     response = llm_instantiated.invoke(state["messages"])
-#     if not (hasattr(response, 'tool_calls') and response.tool_calls):
-#         error_message = AIMessage(content="Error: Model failed to generate tool call.")
-#         return {"messages": [error_message]}
-#     return {"messages": [response]}
-
-
-# Update tool invocation in agent_node
 def agent_node(state: AgentState):
     response = llm_instantiated.invoke(state["messages"])
-    if not hasattr(response, 'tool_calls') or not response.tool_calls:
-        return {"messages": [AIMessage(content="Error: No tool call generated")]}
-    tool_call = response.tool_calls
+    if not (hasattr(response, 'tool_calls') and response.tool_calls):
+        error_message = AIMessage(content="Error: Model failed to generate tool call.")
+        return {"messages": [error_message]}
     return {"messages": [response]}
 
-
-
 # Tool node executes the tool
 tool_node = ToolNode([time_value_tool])
 
@@ -86,46 +51,37 @@ F_MAPPING = {
 
 def format_output(state: AgentState):
     try:
-        # 1. Validate minimum message count
-        if len(state["messages"]) < 2:
-            return {"output": {"error": f"Only {len(state['messages'])} messages, need 2 for processing"}}
+        # 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"}}
         
-        # 2. Safe element access
-        last_msg = state["messages"][-1]
-        second_last_msg = state["messages"][-2]
-
-        # 3. Verify message types
-        if not isinstance(last_msg, ToolMessage):
-            return {"output": {"error": f"Last message is {type(last_msg).__name__}, expected ToolMessage"}}
-        if not isinstance(second_last_msg, AIMessage):
-            return {"output": {"error": f"Second-last message is {type(second_last_msg).__name__}, expected AIMessage"}}
-
-        # 4. Parse tool result
-        try:
-            tool_result = json.loads(last_msg.content)
-        except json.JSONDecodeError:
-            return {"output": {"error": "ToolMessage content is invalid JSON"}}
-
-        # 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"}}
+        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"}}
         
-        # 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", {})
+        ai_message = state["messages"][-2]
+        if not ai_message.tool_calls:
+            return {"output": {"error": "The AI message does not contain tool calls"}}
         
-        # 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/missing factor type: '{factor_type}'"}}
+        # 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"Result key '{result_key}' not found in tool output"}}
-            
-        # 8. Format final output
+            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 Exception as e:
-        return {"output": {"error": f"Unhandled exception: {str(e)}"}}
+    except (KeyError, TypeError, json.JSONDecodeError, IndexError) as e:
+        return {"output": {"error": f"Result formatting failed: {str(e)}"}}