Update app.py

app.py

@@ -127,13 +127,25 @@ def finalize_node(state: AgentState) -> AgentState:
 # ─── 4) Wrap tools in a ToolNode ───
 tool_node = ToolNode([web_search_tool, ocr_image_tool, parse_excel_tool])
 
+# Add a node to store the previous state
+
+def store_prev_state(state: AgentState) -> AgentState:
+    return {**state, "prev_state": state.copy()}
+
+def merge_tool_output(state: AgentState) -> AgentState:
+    prev_state = state.get("prev_state", {})
+    merged = {**prev_state, **state}
+    merged.pop("prev_state", None)
+    return merged
+
 # ─── 5) Build the graph ───
 graph = StateGraph(AgentState)
 
 # 5.a) Register nodes
 graph.add_node("plan", plan_node)
+graph.add_node("store_prev_state", store_prev_state)
 graph.add_node("tools", tool_node)
-graph.add_node("run_tools", run_tools)
+graph.add_node("merge_tool_output", merge_tool_output)
 graph.add_node("finalize", finalize_node)
 
 # 5.b) Wire START → plan
@@ -148,16 +160,19 @@ def route_plan(plan_out: AgentState) -> str:
 graph.add_conditional_edges(
     "plan",
     route_plan,
-    {"tools": "tools", "finalize": "finalize"}
+    {"tools": "store_prev_state", "finalize": "finalize"}
 )
 
-# 5.d) tools → run_tools
-graph.add_edge("tools", "run_tools")
+# 5.d) store_prev_state → tools
+graph.add_edge("store_prev_state", "tools")
+
+# 5.e) tools → merge_tool_output
+graph.add_edge("tools", "merge_tool_output")
 
-# 5.e) run_tools → finalize
-graph.add_edge("run_tools", "finalize")
+# 5.f) merge_tool_output → finalize
+graph.add_edge("merge_tool_output", "finalize")
 
-# 5.f) finalize → END
+# 5.g) finalize → END
 graph.add_edge("finalize", END)
 
 compiled_graph = graph.compile()