Update app.py

app.py

@@ -22,79 +22,118 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
 
 class AgentState(TypedDict):
+    # We store the full chat history as a list of strings.
     messages: Annotated[list[str], add_messages]
-    tool: str  # will store the name of the requested tool (if any)
-    agent_out: str  # raw output from the LLM
+    # If the agent requests a tool, it will fill in:
+    tool_request: dict | None
+    # Whenever a tool runs, its result goes here:
+    tool_result: str | None
 
-# 2) Instantiate the raw LLM and wrap it in a function
+# 2) Wrap ChatOpenAI in a function whose signature is (state, user_input) → new_state
 llm = ChatOpenAI(model_name="gpt-4.1-mini")
 
 def agent_node(state: AgentState, user_input: str) -> AgentState:
-    prev_msgs = state.get("messages", [])
-    messages = prev_msgs + [f"USER: {user_input}"]
-    # Ask the LLM for a response
-    llm_response = llm(messages).content  # returns a string or maybe JSON string
-    # If you expect JSON with {"tool": "...", ...}, parse it:
-    tool_requested = None
+    """
+    This function replaces raw ChatOpenAI. It must accept (state, user_input) 
+    and return a new AgentState dict.
+    """
+    # 2.a) Grab prior chat history (empty list on first turn)
+    prior_msgs = state.get("messages", [])
+    # 2.b) Append the new user_input
+    chat_history = prior_msgs + [f"USER: {user_input}"]
+    # 2.c) Ask the LLM for a response
+    llm_output = llm(chat_history).content
+
+    # 2.d) Check if the LLM output is valid Python dict literal indicating a tool call.
+    #      If it is, parse it and stash in state["tool_request"]. Otherwise, no tool.
+    tool_req = None
     try:
-        parsed = eval(llm_response)  # (use json.loads if the LLM returns valid JSON)
+        parsed = eval(llm_output)
         if isinstance(parsed, dict) and parsed.get("tool"):
-            tool_requested = parsed.get("tool")
-    except:
-        pass
+            tool_req = parsed
+    except Exception:
+        tool_req = None
 
+    # 2.e) Construct the new state:
     return {
-        "messages": messages + [f"ASSISTANT: {llm_response}"],
-        "agent_out": llm_response,
-        "tool": tool_requested or ""
+        "messages": chat_history + [f"ASSISTANT: {llm_output}"],
+        "tool_request": tool_req,
+        "tool_result": None  # will be filled by the tool_node if invoked
     }
 
-# 3) Instantiate a real ToolNode for your three tools
-t_node = ToolNode([ocr_image, parse_excel, web_search])
+# 3) Create a ToolNode for all three tools, then wrap it in a function
+#    whose signature is also (state, tool_request) → new_state.
+underlying_tool_node = ToolNode([ocr_image, parse_excel, web_search])
+
+def tool_node(state: AgentState, tool_request: dict) -> AgentState:
+    """
+    The graph will only call this when tool_request is a dict like
+      {"tool": "...", "path": "...", ...}
+    Use the underlying ToolNode to run it and store the result.
+    """
+    # 3.a) Run the actual ToolNode on that dict:
+    result_text = underlying_tool_node.run(tool_request)
 
-def run_tool_node(state: AgentState, agent_output) -> AgentState:
-    # `agent_output` is the dict that the LLM returned, e.g. {"tool":"ocr_image", "path": "file.png"}
-    tool_result: str = t_node.run(agent_output)
+    # 3.b) Update state.messages to note the tool’s output,
+    #      and clear tool_request so we don’t loop.
     return {
-        "messages": [f"TOOL RESULT: {tool_result}"],
-        "tool": "",  # once a tool has run, clear this so we don’t loop forever
-        "agent_out": tool_result
+        "messages": [f"TOOL ({tool_request['tool']}): {result_text}"],
+        "tool_request": None,
+        "tool_result": result_text
     }
 
-# 4) Build the StateGraph with the corrected node names
+# 4) Build and register nodes exactly as in the tutorial
 graph = StateGraph(AgentState)
 graph.add_node("agent", agent_node)
-graph.add_node("tools", run_tool_node)
+graph.add_node("tools", tool_node)
 
-# 5) START → "agent"
+# 5) Simple START → “agent” edge (no third argument needed)
 graph.add_edge(START, "agent")
 
-# 6) "tools" → "agent"
+# 6) Simple “tools” → “agent” edge (again, no third argument)
 graph.add_edge("tools", "agent")
 
-# 7) Conditional edges out of "agent"
-def route_agent(state: AgentState, agent_output):
-    # If LLM asked for a tool, we go to "tools"; else we terminate
-    if isinstance(agent_output, dict) and agent_output.get("tool") in {"ocr_image", "parse_excel", "web_search"}:
+# 7) Conditional branching out of “agent,” exactly like the tutorial
+def route_agent(state: AgentState, agent_out):
+    """
+    When the LLM (agent_node) runs, it returns an AgentState where
+    - state["tool_request"] is either a dict (if a tool was asked) or None.
+    - state["tool_result"] is always None on entry to agent_node.
+
+    route_agent must look at that returned state (called agent_out)
+    and decide:
+      • If agent_out["tool_request"] is not None, go to "tools".
+      • Otherwise, terminate (go to END).
+    """
+    if agent_out.get("tool_request") is not None:
         return "tools"
     return "final"
 
 graph.add_conditional_edges(
-    "agent",
-    route_agent,
+    "agent",       # source
+    route_agent,   # routing function (signature: (state, agent_out) → str key)
     {
-        "tools": "tools",
-        "final": END
+        "tools": "tools",  # if route_agent(...) == "tools", transition to node "tools"
+        "final": END       # if route_agent(...) == "final", stop execution
     }
 )
 
-# 8) Compile the graph and use run(…), not invoke(…)
+# 8) Compile the graph (now graph.run(...) will work)
 compiled_graph = graph.compile()
 
+# 9) Define respond_to_input so that Gradio (and the Hugging Face submission) can call it
 def respond_to_input(user_input: str) -> str:
-    initial_state: AgentState = {"messages": [], "tool": "", "agent_out": ""}
-    # Use .run() in v0.3.x; if you see an AttributeError, switch to .invoke()
-    return compiled_graph.invoke(initial_state, user_input)
+    # Start with an empty state
+    initial_state: AgentState = {
+        "messages": [],
+        "tool_request": None,
+        "tool_result": None
+    }
+    # Use .run(initial_state, user_input) in v0.3.x
+    final_state = compiled_graph.run(initial_state, user_input)
+    # The “final” on END means agent_out has no more tool calls and finished reasoning
+    # We return the last assistant message from state["messages"]:
+    return final_state["messages"][-1].replace("ASSISTANT: ", "")
 
 
 class BasicAgent: