Update app.py

app.py

@@ -24,17 +24,18 @@ from tools import ocr_image_tool, parse_excel_tool, web_search_tool, run_tools,
 
 llm = ChatOpenAI(model_name="gpt-4o-mini")
 
+# ─── 1) plan_node ───
 # ─── 1) plan_node ───
 def plan_node(state: AgentState) -> AgentState:
     """
     Step 1: Ask GPT to draft a concise direct answer (INTERIM_ANSWER),
             then decide if it's confident enough to stop or if it needs one tool.
-    If confident: return {"final_answer": "<answer>"}
-    Otherwise:   return exactly one of 
-                 {"wiki_query": "..."},
-                 {"ocr_path": "..."},
-                 {"excel_path": "...", "excel_sheet_name": "..."},
-                 {"audio_path": "..."}
+    If confident: return {"final_answer":"<answer>"}
+    Otherwise:   return exactly one of
+                 {"wiki_query":"..."},
+                 {"ocr_path":"..."},
+                 {"excel_path":"...", "excel_sheet_name":"..."},
+                 {"audio_path":"..."}
     """
     prior_msgs = state.get("messages", [])
     user_input = ""
@@ -71,19 +72,18 @@ def plan_node(state: AgentState) -> AgentState:
     try:
         parsed = json.loads(llm_out)
         if isinstance(parsed, dict):
-            # Build a clean partial state that only carries over messages and tool_counter
+            # Build a fresh state that carries only messages + tool_counter
             partial: AgentState = {
                 "messages": new_msgs,
-                "tool_counter": state.get("tool_counter", 0)
+                "tool_counter": state.get("tool_counter", 0),
             }
-            # Only copy allowed keys (LMM won't know about web_search_query, so it won't appear)
             allowed_keys = {
                 "final_answer",
                 "wiki_query",
                 "ocr_path",
                 "excel_path",
                 "excel_sheet_name",
-                "audio_path"
+                "audio_path",
             }
             for k, v in parsed.items():
                 if k in allowed_keys:
@@ -92,10 +92,10 @@ def plan_node(state: AgentState) -> AgentState:
     except json.JSONDecodeError:
         pass
 
-    # Fallback: treat as a final answer (no further tools)
+    # Fallback: interpret as a final answer (no further tools)
     return {
         "messages": new_msgs,
-        "final_answer": "Sorry, I could not parse your intent."
+        "final_answer": "Sorry, I could not parse your intent.",
     }
 
 
@@ -112,7 +112,7 @@ def tool_node(state: AgentState) -> AgentState:
       - ocr_path → ocr_image_tool
       - excel_path → parse_excel_tool
       - audio_path → audio_transcriber_tool
-      - (web_search_query path kept but not exposed to LLM)
+      - (web_search_query path is still here but not exposed to the LLM)
     """
     tool_counter = state.get("tool_counter", 0)
     if tool_counter > 5:
@@ -120,7 +120,6 @@ def tool_node(state: AgentState) -> AgentState:
     tool_counter += 1
     state["tool_counter"] = tool_counter
 
-    # Only one of these keys should be present at a time
     if state.get("wiki_query"):
         return wikipedia_search_tool(state)
     if state.get("web_search_query"):
@@ -131,7 +130,8 @@ def tool_node(state: AgentState) -> AgentState:
         return parse_excel_tool(state)
     if state.get("audio_path"):
         return audio_transcriber_tool(state)
-    return {}
+
+    return {}  # nothing to do
 
 
 # ─── 4) merge_tool_output ───
@@ -142,10 +142,28 @@ def merge_tool_output(state: AgentState) -> AgentState:
     prev = state.get("prev_state", {}).copy()
 
     # Drop any lingering request keys so they don't persist
-    for dead in ["wiki_query", "web_search_query", "ocr_path", "excel_path", "excel_sheet_name", "audio_path"]:
+    for dead in [
+        "wiki_query",
+        "web_search_query",
+        "ocr_path",
+        "excel_path",
+        "excel_sheet_name",
+        "audio_path",
+    ]:
         prev.pop(dead, None)
 
     merged = {**prev, **state}
+    # Also drop them from the merged result
+    for dead in [
+        "wiki_query",
+        "web_search_query",
+        "ocr_path",
+        "excel_path",
+        "excel_sheet_name",
+        "audio_path",
+    ]:
+        merged.pop(dead, None)
+
     merged.pop("prev_state", None)
     return merged
 
@@ -161,11 +179,14 @@ def inspect_node(state: AgentState) -> AgentState:
       • Return {"final_answer":"<final>"} if done, OR
       • Return exactly one tool key to run next (wiki_query / ocr_path / excel_path & excel_sheet_name / audio_path).
     """
