Update app.py

app.py

@@ -126,18 +126,23 @@ def init_state(question: str) -> AgentState:
 
 
 def should_continue(state: AgentState) -> str:
-    last_message = state["history"][-1]
+    history = state.get("history", [])
 
-    # If the last message is FINAL ANSWER, stop
+    if not history:
+        return "reason"  # No history yet, reason first
+
+    last_message = history[-1]
+
+    # End if agent has produced a final answer
     if isinstance(last_message, AIMessage) and "FINAL ANSWER:" in last_message.content:
         return "end"
 
-    # If an action_request was just added, continue to tool
-    for msg in reversed(state["history"]):
+    # If an action_request exists, trigger tool use
+    for msg in reversed(history):
         if isinstance(msg, dict) and msg.get("role") == "action_request":
             return "continue"
 
-    # Otherwise, reason again
+    # Otherwise, go back to reasoning
     return "reason"
 
 
@@ -153,23 +158,28 @@ def reasoning_node(state: AgentState) -> AgentState:
     if not GOOGLE_API_KEY:
         raise ValueError("GOOGLE_API_KEY not set in environment variables.")
 
-    # Ensure history ends with a HumanMessage
-    if not state.get("history") or not isinstance(state["history"][-1], HumanMessage):
-        state["history"] = state.get("history", [])
+    # Ensure history is initialized and ends with a HumanMessage
+    if "history" not in state or not isinstance(state["history"], list):
+        state["history"] = []
+    if not state["history"] or not isinstance(state["history"][-1], HumanMessage):
         state["history"].append(HumanMessage(content="Continue."))
 
     # Ensure context is a dictionary
     if not isinstance(state.get("context"), dict):
         state["context"] = {}
 
-    # Initialize the Gemini model (via LangChain wrapper)
+    # Ensure reasoning and iterations keys are present
+    state.setdefault("reasoning", "")
+    state.setdefault("iterations", 0)
+
+    # Initialize Gemini model via LangChain
     llm = ChatGoogleGenerativeAI(
         model="gemini-1.5-flash",
         temperature=0.1,
         google_api_key=GOOGLE_API_KEY
     )
 
-    # Create prompt from messages
+    # Create prompt
     prompt = ChatPromptTemplate.from_messages([
         ("system", (
             "You're an expert problem solver. Analyze the question, select the best tool, "
@@ -183,11 +193,11 @@ def reasoning_node(state: AgentState) -> AgentState:
         *state["history"]
     ])
 
-    # Build and invoke the chain
+    # Invoke model
     chain = prompt | llm
     response = chain.invoke({
         "context": state["context"],
-        "reasoning": state.get("reasoning", ""),
+        "reasoning": state["reasoning"],
         "question": state["question"]
     })
 
@@ -196,10 +206,10 @@ def reasoning_node(state: AgentState) -> AgentState:
 
     # Update state
     state["history"].append(AIMessage(content=content))
-    state["reasoning"] += f"\nStep {state['iterations']+1}: {reasoning}"
+    state["reasoning"] += f"\nStep {state['iterations'] + 1}: {reasoning}"
     state["iterations"] += 1
 
-    # Decide next step based on action
+    # Store either final answer or tool to call
     if "final answer" in action.lower():
         state["history"].append(AIMessage(content=f"FINAL ANSWER: {action_input}"))
     else:
@@ -211,12 +221,17 @@ def reasoning_node(state: AgentState) -> AgentState:
     return state
 
 
+
     
 
 def tool_node(state: AgentState) -> AgentState:
     from langchain.schema import AIMessage
 
-    # Get the last tool action request
+    # Ensure history exists
+    if "history" not in state or not isinstance(state["history"], list):
+        raise ValueError("Invalid or missing history in state")
+
+    # Find the most recent action request in history
     tool_call = None
     for msg in reversed(state["history"]):
         if isinstance(msg, dict) and msg.get("role") == "action_request":
@@ -226,11 +241,17 @@ def tool_node(state: AgentState) -> AgentState:
     if not tool_call:
         raise ValueError("No tool call found in history")
 
-    tool_name = tool_call["tool"]
-    tool_input = tool_call["input"]
+    tool_name = tool_call.get("tool")
+    tool_input = tool_call.get("input")
+
+    # Defensive check for missing tool or input
+    if not tool_name or tool_input is None:
+        raise ValueError("Tool name or input missing from action request")
 
     # Look up and invoke the tool
-    tool_fn = next((t for t in BasicAgent().tools if t.__name__ == tool_name), None)
+    agent = BasicAgent()  # Create agent to access tools
+    tool_fn = next((t for t in agent.tools if t.__name__ == tool_name), None)
+
     if tool_fn is None:
         raise ValueError(f"Tool '{tool_name}' not found")
 
@@ -239,12 +260,13 @@ def tool_node(state: AgentState) -> AgentState:
     except Exception as e:
         tool_output = f"[Tool Error] {str(e)}"
 
-    # Store tool result as AIMessage
+    # Add output to history as an 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()