Update app.py

app.py

@@ -5,61 +5,87 @@ import inspect
 import pandas as pd
 
 
+from typing import Sequence, Annotated, TypedDict, Union
 from langgraph.graph import StateGraph, END
-from langchain_core.runnables import Runnable
-from langchain_core.agents import AgentFinish
-from langchain.agents import create_react_agent
-from langchain.agents.agent import RunnableAgent
-from langchain.tools import Tool
-from langchain_core.messages import HumanMessage
+from langgraph.graph.message import add_messages
+from langgraph.prebuilt import ToolNode
+from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage
 from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_community.tools import DuckDuckGoSearchRun
+from langchain.tools import tool
+from dotenv import load_dotenv
+load_dotenv()
+# ----------- Agent State -----------
+class AgentState(TypedDict):
+    messages: Annotated[Sequence[BaseMessage], add_messages]
 
 
-# (Keep Constants as is)
-# --- Constants ---
-DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
+# ----------- Math Tools ------------
+@tool
+def add(a: int, b: int):
+    """Adds two numbers."""
+    return a + b
 
-GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
+@tool
+def subtract(a: int, b: int):
+    """Subtracts two numbers."""
+    return a - b
 
+@tool
+def multiply(a: int, b: int):
+    """Multiplies two numbers."""
+    return a * b
 
-# --- Basic Agent Definition ---
-# ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
-class BasicAgent:
-    def __init__(self):
-        print("LangGraph ReAct Agent initialized.")
-        if not GOOGLE_API_KEY:
-            raise ValueError("GOOGLE_API_KEY environment variable not found.")
-
-        self.llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash")
+# ----------- DuckDuckGo Tool (LangChain built-in) -----------
+ddg_tool = DuckDuckGoSearchRun(name="duckduckgo_search")
 
+# ----------- Combine all tools -----------
+tools = [add, subtract, multiply, ddg_tool]
 
-        # Example tools
-        self.tools = [
-            Tool(
-                name="DuckDuckGo Search",
-                func=lambda q: f"Pretend search result for: {q}",
-                description="Search the web using DuckDuckGo. Useful for recent events."
-            )
-        ]
-
-        react_agent = create_react_agent(self.llm, self.tools)
 
-        workflow = StateGraph({"agent_outcome": Runnable})
-        workflow.add_node("agent", react_agent)
-        workflow.set_entry_point("agent")
-        workflow.add_edge("agent", END)
+# ----------- BasicAgent Class -----------
+class BasicAgent:
+    def __init__(self):
+        print("LangGraph Gemini Agent with LangChain DuckDuckGo Tool Initialized.")
 
-        app = workflow.compile()
-        self.agent = RunnableAgent(app)
+        self.model = ChatGoogleGenerativeAI(model="gemini-2.0-flash").bind_tools(tools)
 
-    def __call__(self, question: str) -> str:
-        result = self.agent.invoke({"input": HumanMessage(content=question)})
-        if isinstance(result, AgentFinish):
-            return result.return_values["output"]
-        elif isinstance(result, dict) and "output" in result:
-            return result["output"]
+        def model_call(state: AgentState) -> AgentState:
+            system_prompt = SystemMessage(
+                content="You are an AI assistant. Use tools like math functions and web search (DuckDuckGo) to answer queries."
+            )
+            response = self.model.invoke([system_prompt] + state["messages"])
+            return {"messages": [response]}
+
+        def should_continue(state: AgentState):
+            last_message = state["messages"][-1]
+            if not getattr(last_message, "tool_calls", None):
+                return "end"
+            return "continue"
+
+        graph = StateGraph(AgentState)
+        graph.add_node("our_agent", model_call)
+        graph.add_node("tools", ToolNode(tools=tools))
+        graph.set_entry_point("our_agent")
+        graph.add_conditional_edges("our_agent", should_continue, {
+            "continue": "tools",
+            "end": END
+        })
+        graph.add_edge("tools", "our_agent")
+
+        self.agent = graph.compile()
+
+    def __call__(self, query: Union[str, Sequence[BaseMessage]]) -> str:
+        if isinstance(query, str):
+            messages = [HumanMessage(content=query)]
+        elif isinstance(query, list):
+            messages = query
         else:
-            return "No valid output returned."
+            raise ValueError("Invalid input: Must be a string or list of messages.")
+
+        result = self.agent.invoke({"messages": messages})
+        last_message = result["messages"][-1]
+        return last_message.content if hasattr(last_message, "content") else str(last_message)
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """