Update app.py

app.py

@@ -123,18 +123,23 @@ def init_state(question: str) -> AgentState:
         "iterations": 0
     }
 
+
+
 def should_continue(state: AgentState) -> str:
-    """Determine if agent should continue or finish"""
-    last_msg = state['history'][-1]
-    
-    # Stop conditions
-    if state['iterations'] >= 5:
-        return "end"
-    if "FINAL ANSWER:" in last_msg.get('content', ''):
+    last_message = state["history"][-1]
+
+    # If the last message is FINAL ANSWER, stop
+    if isinstance(last_message, AIMessage) and "FINAL ANSWER:" in last_message.content:
         return "end"
-    if last_msg['role'] == 'tool':
-        return "reason"
-    return "continue"
+
+    # If an action_request was just added, continue to tool
+    for msg in reversed(state["history"]):
+        if isinstance(msg, dict) and msg.get("role") == "action_request":
+            return "continue"
+
+    # Otherwise, reason again
+    return "reason"
+
 
 
 def reasoning_node(state: AgentState) -> AgentState:
@@ -209,33 +214,37 @@ def reasoning_node(state: AgentState) -> AgentState:
     
 
 def tool_node(state: AgentState) -> AgentState:
-    """Execute selected tool and update state"""
-    last_action = state['history'][-1]
-    tool_name = last_action['tool']
-    tool_input = last_action['input']
-    
-    # Tool mapping
-    tools = {
-        "duckduckgo_search": duckduckgo_search,
-        "wikipedia_search": wikipedia_search,
-        "arxiv_search": arxiv_search,
-        "document_qa": document_qa,
-        "python_execution": python_execution
-    }
-    
-    # Execute tool
-    tool_result = tools[tool_name].invoke(tool_input)
-    
-    # Update state
-    state['history'].append(ToolMessage(
-        content=tool_result,
-        tool_call_id=tool_name
-    ))
-    state['context'] = f"Tool Result ({tool_name}): {tool_result}"
-    state['iterations'] += 1
-    
+    from langchain.schema import AIMessage
+
+    # Get the last tool action request
+    tool_call = None
+    for msg in reversed(state["history"]):
+        if isinstance(msg, dict) and msg.get("role") == "action_request":
+            tool_call = msg
+            break
+
+    if not tool_call:
+        raise ValueError("No tool call found in history")
+
+    tool_name = tool_call["tool"]
+    tool_input = tool_call["input"]
+
+    # Look up and invoke the tool
+    tool_fn = next((t for t in BasicAgent().tools if t.__name__ == tool_name), None)
+    if tool_fn is None:
+        raise ValueError(f"Tool '{tool_name}' not found")
+
+    try:
+        tool_output = tool_fn(tool_input)
+    except Exception as e:
+        tool_output = f"[Tool Error] {str(e)}"
+
+    # Store tool result as AIMessage
+    state["history"].append(AIMessage(content=f"[{tool_name} output]\n{tool_output}"))
+
     return state
 
+
 def parse_agent_response(response: str) -> tuple:
     """Extract reasoning, action, and input from response"""
     reasoning = response.split("Reasoning:")[1].split("Action:")[0].strip()