Update app.py

app.py

@@ -21,65 +21,61 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
 
-
 class AgentState(TypedDict):
-    messages: Annotated[list, add_messages]
-    # question: str
-    # answer: str
+    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
 
+# 2) Instantiate the raw LLM and wrap it in a function
 llm = ChatOpenAI(model_name="gpt-4.1-mini")
 
-# 4) Wrap the tools in a single ToolNode
-#    You can list as many @tool functions here as you like.
-# search_node = ToolNode([web_search])
-
-# excel_tool_node = ToolNode([parse_excel])
+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
+    try:
+        parsed = eval(llm_response)  # (use json.loads if the LLM returns valid JSON)
+        if isinstance(parsed, dict) and parsed.get("tool"):
+            tool_requested = parsed.get("tool")
+    except:
+        pass
+
+    return {
+        "messages": messages + [f"ASSISTANT: {llm_response}"],
+        "agent_out": llm_response,
+        "tool": tool_requested or ""
+    }
 
-# image_tool_node = ToolNode([ocr_image])
+# 3) Instantiate a real ToolNode for your three tools
 t_node = ToolNode([ocr_image, parse_excel, web_search])
 
-def tool_node(state: AgentState, agent_output)-> AgentState:
-    """
-    Wrap ToolNode so it matches (state, agent_output) → next_state.
-    `agent_output` is expected to be a dict like {"tool": "...", ...}.
-    """
-    print(f"ToolNode received agent_output: {agent_output}")
-    # Let the ToolNode run with that dict (ToolNode.run(dict) returns a string)
-    tool_result: str = tool_node.run(agent_output)
-    # Now you might want the LLM to reason over tool_result, but for simplicity...
-    # We'll just store tool_result in messages.
-    return {"messages": [tool_result]}
-
+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)
+    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
+    }
 
-# 5) Build the StateGraph
+# 4) Build the StateGraph with the corrected node names
 graph = StateGraph(AgentState)
+graph.add_node("agent", agent_node)
+graph.add_node("tools", run_tool_node)
 
-# ────────────────────────
-graph.add_node("agent", llm)
-graph.add_node("tools", tool_node)
-# Edge A: START → "agent"
-#    Wrap the user_input into state["messages"]
-graph.add_edge(
-    START,
-    "agent"
-)
-
-# Edge C: "tools" → "agent"
-#    Whatever string the tool returns becomes the next prompt to the LLM
-graph.add_edge(
-    "tools",
-    "agent"
-)
+# 5) START → "agent"
+graph.add_edge(START, "agent")
 
-# 7) Use add_conditional_edges out of "agent" instead of two separate edges
+# 6) "tools" → "agent"
+graph.add_edge("tools", "agent")
 
-def route_agent(state: AgentState, agent_out):
-    """
-    Return "tools" if the LLM output is a dict with a valid "tool" key,
-    otherwise return "final".
-    """
-    print(f"Route agent received agent_out: {agent_out}")
-    if isinstance(agent_out, dict) and agent_out.get("tool") in {"ocr_image", "parse_excel", "web_search"}:
+# 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"}:
         return "tools"
     return "final"
 
@@ -87,20 +83,20 @@ graph.add_conditional_edges(
     "agent",
     route_agent,
     {
-        "tools": "tools",  # if route_agent(...) == "tools", go to node "tools"
-        "final": END       # if route_agent(...) == "final", go to END
+        "tools": "tools",
+        "final": END
     }
 )
 
+# 8) Compile the graph and use run(…), not invoke(…)
 compiled_graph = graph.compile()
+
 def respond_to_input(user_input: str) -> str:
-    print(f"Reached: User input: {user_input}")
-    initial_state: AgentState = {"messages": []}
-    print("reached respond_to_input")
+    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)
 
 
-
 class BasicAgent:
     def __init__(self):
         print("BasicAgent initialized.")