Update app.py

app.py

@@ -26,49 +26,49 @@ tool_node = ToolNode([ocr_image_tool, parse_excel_tool, web_search_tool])
 llm = ChatOpenAI(model_name="gpt-4.1-mini", temperature=0.0)
 
 # agent = create_react_agent(model=llm, tools=tool_node)
-
-# ─── Revised plan_node with NO extra arguments ───
 def plan_node(state: AgentState) -> AgentState:
     """
-    Assumes that `state["messages"]` already ends with a HumanMessage of the user’s question.
-    We look at that last HumanMessage, append it to our new history, and ask the LLM
-    to set exactly one key in a Python dict: web_search_query, ocr_path,
-    excel_path (+ excel_sheet_name), or final_answer.
+    `state["messages"]` must already end in a HumanMessage containing the user’s question.
+    We inspect that last HumanMessage and ask the LLM to set exactly one key:
+      • web_search_query
+      • ocr_path
+      • excel_path (and excel_sheet_name)
+      • final_answer
+    The LLM must return a bare Python‐dict literal containing exactly that one key.
     """
-    # 1) Grab all prior BaseMessage objects (SystemMessage/HumanMessage/AIMessage) from state
+    # 1) Grab prior BaseMessage list
     prior_msgs = state.get("messages", [])
 
-    # 2) Find the very last HumanMessage (the user_input). We assume the last message is one.
-    #    If there is no HumanMessage, we treat user_input as empty.
+    # 2) Extract the last HumanMessage content (the user question)
     user_input = ""
     for msg in reversed(prior_msgs):
         if isinstance(msg, HumanMessage):
             user_input = msg.content
             break
 
-    # 3) Build our new chat history by re‐using prior_msgs. It already includes that HumanMessage.
+    # 3) Build new_history = copy of prior_msgs (it already contains that HumanMessage)
     new_history = prior_msgs.copy()
 
-    # 4) Add a SystemMessage that instructs the LLM how to choose exactly one key
+    # 4) Append a SystemMessage explaining how to return exactly one key
     explanation = SystemMessage(
         content=(
-            "You can set exactly one of the following keys in a Python dict, and nothing else:\n"
-            "  • web_search_query: <search terms>  \n"
-            "  • ocr_path: <path to an image file>  \n"
-            "  • excel_path: <path to a .xlsx file>  \n"
-            "  • excel_sheet_name: <sheet name>  \n"
+            "You can set exactly one of these keys in a Python dict (and nothing else):\n"
+            "  • web_search_query: <search terms>\n"
+            "  • ocr_path: <path to an image file>\n"
+            "  • excel_path: <path to a .xlsx file>\n"
+            "  • excel_sheet_name: <sheet name>\n"
             "Or, if no tool is needed, set final_answer: <your answer>.\n"
             "Example: {'web_search_query':'Mercedes Sosa discography'}\n"
             "Respond with only that Python dict literal—no extra text or explanation."
         )
     )
 
-    # 5) Compose the prompt as a list of BaseMessage, then call the LLM
+    # 5) Call the LLM with [ all previous BaseMessages ] + explanation
     prompt_messages = new_history + [explanation]
     llm_response = llm(prompt_messages)
     llm_out = llm_response.content.strip()
 
-    # 6) Parse the LLM’s output as a dict
+    # 6) Try to parse the LLM output as a dict
     try:
         parsed = eval(llm_out, {}, {})
         if isinstance(parsed, dict):
@@ -93,18 +93,20 @@ def plan_node(state: AgentState) -> AgentState:
         "final_answer": "Sorry, I could not parse your intent."
     }
 
-
-# ─── Revised finalize_node with NO extra arguments ───
+# ─── 3) Define finalize_node (only takes state) ───
 def finalize_node(state: AgentState) -> AgentState:
     """
-    Assumes that `state['messages']` is a list of BaseMessage, possibly ending in an AIMessage
-    (or plan_node may have set final_answer directly). We append any tool results
-    as SystemMessages, then prompt the LLM for one final answer.
+    By this time:
+      - state['messages'] is a list of BaseMessage (SystemMessage/HumanMessage/AIMessage).
+      - Possibly state['web_search_result'] or state['ocr_result'] or state['excel_result'] is set.
+      - Or state['final_answer'] is already set (if plan_node decided no tool was needed).
+
+    We append any tool results as SystemMessages, then prompt the LLM for one final answer.
     """
     # 1) Copy the existing BaseMessage list
     history = state.get("messages", []).copy()
 
-    # 2) If any tool-result fields exist, append them as SystemMessages
+    # 2) Append each tool result as a SystemMessage, if present
     if "web_search_result" in state and state["web_search_result"] is not None:
         history.append(SystemMessage(content=f"WEB_SEARCH_RESULT: {state['web_search_result']}"))
     if "ocr_result" in state and state["ocr_result"] is not None:
@@ -112,39 +114,32 @@ def finalize_node(state: AgentState) -> AgentState:
     if "excel_result" in state and state["excel_result"] is not None:
         history.append(SystemMessage(content=f"EXCEL_RESULT: {state['excel_result']}"))
 
