Update app.py

app.py

@@ -42,51 +42,44 @@ tool_node = ToolNode([ocr_image, parse_excel, web_search])
 graph = StateGraph(AgentState)
 
 # ────────────────────────
-# Edge 1: START → LLM
-# Wrap the user_input into state.messages
+graph.add_node("agent", llm)
+graph.add_node("tools", tool_node)
+# Edge A: START → "agent"
+#    Wrap the user_input into state["messages"]
 graph.add_edge(
     START,
-    llm,
-    lambda state, user_input: {"messages": [user_input]},
+    "agent",
+    transition=lambda state, user_input: {"messages": [user_input]},
 )
 
-# ────────────────────────
-# Edge 2: LLM → ToolNode
-# Only fire when the LLM returns a dict with exactly "tool":"ocr_image" or "tool":"parse_excel"
-# The lambda must return that dict so the ToolNode can extract its arguments.
-def route_to_tool(state: AgentState, llm_out):
-    # Expecting llm_out to be a dict like:
-    #   {"tool": "ocr_image", "path": "invoice.png"}
-    # or {"tool": "parse_excel", "path":"sales.xlsx", "sheet_name":"Sheet1"}
-    if isinstance(llm_out, dict) and llm_out.get("tool") in {"ocr_image", "parse_excel", "web_search"}:
-        return llm_out
-    return None  # → do not invoke the tool if it's not matching
-
+# Edge C: "tools" → "agent"
+#    Whatever string the tool returns becomes the next prompt to the LLM
 graph.add_edge(
-    llm,
-    tool_node,
-    route_to_tool,
+    "tools",
+    "agent",
+    transition=lambda state, tool_output: tool_output,
 )
 
-# ────────────────────────
-# Edge 3: ToolNode → LLM
-# Whatever the tool returns (a string), feed that straight back into the LLM as the next turn.
-graph.add_edge(
-    tool_node,
-    llm,
-    lambda state, tool_out: tool_out
-)
+# 7) Use add_conditional_edges out of "agent" instead of two separate edges
 
-# ────────────────────────
-# Edge 4: LLM → END
-# Once the LLM is done reasoning (without calling any more tools), return its output string.
-graph.add_edge(
-    llm,
-    END,
-    lambda state, final_str: final_str
+def route_agent(state: AgentState, agent_out):
+    """
+    Return "tools" if the LLM output is a dict with a valid "tool" key,
+    otherwise return "final".
+    """
+    if isinstance(agent_out, dict) and agent_out.get("tool") in {"ocr_image", "parse_excel"}:
+        return "tools"
+    return "final"
+
+graph.add_conditional_edges(
+    "agent",
+    route_agent,
+    {
+        "tools": "tools",  # if route_agent(...) == "tools", go to the "tools" node
+        "final": END,      # if route_agent(...) == "final", transition to END
+    },
 )
 
-
 def respond_to_input(user_input: str) -> str:
     initial_state: AgentState = {"messages": []}
     return graph.run(initial_state, user_input)