+
     # 0) If we've already called tools too many times, force a final answer:
     if state.get("tool_counter", 0) >= 5:
         return {
             "messages": state["messages"],
-            "final_answer": state.get("final_answer", "ERROR: no interim_answer to finalize.")
+            "final_answer": state.get(
+                "final_answer", "ERROR: no interim_answer to finalize."
+            ),
         }
 
     messages_for_llm = []
@@ -219,14 +240,25 @@ def inspect_node(state: AgentState) -> AgentState:
         if isinstance(parsed, dict):
             # If GPT gave a final_answer, we finish here
             if "final_answer" in parsed:
-                return {"messages": new_msgs, "final_answer": parsed["final_answer"]}
+                return {
+                    "messages": new_msgs,
+                    "final_answer": parsed["final_answer"],
+                }
 
-            # If GPT requested exactly one valid tool, we return only that key
-            valid_keys = {"wiki_query", "ocr_path", "excel_path", "excel_sheet_name", "audio_path"}
+            # If GPT requested exactly one valid tool, return only that key + carry tool_counter
+            valid_keys = {
+                "wiki_query",
+                "ocr_path",
+                "excel_path",
+                "excel_sheet_name",
+                "audio_path",
+            }
             requested_keys = set(parsed.keys()) & valid_keys
             if len(requested_keys) == 1:
-                # Build a fresh dict with only messages + that single tool key
-                clean = {"messages": new_msgs}
+                clean: AgentState = {
+                    "messages": new_msgs,
+                    "tool_counter": state.get("tool_counter", 0),
+                }
                 for k in requested_keys:
                     clean[k] = parsed[k]
                 return clean
@@ -238,7 +270,10 @@ def inspect_node(state: AgentState) -> AgentState:
         return {"messages": new_msgs, "final_answer": ia}
 
     # If there is no interim either, we cannot proceed
-    return {"messages": new_msgs, "final_answer": "ERROR: could not parse inspect decision."}
+    return {
+        "messages": new_msgs,
+        "final_answer": "ERROR: could not parse inspect decision.",
+    }
 
 
 # ─── 6) finalize_node ───
@@ -274,7 +309,7 @@ def route_plan(plan_out: AgentState) -> str:
 graph.add_conditional_edges(
     "plan",
     route_plan,
-    {"store_prev_state": "store_prev_state", "finalize": "finalize"}
+    {"store_prev_state": "store_prev_state", "finalize": "finalize"},
 )
 
 # store_prev_state → tools
@@ -295,7 +330,7 @@ def route_inspect(inspect_out: AgentState) -> str:
 graph.add_conditional_edges(
     "inspect",
     route_inspect,
-    {"store_prev_state": "store_prev_state", "finalize": "finalize"}
+    {"store_prev_state": "store_prev_state", "finalize": "finalize"},
 )
 
 # finalize → END
@@ -327,7 +362,7 @@ def respond_to_input(user_input: str, task_id) -> str:
     initial_state: AgentState = {
         "messages": [system_msg, human_msg],
         "task_id": task_id,
-        "tool_counter": 0
+        "tool_counter": 0,
     }
     final_state = compiled_graph.invoke(initial_state)
     return final_state.get("final_answer", "Error: No final answer generated.")