Update app.py

app.py

@@ -1,17 +1,17 @@
 import os
 import gradio as gr
 import requests
-import inspect
 import pandas as pd
 
-from typing import TypedDict, Annotated
-from langgraph.graph.message import add_messages
-from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, SystemMessage
-from langgraph.prebuilt import ToolNode
-from langgraph.graph import START, StateGraph
-from langgraph.prebuilt import tools_condition
 from langchain_openai import ChatOpenAI
-from langchain_community.tools import DuckDuckGoSearchRun
+from langchain.agents import create_agent
+from langchain.messages import HumanMessage
+from langchain.tools import tool
+from typing import Dict, Any
+from tavily import TavilyClient
+
+from dotenv import load_dotenv
+load_dotenv()
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -28,45 +28,33 @@ Your final answer must include only the answer to the user's question, without a
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
 
 # Generate the AgentState and Agent graph
-class AgentState(TypedDict):
-    messages: Annotated[list[AnyMessage], add_messages]
 
 class BasicAgent:
     def __init__(self):
         self.llm = ChatOpenAI(
             model="nvidia/nemotron-3-super-120b-a12b:free",
             base_url="https://openrouter.ai/api/v1",
-            api_key=os.environ.get("OPENROUTER_API_KEY")
+            api_key=os.getenv("OPENROUTER_API_KEY")
         )
-        search_tool = DuckDuckGoSearchRun()
-        tools = [search_tool]
-        self.llm = self.llm.bind_tools(tools)
-
-        builder = StateGraph(AgentState)
-
-        builder.add_node("assistant", self.assistant)
-        builder.add_node("tools", ToolNode(tools))
-
-        builder.add_edge(START, "assistant")
-        builder.add_conditional_edges(
-            "assistant", 
-            tools_condition
+        self.tavily_client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
+        tools = [self.web_search]
+        self.agent = create_agent(
+            model=self.llm,
+            tools=tools,
+            system_prompt=SYSTEM_PROMPT
         )
-        builder.add_edge("tools", "assistant")
-        self.agent = builder.compile()
         print("BasicAgent initialized.")
-    def assistant(self, state: AgentState) -> dict:
-        return {
-            "messages": [self.llm.invoke(state["messages"])],
-        }
+    @tool
+    def web_search(self, query: str) -> Dict[str, Any]:
+        """Search the web for information"""
+        return self.tavily_client.search(query)
     def __call__(self, question: str) -> str:
-        print(f"Agent received question: {question}")
-        messages = [SystemMessage(content=SYSTEM_PROMPT), HumanMessage(content=question)]
-        initial_state = AgentState(messages=messages)
-        response = self.agent.invoke(initial_state)
-        final_answer = response['messages'][-1].content
-        print(f"Agent's answer: {final_answer}")
-        return final_answer
+        user_prompt = HumanMessage(content=question)
+        response = self.agent.invoke(
+            {"messages": [user_prompt]}
+        )
+        answer = response['messages'][-1].content.strip()
+        return answer
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
@@ -122,13 +110,14 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
     results_log = []
     answers_payload = []
     print(f"Running agent on {len(questions_data)} questions...")
-    for item in questions_data:
+    for i, item in enumerate(questions_data):
         task_id = item.get("task_id")
         question_text = item.get("question")
         if not task_id or question_text is None:
             print(f"Skipping item with missing task_id or question: {item}")
             continue
         try:
+            print(f"Answering question {i+1}/{len(questions_data)}")
             submitted_answer = agent(question_text)
             answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
             results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})