Update app.py

app.py

@@ -2,65 +2,74 @@ import os
 import gradio as gr
 import requests
 import pandas as pd
-from typing import TypedDict, Annotated
-from langchain_core.messages import BaseMessage
-from langchain.agents import create_react_agent, AgentExecutor
-from langchain import hub
-from langchain_community.tools import DuckDuckGoSearchRun
+from typing import TypedDict, Annotated, Sequence
+import operator
+from langchain_core.messages import BaseMessage, HumanMessage, ToolMessage
+from langchain.agents import AgentExecutor
+from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_huggingface import HuggingFaceEndpoint
+from langgraph.graph import StateGraph, END
+from langgraph.prebuilt import ToolNode, tools_condition
 
 # --- Main Application Logic ---
 
-# This function builds our agent and the executor that runs it.
-def create_agent_executor():
-    print("Initializing ReAct Agent...")
+# This defines the "memory" or state of our agent.
+class AgentState(TypedDict):
+    messages: Annotated[Sequence[BaseMessage], operator.add]
+
+# This function builds our final, robust agent using LangGraph.
+def create_langgraph_agent():
+    print("Initializing LangGraph Agent...")
     
     # 1. Set up the LLM (The "Brain")
+    # We use the powerful Mistral model with the correct 'conversational' task
     llm = HuggingFaceEndpoint(
         repo_id="mistralai/Mistral-7B-Instruct-v0.2",
         task="conversational",
         max_new_tokens=512,
         do_sample=False,
     )
-    print("LLM initialized.")
+    # We bind the tools to the LLM so it knows about them
+    tools = [TavilySearchResults(max_results=3)]
+    llm_with_tools = llm.bind_tools(tools)
+    print("LLM and tools initialized.")
 
-    # 2. Define the Tools
-    tools = [DuckDuckGoSearchRun()]
-    print("Tools initialized.")
+    # 2. Define the Graph Nodes
+    # The 'agent' node calls the LLM
+    def agent_node(state):
+        print("Calling agent node...")
+        response = llm_with_tools.invoke(state["messages"])
+        return {"messages": [response]}
     
-    # 3. Get the ReAct Prompt Template
-    # This prompt is specifically designed to make models think step-by-step
-    prompt = hub.pull("hwchase17/react-chat")
-    print("Prompt template pulled.")
-
-    # 4. Create the Agent's logic using the ReAct framework
-    agent_runnable = create_react_agent(llm, tools, prompt)
-    print("Agent logic created.")
+    # The 'tool' node executes the tools
+    tool_node = ToolNode(tools)
+    print("Graph nodes defined.")
 
-    # 5. Create the Agent Executor
-    # This is the runtime that will loop through the agent's thoughts and tool uses.
-    # We are going back to the standard AgentExecutor, as the LangGraph build was overly complex
-    # and causing the new errors.
-    agent_executor = AgentExecutor(
-        agent=agent_runnable,
-        tools=tools,
-        verbose=True,
-        handle_parsing_errors=True,
-        max_iterations=5,
+    # 3. Define the Graph
+    graph = StateGraph(AgentState)
+    graph.add_node("agent", agent_node)
+    graph.add_node("tools", tool_node)
+    
+    graph.set_entry_point("agent")
+    
+    # This conditional edge decides whether to call a tool or end
+    graph.add_conditional_edges(
+        "agent",
+        tools_condition,
     )
-    print("Agent Executor created. Initialization complete.")
-    return agent_executor
+    graph.add_edge("tools", "agent")
+
+    # 4. Compile the graph into a runnable app
+    app = graph.compile()
+    print("LangGraph agent compiled and ready.")
+    return app
 
 # This function runs the agent for a single question.
-def run_agent_for_task(agent_executor, question: str) -> str:
+def run_agent(agent_executor, question: str) -> str:
     print(f"Agent received question: {question}")
     try:
-        # For this agent, the input is a dictionary with "input" and an empty "chat_history"
-        response = agent_executor.invoke({
-            "input": question,
-            "chat_history": []
-        })
-        final_answer = response.get("output", "Error: Could not parse final answer.")
+        response = agent_executor.invoke({"messages": [HumanMessage(content=question)]})
+        final_answer = response['messages'][-1].content
     except Exception as e:
         print(f"Error during agent execution: {e}")
         final_answer = f"Error: Agent failed to execute. {e}"
@@ -68,7 +77,7 @@ def run_agent_for_task(agent_executor, question: str) -> str:
     print(f"Agent returning answer: {final_answer}")
     return str(final_answer)
 
-# The rest of the file runs the evaluation and is mostly unchanged from the template.
+# The rest of the file runs the evaluation
 def run_and_submit_all(profile: gr.OAuthProfile | None):
     space_id = os.getenv("SPACE_ID")
     if not profile:
@@ -76,7 +85,7 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
     username = f"{profile.username}"
     
     try:
-        agent_executor = create_agent_executor()
+        agent_executor = create_langgraph_agent()
     except Exception as e:
         return f"Error initializing agent: {e}", None
         
@@ -94,7 +103,7 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
     for item in questions_data:
         task_id, question_text = item.get("task_id"), item.get("question")
         if task_id and question_text:
-            submitted_answer = run_agent_for_task(agent_executor, question_text)
+            submitted_answer = run_agent(agent_executor, question_text)
             answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
 
     submission_data = {"username": username.strip(), "agent_code": agent_code, "answers": answers_payload}
@@ -114,9 +123,9 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
     except Exception as e:
         return f"Error during submission: {e}", pd.DataFrame(answers_payload)
 
-# --- Gradio Interface ---
+# Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Agent Evaluation Runner (Final ReAct Version)")
+    gr.Markdown("# Agent Evaluation Runner (Final LangGraph v2)")
     gr.LoginButton()
     run_button = gr.Button("Run Evaluation & Submit All Answers")
     status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)