Update app.py

app.py

@@ -32,7 +32,7 @@ def plan_node(state: AgentState) -> AgentState:
     If confident: return {"final_answer": "<answer>"}
     Otherwise:   return exactly one of 
                  {"wiki_query": "..."},
-                 
+                 {"web_search_query": "..."},
                  {"ocr_path": "..."},
                  {"excel_path": "...", "excel_sheet_name": "..."},
                  {"audio_path": "..."}
@@ -54,7 +54,7 @@ def plan_node(state: AgentState) -> AgentState:
             "       and nothing else.\n"
             "     • Otherwise, return exactly one of:\n"
             "         {\"wiki_query\":\"<Wikipedia search>\"}\n"
-            # "         {\"web_search_query\":\"<search terms>\"}\n"
+            "         {\"web_search_query\":\"<search terms>\"}\n"
             "         {\"ocr_path\":\"<image path or task_id>\"}\n"
             "         {\"excel_path\":\"<xlsx path>\", \"excel_sheet_name\":\"<sheet name>\"}\n"
             "         {\"audio_path\":\"<audio path or task_id>\"}\n"
@@ -185,7 +185,7 @@ def inspect_node(state: AgentState) -> AgentState:
         "and nothing else.\n"
         "Otherwise, return exactly one of these JSON literals to fetch another tool:\n"
         "  {\"wiki_query\":\"<query for Wikipedia>\"}\n"
-        # "  {\"web_search_query\":\"<search terms>\"}\n"
+    
         "  {\"ocr_path\":\"<image path or task_id>\"}\n"
         "  {\"excel_path\":\"<xls path>\", \"excel_sheet_name\":\"<sheet name>\"}\n"
         "  {\"audio_path\":\"<audio path or task_id>\"}\n"
@@ -196,30 +196,32 @@ def inspect_node(state: AgentState) -> AgentState:
     raw = llm_response.content.strip()
 
     new_msgs = state["messages"] + [AIMessage(content=raw)]
+
+    # Try to parse the LLM’s JSON
     try:
         parsed = json.loads(raw)
         if isinstance(parsed, dict):
-            partial: AgentState = {"messages": new_msgs}
-            allowed = {
-                "final_answer",
-                "wiki_query",
-                "web_search_query",
-                "ocr_path",
-                "excel_path",
-                "excel_sheet_name",
-                "audio_path"
-            }
-            for k, v in parsed.items():
-                if k in allowed:
-                    partial[k] = v
-            return partial
+            # Check if the LLM explicitly provided "final_answer"
+            if "final_answer" in parsed:
+                return {"messages": new_msgs, "final_answer": parsed["final_answer"]}
+
+            # Check if it requested exactly one tool
+            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:
+                # Return exactly that key so the graph goes to store_prev_state → tools
+                return {**{"messages": new_msgs}, **{k: parsed[k] for k in requested_keys}}
+
     except json.JSONDecodeError:
         pass
 
-    return {
-        "messages": new_msgs,
-        "final_answer": "ERROR: could not parse inspect decision."
-    }
+    # ----- NEW FALLBACK: If LLM didn’t give a valid tool key or final_answer,
+    # return the existing INTERIM_ANSWER as the final answer. -----
+    if ia := state.get("final_answer"):
+        return {"messages": new_msgs, "final_answer": ia}
+
+    # If there is no INTERIM_ANSWER either, then we genuinely cannot proceed
+    return {"messages": new_msgs, "final_answer": "ERROR: could not parse inspect decision."}
 
 
 # ─── 6) finalize_node ───