Update app.py

app.py

@@ -260,8 +260,10 @@ import operator
 # --- LangChain & LangGraph Imports ---
 from langchain_core.messages import BaseMessage, HumanMessage, ToolMessage, AIMessage, SystemMessage
 from langchain_core.tools import tool
-# <<<--- CHANGE: Import the HuggingFaceEndpoint for open-source models --->>>
 from langchain_huggingface import HuggingFaceEndpoint
+# <<<--- CHANGE 1: Import new components for building the agent --->>>
+from langchain.agents import AgentExecutor, create_tool_calling_agent
+from langchain_core.prompts import ChatPromptTemplate
 from langgraph.graph import StateGraph, END
 from langgraph.prebuilt import ToolNode
 from tavily import TavilyClient
@@ -272,7 +274,7 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 FILES_DIR = "./files"
 os.makedirs(FILES_DIR, exist_ok=True)
 
-# --- System Prompt (Unchanged, it's strong and model-agnostic) ---
+# --- System Prompt (Unchanged) ---
 AGENT_SYSTEM_PROMPT = """You are a world-class AI agent, specialized in solving complex problems from the GAIA benchmark.
 Your task is to analyze the user's question, think step-by-step, and use the provided tools to find the correct answer.
 CRITICAL INSTRUCTIONS:
@@ -298,9 +300,7 @@ tavily = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
 
 @tool
 def tavily_search(query: str) -> str:
-    """
-    Uses the Tavily Search API to find information on the web.
-    """
+    """Uses the Tavily Search API to find information on the web."""
     print(f"--- Calling Tavily Search Tool with query: {query} ---")
     try:
         result = tavily.search(query=query, search_depth="advanced")
@@ -310,108 +310,75 @@ def tavily_search(query: str) -> str:
 
 @tool
 def read_file(url: str) -> str:
-    """
-    Downloads a file from a given URL, saves it locally, and returns its content.
-    It can handle both plain text files and PDF files.
-    """
+    """Downloads and reads the content of a file (text or PDF) from a URL."""
     print(f"--- Calling Read File Tool with URL: {url} ---")
     try:
         filename = os.path.join(FILES_DIR, os.path.basename(url))
         response = requests.get(url)
         response.raise_for_status()
-        
-        with open(filename, 'wb') as f:
-            f.write(response.content)
+        with open(filename, 'wb') as f: f.write(response.content)
             
         if url.lower().endswith('.pdf'):
-            print(f"--- File identified as PDF. Reading with pypdf. ---")
             try:
                 pdf_reader = pypdf.PdfReader(filename)
-                content = ""
-                for page in pdf_reader.pages:
-                    content += page.extract_text()
-                return f"Successfully read PDF file '{filename}'. Content:\n\n{content}"
-            except Exception as e:
-                return f"Error reading PDF file: {e}"
+                return f"Successfully read PDF file '{filename}'. Content:\n\n{''.join(p.extract_text() for p in pdf_reader.pages)}"
+            except Exception as e: return f"Error reading PDF file: {e}"
         else:
-            print(f"--- File identified as text. Reading normally. ---")
             try:
-                with open(filename, 'r', encoding='utf-8') as f:
-                    content = f.read()
-                return f"Successfully read text file '{filename}'. Content:\n\n{content}"
-            except UnicodeDecodeError:
-                return f"Successfully downloaded binary file '{filename}'. Cannot display content as text."
-
-    except requests.exceptions.RequestException as e:
-        return f"Error downloading or reading file: {e}"
+                with open(filename, 'r', encoding='utf-8') as f: return f"Successfully read text file '{filename}'. Content:\n\n{f.read()}"
+            except UnicodeDecodeError: return f"Successfully downloaded binary file '{filename}'. Cannot display content as text."
+    except requests.exceptions.RequestException as e: return f"Error downloading or reading file: {e}"
 
 @tool
 def python_interpreter(code: str) -> str:
-    """
-    Executes a given string of Python code and returns the output from stdout.
-    """
+    """Executes Python code and returns its stdout."""
     print(f"--- Calling Python Interpreter Tool with code:\n{code} ---")
     output_buffer = io.StringIO()
     try:
-        with redirect_stdout(output_buffer):
-            exec(code, globals())
+        with redirect_stdout(output_buffer): exec(code, globals())
         return f"Code executed successfully. Output:\n{output_buffer.getvalue()}"
-    except Exception as e:
-        return f"Error executing Python code: {e}"
+    except Exception as e: return f"Error executing Python code: {e}"
 
 #
 # ================================================================================================
-#  ✅ 2. CONFIGURE AND BUILD THE AGENT GRAPH (WITH HUGGING FACE)
+#  ✅ 2. CONFIGURE AND BUILD THE AGENT GRAPH (CORRECTED IMPLEMENTATION)
 # ================================================================================================
 #
 class AgentState(TypedDict):
-    messages: Annotated[List[BaseMessage], operator.add]
-
+    # <<<--- CHANGE 2: The state is now simpler. It tracks the input and the agent's output. --->>>
+    input: str
+    agent_outcome: dict | None
+    
 def build_agent_graph():
     """Builds the LangGraph agent."""
     tools = [tavily_search, read_file, python_interpreter]
     
-    # <<<--- CHANGE: Instantiate the Hugging Face Model Endpoint --->>>
-    # This uses the recommended Command R+ model for its excellent tool-use capabilities.
-    # It will automatically use the HUGGINGFACEHUB_API_TOKEN secret.
     repo_id = "CohereForAI/c4ai-command-r-plus"
     llm = HuggingFaceEndpoint(
-        repo_id=repo_id,
-        max_new_tokens=1024,
-        temperature=0, # Keep temperature low for fact-based tasks
+        repo_id=repo_id, max_new_tokens=1024, temperature=0.1,
         huggingfacehub_api_token=os.getenv("HUGGINGFACEHUB_API_TOKEN")
     )
     
-    llm_with_tools = llm.bind_tools(tools)
-
-    def call_model(state: AgentState) -> dict:
-        """Helper function to prepare messages and call the model."""
-        messages = state['messages']
-        # The HuggingFaceEndpoint doesn't support a separate SystemMessage.
-        # We'll format the system prompt and the latest human message together.
-        if isinstance(messages[0], SystemMessage):
-            # Start with the system message content
-            formatted_messages = [HumanMessage(content=messages[0].content + "\n\nHere is the user's question:\n" + messages[-1].content)]
-            # Add any previous tool outputs
-            formatted_messages.extend(messages[1:-1])
-        else:
-            formatted_messages = messages
-
-        response = llm_with_tools.invoke(formatted_messages)
-        return {"messages": [response]}
-
-    def should_continue(state: AgentState) -> str:
-        """Determines whether to continue the loop or end."""
-        return "action" if state['messages'][-1].tool_calls else "end"
+    # <<<--- CHANGE 3: Create a proper agent prompt template --->>>
+    # This prompt is specifically designed to instruct the model on how to use tools.
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", AGENT_SYSTEM_PROMPT),
+        ("human", "{input}"),
+        # The 'agent_scratchpad' placeholder is crucial for the agent to remember previous tool calls.
+        ("placeholder", "{agent_scratchpad}"), 
+    ])
+
+    # <<<--- CHANGE 4: Create the tool-calling agent runnable --->>>
+    agent = create_tool_calling_agent(llm, tools, prompt)
+    
+    # <<<--- CHANGE 5: Create the AgentExecutor which will run the agent loop --->>>
+    # This replaces the manual LangGraph loop for calling the model and tools.
+    # It's a robust, pre-built component for this exact purpose.
+    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
 
-    tool_node = ToolNode(tools)
-    workflow = StateGraph(AgentState)
-    workflow.add_node("agent", call_model)
-    workflow.add_node("action", tool_node)
-    workflow.set_entry_point("agent")
-    workflow.add_conditional_edges("agent", should_continue, {"action": "action", "end": END})
-    workflow.add_edge('action', 'agent')
-    return workflow.compile()
+    # We are not using LangGraph for the agent loop anymore, as AgentExecutor handles it.
+    # We just need a simple callable class that invokes it.
+    return agent_executor
 
 #
 # ================================================================================================
@@ -420,28 +387,26 @@ def build_agent_graph():
 #
 class GaiaAgent:
     def __init__(self):
-        # <<<--- CHANGE: Update print statement for new model --->>>
-        print("GaiaAgent initialized. Building fresh Command R+ agent graph...")
+        print("GaiaAgent initialized. Building Command R+ agent with AgentExecutor...")
+        # The agent_app is now the fully-formed AgentExecutor
         self.agent_app = build_agent_graph()
 
     def __call__(self, question: str) -> str:
         print(f"\n{'='*60}\nAgent received question: {question[:100]}...\n{'='*60}")
-        initial_input = {
-            "messages": [
-                SystemMessage(content=AGENT_SYSTEM_PROMPT),
-                HumanMessage(content=question)
-            ]
-        }
-        final_state = None
-        for i, step in enumerate(self.agent_app.stream(initial_input, {"recursion_limit": 15})):
-            if i == 0: print("--- Starting Agentic Loop ---")
-            final_state = step
-
-        final_answer_message = final_state['agent']['messages'][-1]
-        final_answer = str(final_answer_message.content).strip()
-        print(f"\n--- Agent finished. Final Answer: {final_answer} ---\n")
-        return final_answer
+        
+        # <<<--- CHANGE 6: Invoke the AgentExecutor and extract the final answer --->>>
+        try:
+            # The AgentExecutor takes a dictionary and returns the final output in the 'output' key.
+            response = self.agent_app.invoke({"input": question})
+            final_answer = str(response.get("output", "")).strip()
+            print(f"\n--- Agent finished. Final Answer: {final_answer} ---\n")
+            return final_answer
+        except Exception as e:
+            print(f"An error occurred during agent execution: {e}")
+            return f"AGENT_EXECUTION_ERROR: {e}"
+
 
+# --- The rest of the file (run_and_submit_all, Gradio UI) remains the same ---
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     space_id = os.getenv("SPACE_ID")
     if not profile: return "Please Login to Hugging Face with the button.", None
@@ -466,13 +431,16 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
     answers_payload = []
     print(f"Running agent on {len(questions_data)} questions...")
 
+    # Instantiate the agent once to save time
+    agent_instance = GaiaAgent()
+
     for item in questions_data:
         task_id = item.get("task_id")
         question_text = item.get("question")
         if not task_id or question_text is None: continue
         try:
-            agent = GaiaAgent()
-            submitted_answer = agent(question_text)
+            # Reuse the same agent instance
+            submitted_answer = agent_instance(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})
         except Exception as e:
@@ -502,17 +470,15 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
         results_df = pd.DataFrame(results_log)
         return status_message, results_df
 
-# --- Gradio Interface ---
 with gr.Blocks() as demo:
-    # <<<--- CHANGE: Update UI titles and descriptions for the new model --->>>
-    gr.Markdown("# GAIA Agent Final Assessment (Open Source: Command R+)")
+    gr.Markdown("# GAIA Agent Final Assessment (Open Source: Command R+ - Corrected)")
     gr.Markdown(
         """
-        **Instructor's Note:** This version runs a top-tier open-source model from the Hugging Face Hub: **`CohereForAI/c4ai-command-r-plus`**.
-        This model is state-of-the-art for agentic tool use.
+        **Instructor's Note:** This version corrects the agent construction logic to be compatible with the `HuggingFaceEndpoint`.
+        It now uses the standard `create_tool_calling_agent` and `AgentExecutor` from LangChain for robust tool use.
         1.  Ensure you have a **`HUGGINGFACEHUB_API_TOKEN`** and a **`TAVILY_API_KEY`** set in your Space secrets.
-        2.  Ensure your `requirements.txt` includes `langchain-huggingface`.
-        3.  Good luck! Let's see how this powerful open model performs.
+        2.  Your `requirements.txt` should include `langchain`, `langchain-huggingface`, and `langchain-core`.
+        3.  Let's run the evaluation again!
         """
     )
     gr.LoginButton()