Update app.py

app.py

@@ -13,11 +13,7 @@ from langgraph.graph import StateGraph, START, END
 from langgraph.graph.message import add_messages
 
 # Create a ToolNode that knows about your web_search function
-search_node = ToolNode([web_search])
 
-excel_tool_node = ToolNode([parse_excel])
-
-image_tool_node = ToolNode([ocr_image])
 # (Keep Constants as is)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
@@ -26,9 +22,79 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
 
 
-
-
-
+class AgentState(TypedDict):
+    messages: Annotated[list, add_messages]
+    question: str
+    answer: str
+
+llm = ChatOpenAI(model_name="gpt-4.1-mini")
+llm_node = LLMNode(llm)
+
+# 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])
+
+# image_tool_node = ToolNode([ocr_image])
+tool_node = ToolNode([ocr_image, parse_excel, web_search])
+
+# 5) Build the StateGraph
+graph = StateGraph[AgentState]()
+
+# ────────────────────────
+# Edge 1: START → LLM
+# Wrap the user_input into state.messages
+graph.add_edge(
+    START,
+    llm_node,
+    lambda state, user_input: {"messages": [user_input]},
+    name="start_to_llm",
+)
+
+# ────────────────────────
+# Edge 2: LLM → ToolNode
+# Only fire when the LLM returns a dict with exactly "tool":"ocr_image" or "tool":"parse_excel"
+# The lambda must return that dict so the ToolNode can extract its arguments.
+def route_to_tool(state: AgentState, llm_out):
+    # Expecting llm_out to be a dict like:
+    #   {"tool": "ocr_image", "path": "invoice.png"}
+    # or {"tool": "parse_excel", "path":"sales.xlsx", "sheet_name":"Sheet1"}
+    if isinstance(llm_out, dict) and llm_out.get("tool") in {"ocr_image", "parse_excel", "web_search"}:
+        return llm_out
+    return None  # → do not invoke the tool if it's not matching
+
+graph.add_edge(
+    llm_node,
+    tool_node,
+    route_to_tool,
+    name="llm_to_tool",
+)
+
+# ────────────────────────
+# Edge 3: ToolNode → LLM
+# Whatever the tool returns (a string), feed that straight back into the LLM as the next turn.
+graph.add_edge(
+    tool_node,
+    llm_node,
+    lambda state, tool_out: tool_out,
+    name="tool_to_llm",
+)
+
+# ────────────────────────
+# Edge 4: LLM → END
+# Once the LLM is done reasoning (without calling any more tools), return its output string.
+graph.add_edge(
+    llm_node,
+    END,
+    lambda state, final_str: final_str,
+    name="llm_to_end",
+)
+
+
+def respond_to_input(user_input: str) -> str:
+    initial_state: AgentState = {"messages": []}
+    return graph.run(initial_state, user_input)
 
 
 
@@ -36,10 +102,11 @@ class BasicAgent:
     def __init__(self):
         print("BasicAgent initialized.")
     def __call__(self, question: str) -> str:
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
-        fixed_answer = "This is a default answer."
-        print(f"Agent returning fixed answer: {fixed_answer}")
-        return fixed_answer
+        # print(f"Agent received question (first 50 chars): {question[:50]}...")
+        # fixed_answer = "This is a default answer."
+        # print(f"Agent returning fixed answer: {fixed_answer}")
+        return respond_to_input(question)
+        # return fixed_answer