-    # 3) If plan_node already set final_answer, just return it:
+    # 3) If plan_node already set a final_answer, just return it directly
     if state.get("final_answer") is not None:
         return {"final_answer": state["final_answer"]}
 
-    # 4) Otherwise, ask the LLM to give the final answer now
+    # 4) Otherwise, ask the LLM to produce the final answer
     history.append(SystemMessage(content="Please provide the final answer now."))
     llm_response = llm(history)
     return {"final_answer": llm_response.content.strip()}
 
+# ─── 4) Wrap the low‐level tool wrappers in a ToolNode ───
 tool_node = ToolNode([web_search_tool, ocr_image_tool, parse_excel_tool])
 
-# ─── 5) Build the StateGraph ───
+# ─── 5) Build and compile the StateGraph ───
 graph = StateGraph(AgentState)
 
-# 5.a) Register nodes
+# 5.a) Register each node
 graph.add_node("plan", plan_node)
 graph.add_node("tools", tool_node)
 graph.add_node("run_tools", run_tools)
 graph.add_node("finalize", finalize_node)
 
-# 5.b) START → plan
+# 5.b) Wire START → plan
 graph.add_edge(START, "plan")
 
-
-
-
-
-
+# 5.c) plan → conditional: if any tool key is set, go to "tools"; otherwise "finalize"
 def route_plan(plan_out: AgentState) -> str:
-    """
-    plan_out is exactly what plan_node returned (a partial AgentState).
-    If it set any of the tool-request keys, route to 'tools'; otherwise 'finalize'.
-    """
     if plan_out.get("web_search_query") or plan_out.get("ocr_path") or plan_out.get("excel_path"):
         return "tools"
     return "finalize"
@@ -155,57 +150,57 @@ graph.add_conditional_edges(
     {"tools": "tools", "finalize": "finalize"}
 )
 
-
-
+# 5.d) Wire tools → run_tools
 graph.add_edge("tools", "run_tools")
 
-# 5.e) run_tools → finalize
+# 5.e) Wire run_tools → finalize
 graph.add_edge("run_tools", "finalize")
 
-# 5.f) finalize → END
+# 5.f) Wire finalize → END
 graph.add_edge("finalize", END)
 
 compiled_graph = graph.compile()
 
+# ─── 6) Define respond_to_input ───
 def respond_to_input(user_input: str) -> str:
     """
-    Initialize with a SystemMessage (tools description) and the user’s question as a HumanMessage.
-    Then run through plan → tools → run_tools → finalize. Return the "final_answer" from final_state.
+    Start with a SystemMessage + HumanMessage; then let the graph run:
+    plan_node → tools → run_tools → finalize_node. Return final_answer.
     """
-    # 1) Create a SystemMessage that tells the agent about its tools
+    # 1) SystemMessage describing the tools
     system_msg = SystemMessage(
         content=(
             "You have access to exactly these tools:\n"
-            "  1) web_search(query:str) → Returns the top search results for the query.\n"
-            "  2) parse_excel(path:str, sheet_name:str) → Reads an Excel file and returns its contents.\n"
-            "  3) ocr_image(path:str) → Runs OCR on an image and returns any detected text.\n\n"
-            "If you need a tool, set exactly one of these keys in a Python‐dict:\n"
+            "  1) web_search(query:str) → Returns DuckDuckGo results.\n"
+            "  2) parse_excel(path:str, sheet_name:str) → Reads an Excel file.\n"
+            "  3) ocr_image(path:str) → Runs OCR on an image.\n\n"
+            "If you need a tool, set exactly one of these keys in a Python dict:\n"
             "  • web_search_query: <search terms>\n"
             "  • ocr_path: <path to image>\n"
             "  • excel_path: <path to xlsx>\n"
             "  • excel_sheet_name: <sheet name>\n"
             "Otherwise, set final_answer: <your answer>.\n"
-            "Respond with that Python dict literal—no extra text or explanation."
+            "Respond with only that Python dict literal—no extra text."
         )
     )
-
-    # 2) Wrap the user_input in a HumanMessage
+    # 2) HumanMessage wrapping the user’s question
     human_msg = HumanMessage(content=user_input)
 
-    # 3) Build the initial state so that "messages" contains both messages
-    initial_state: AgentState = {
-        "messages": [system_msg, human_msg],
-        "user_input": user_input
-    }
+    # 3) Build initial_state so that "messages" = [system_msg, human_msg]
+    initial_state: AgentState = {"messages": [system_msg, human_msg]}
 
-    # 4) Invoke the compiled graph (no second argument needed)
+    # 4) Invoke the graph (no second argument needed)
     final_state = compiled_graph.invoke(initial_state)
 
-    # 5) Return the final answer (or a fallback if missing)
+    # 5) Return the "final_answer" or a fallback
     return final_state.get("final_answer", "Error: No final answer generated.")
 
-
-
+# ─── 7) BasicAgent wrapper ───
+class BasicAgent:
+    def __init__(self):
+        print("BasicAgent initialized.")
+    def __call__(self, question: str) -> str:
+        return respond_to_input(question)
 class BasicAgent:
     def __init__(self):
         print("BasicAgent initialized.")