Update app.py

app.py

@@ -26,8 +26,14 @@ llm = ChatOpenAI(model_name="gpt-4.1")
 
 # agent = create_react_agent(model=llm, tools=tool_node)
 
-# ─── 2) Revised plan_node ───
 def plan_node(state: AgentState) -> AgentState:
+    """
+    This plan_node will ask GPT to:
+      1) First write a concise *direct* answer.
+      2) Then decide whether it’s confident enough to stop (return {"final_answer": ...})
+         or if it needs to verify via one tool (return exactly one of {"wiki_query":...}, 
+         {"web_search_query":...}, {"ocr_path":...}, {"excel_path":...,"excel_sheet_name":...}, or {"audio_path":...}).
+    """
     prior_msgs = state.get("messages", [])
     user_input = ""
     for msg in reversed(prior_msgs):
@@ -35,62 +41,63 @@ def plan_node(state: AgentState) -> AgentState:
             user_input = msg.content
             break
 
-    # (1) Build your system/human messages exactly as before
+    # (1) Build a fresh SystemMessage that tells the LLM exactly how to self‐evaluate
     system_msg = SystemMessage(
         content=(
-            "You are an agent that decides whether to call a tool or answer directly.\n"
-            "User's question: \"" + user_input + "\"\n\n"
-            "• If you can answer directly, return exactly {\"final_answer\":\"<your answer>\"}.\n"
-            "• Otherwise, respond with exactly one of:\n"
-            "    {\"web_search_query\":\"<search terms>\"}\n"
-            "    {\"ocr_path\":\"<path to image>\"}\n"
-            "    {\"excel_path\":\"<path to xlsx>\", \"excel_sheet_name\":\"<sheet>\"}\n"
-            "    {\"audio_path\":\"<path to audio file>\"}\n"
-            "    {\"wiki_query\":\"<wikipedia search terms>\"}\n"
-            "Do not include any extra characters or markdown—only the JSON literal."
+            "You are an agent that must do two things in a single JSON output:\n\n"
+            "  1) Produce a concise, direct answer to the user’s question (no explanation, just the answer).  \n"
+            "  2) Judge whether that answer is reliable.  \n"
+            "     • If you are fully confident and do NOT need any external verification, return exactly:\n"
+            "         {\"final_answer\":\"<your concise answer>\"}\n"
+            "       and nothing else.\n"
+            "     • If you think you need to verify or look something up first, return exactly one of the following (and nothing else):\n"
+            "         {\"wiki_query\":\"<search terms for Wikipedia>\"}\n"
+            "         {\"web_search_query\":\"<search terms>\"}\n"
+            "         {\"ocr_path\":\"<local image path or task_id>\"}\n"
+            "         {\"excel_path\":\"<local .xlsx path>\", \"excel_sheet_name\":\"<sheet name>\"}\n"
+            "         {\"audio_path\":\"<local audio path or task_id>\"}\n\n"
+            "       You must pick exactly one key—either final_answer or exactly one tool key.\n"
+            "       Do NOT wrap it in any markdown or extra text.  Only output a single JSON object.\n"
+            "\n"
+            f"User’s question: \"{user_input}\"\n"
         )
     )
     human_msg = HumanMessage(content=user_input)
 
-    # (2) Call the LLM
-    llm_response = llm.invoke([system_msg, human_msg])
+    # (2) Call the LLM with this single system/human pair
+    llm_response = llm([system_msg, human_msg])
     llm_out = llm_response.content.strip()
 
-    # ── DEBUG: print raw LLM output ──
-    # print("\n>>> plan_node got raw LLM output:")
-    # print(llm_out)
-    # print("<<< end raw output\n")
-
-    # (3) Append the LLM output to the message history
+    # (3) Append the LLM output into the message history
     ai_msg = AIMessage(content=llm_out)
     new_msgs = prior_msgs.copy() + [ai_msg]
 
-    # (4) Try parsing as JSON
+    # (4) Attempt to parse that JSON
     try:
         parsed = json.loads(llm_out)
-        # print(">>> plan_node parsed JSON:", parsed)
         if isinstance(parsed, dict):
             partial: AgentState = {"messages": new_msgs}
-            allowed = {
+            allowed_keys = {
+                "final_answer",
+                "wiki_query",
                 "web_search_query",
                 "ocr_path",
                 "excel_path",
                 "excel_sheet_name",
-                "audio_path",
-                "wiki_query",
-                "final_answer"
+                "audio_path"
             }
             for k, v in parsed.items():
-                if k in allowed:
+                if k in allowed_keys:
                     partial[k] = v
-                    # print(f">>> plan_node is setting {k!r} → {v!r}")
             return partial
-    except json.JSONDecodeError as e:
-        # print(">>> plan_node JSON parse error:", e)
+    except json.JSONDecodeError:
         pass
-    # (5) Fallback
-    # print(">>> plan_node falling back to final_answer alone\n")
-    return {"messages": new_msgs, "final_answer": "Sorry, I could not parse your intent."}
+
+    # (5) If parsing failed, fall back to a safe “sorry” answer
+    return {
+        "messages": new_msgs,
+        "final_answer": "Sorry, I could not parse your intent."
+    }