Update app.py

app.py

@@ -4,8 +4,8 @@ import requests
 import pandas as pd
 from typing import TypedDict, Annotated, Sequence
 import operator
-from langchain_core.messages import BaseMessage, HumanMessage
-from langchain_community.tools import DuckDuckGoSearchRun
+from langchain_core.messages import BaseMessage, HumanMessage, ToolMessage
+from langchain.tools import tool
 from langchain_openai import ChatOpenAI
 from langgraph.graph import StateGraph, END
 from langgraph.prebuilt import ToolNode, tools_condition
@@ -16,22 +16,68 @@ from langgraph.prebuilt import ToolNode, tools_condition
 class AgentState(TypedDict):
     messages: Annotated[Sequence[BaseMessage], operator.add]
 
-# This function builds our final, robust agent using LangGraph and OpenAI
+# NEW: Custom tool for downloading files from the course API
+@tool
+def download_file(task_id: str) -> str:
+    """
+    Downloads a file associated with a specific task_id from the course API.
+    The file will be saved in the /tmp/ directory.
+    Use this tool when a question mentions an attached file (e.g., image, excel, audio).
+    """
+    try:
+        # NOTE: This assumes the Space is running on the Hugging Face platform
+        # and can access the scoring API directly.
+        file_url = f"https://agents-course-unit4-scoring.hf.space/files/{task_id}"
+        response = requests.get(file_url)
+        response.raise_for_status()
+        
+        # Extract filename from headers, or default to task_id
+        content_disposition = response.headers.get('content-disposition')
+        if content_disposition:
+            parts = content_disposition.split(';')
+            for part in parts:
+                if 'filename=' in part:
+                    filename = part.split('=')[1].strip('"')
+                    break
+        else:
+            # Fallback filename if header is not present
+            # This is a simplification; a real app might need better content-type handling
+            filename = task_id 
+            
+        # Save the file to a temporary directory within the Space
+        save_path = f"/tmp/{filename}"
+        with open(save_path, "wb") as f:
+            f.write(response.content)
+            
+        return f"Successfully downloaded file for task {task_id}. It is available at path: {save_path}"
+    except Exception as e:
+        return f"An error occurred while downloading the file: {str(e)}"
+
+# This function builds our final, robust agent.
 def create_langgraph_agent():
     print("Initializing LangGraph Agent with OpenAI...")
     
-    # 1. Set up the LLM (The "Brain") using OpenAI's GPT-3.5 Turbo
-    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
+    # 1. Set up the LLM (The "Brain")
+    llm = ChatOpenAI(model="gpt-4o", temperature=0) # Switched to gpt-4o for better performance
     
-    # 2. Define the Tools, now using DuckDuckGo
-    tools = [DuckDuckGoSearchRun()]
+    # 2. Define the Tools, including our new file downloader
+    tools = [TavilySearchResults(max_results=3), download_file]
     llm_with_tools = llm.bind_tools(tools)
     print("LLM and tools initialized.")
 
-    # 3. Define the Graph Nodes
+    # 3. Define the agent's logic (the "agent" node)
+    # NEW: We are using a much more specific system prompt based on the GAIA paper
+    system_prompt = """You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER].
+YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
+If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise.
+If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise.
+If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string."""
+
     def agent_node(state):
         print("Calling agent node...")
-        response = llm_with_tools.invoke(state["messages"])
+        # Add the system prompt to the beginning of the conversation
+        messages_with_system_prompt = [("system", system_prompt)] + state["messages"]
+        response = llm_with_tools.invoke(messages_with_system_prompt)
         return {"messages": [response]}
     
     tool_node = ToolNode(tools)
@@ -41,7 +87,6 @@ def create_langgraph_agent():
     graph = StateGraph(AgentState)
     graph.add_node("agent", agent_node)
     graph.add_node("tools", tool_node)
-    
     graph.set_entry_point("agent")
     graph.add_conditional_edges("agent", tools_condition)
     graph.add_edge("tools", "agent")
@@ -51,12 +96,20 @@ def create_langgraph_agent():
     print("LangGraph agent compiled and ready.")
     return app
 
-# This function runs the agent for a single question.
-def run_agent(agent_executor, question: str) -> str:
+# This function runs the agent and extracts the final answer.
+def run_agent(agent_executor, question: str, task_id: str) -> str:
     print(f"Agent received question: {question}")
     try:
-        response = agent_executor.invoke({"messages": [HumanMessage(content=question)]})
-        final_answer = response['messages'][-1].content
+        # Pass the task_id to the agent's input so the file downloader can use it
+        response = agent_executor.invoke({"messages": [HumanMessage(content=f"Task ID: {task_id}\n\nQuestion: {question}")]})
+        
+        # NEW: Extract only the part after "FINAL ANSWER:"
+        raw_answer = response['messages'][-1].content
+        if "FINAL ANSWER:" in raw_answer:
+            final_answer = raw_answer.split("FINAL ANSWER:")[-1].strip()
+        else:
+            final_answer = raw_answer # Fallback if the model doesn't follow the format
+            
     except Exception as e:
         print(f"Error during agent execution: {e}")
         final_answer = f"Error: Agent failed to execute. {e}"
@@ -71,8 +124,9 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
         return "Please Login to Hugging Face with the button.", None
     username = f"{profile.username}"
     
-    if not os.getenv("OPENAI_API_KEY"):
-         return "OpenAI API key not found. Please set the OPENAI_API_KEY secret in your Space settings.", None
+    # API key checks
+    if not os.getenv("TAVILY_API_KEY") or not os.getenv("OPENAI_API_KEY"):
+         return "One or more API keys (TAVILY_API_KEY, OPENAI_API_KEY) are not set. Please set them in your Space secrets.", None
 
     try:
         agent_executor = create_langgraph_agent()
@@ -93,13 +147,13 @@ 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(agent_executor, question_text)
+            submitted_answer = run_agent(agent_executor, question_text, task_id)
             answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
 
     submission_data = {"username": username.strip(), "agent_code": agent_code, "answers": answers_payload}
     submit_url = f"https://agents-course-unit4-scoring.hf.space/submit"
     try:
-        response = requests.post(submit_url, json=submission_data, timeout=120)
+        response = requests.post(submit_url, json=submission_data, timeout=180) # Increased timeout for gpt-4
         response.raise_for_status()
         result_data = response.json()
         final_status = (
@@ -115,7 +169,7 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
 
 # Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Agent Evaluation Runner (OpenAI + DuckDuckGo)")
+    gr.Markdown("# Agent Evaluation Runner (GPT-4o + LangGraph)")
     gr.LoginButton()
     run_button = gr.Button("Run Evaluation & Submit All Answers")
     status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)