Update agent.py

agent.py

@@ -23,7 +23,7 @@ configure(api_key=os.getenv("GOOGLE_API_KEY"))
 #logger = logging.getLogger(__name__)
 
 # --- Model Configuration ---
-GEMINI_MODEL_NAME = "gemini/gemini-1.5-flash"
+GEMINI_MODEL_NAME = "gemini/gemini-2.0-flash"
 OPENAI_MODEL_NAME = "openai/gpt-4o"
 GROQ_MODEL_NAME = "groq/llama3-70b-8192"
 DEEPSEEK_MODEL_NAME = "deepseek/deepseek-chat"
@@ -123,6 +123,46 @@ class WikiContentFetcher(Tool):
         except wiki.exceptions.PageError:
             return f"'{page_title}' not found."
 
+class FileAttachmentQueryTool(Tool):
+    name = "run_query_with_file"
+    description = """
+    Downloads a file mentioned in a user prompt, adds it to the context, and runs a query on it.
+    This assumes the file is 20MB or less.
+    """
+    inputs = {
+        "task_id": {
+            "type": "string",
+            "description": "A unique identifier for the task related to this file, used to download it."
+        },
+        "mime_type": {
+            "type": "string",
+            "nullable": True,
+            "description": "The MIME type of the file, or the best guess if unknown."
+        },
+        "user_query": {
+            "type": "string",
+            "description": "The question to answer about the file."
+        }
+    }
+    output_type = "string"
+
+    def forward(self, task_id: str, mime_type: str | None, user_query: str) -> str:
+        file_url = f"https://agents-course-unit4-scoring.hf.space/files/{task_id}"
+        file_response = requests.get(file_url)
+        if file_response.status_code != 200:
+            return f"Failed to download file: {file_response.status_code} - {file_response.text}"
+        file_data = file_response.content
+        mime_type = mime_type or file_response.headers.get('Content-Type', 'application/octet-stream')
+
+        from google.generativeai import GenerativeModel
+        model = GenerativeModel(self.model_name)
+        response = model.generate_content([
+            types.Part.from_bytes(data=file_data, mime_type=mime_type),
+            user_query
+        ])
+
+        return response.text
+
 # --- Basic Agent Definition ---
 class BasicAgent:
     def __init__(self, provider="deepseek"):
@@ -137,6 +177,7 @@ class BasicAgent:
             MathSolver(),
             RiddleSolver(),
             TextTransformer(),
+            FileAttachmentQueryTool(model_name=GEMINI_MODEL_NAME),
         ]
         self.agent = ToolCallingAgent(
             model=model,
@@ -146,21 +187,24 @@ class BasicAgent:
         )
         self.agent.system_prompt = (
             """
-            You are a general AI assistant. I will ask you a question. 
-            YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. 
-            If your answer is a number and you are not explicitly asked for a string, write it in numerals instead of words, and don't use comma to write your number nor 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.
-
-            Answer questions as literally as you can, making as few assumptions as possible. Restrict the answer to the narrowest definition that still satifies the question.
-            If you are provied with a video, please watch and summarize the entire video before answering the question. The correct answer may be present only in a few frames of the video.
-            If you have difficulty finding an answer on Wikipedia, you may search the internet using Google Search or Duckduckgo search.
-            If you are asked to prove something, first state your assumptions and think step by step before giving your final answer.
-
-            Your final answer must strictly follow this format:
-            FINAL ANSWER: [ANSWER]
+            You are a GAIA benchmark AI assistant. Your sole purpose is to provide exact, minimal answers in the format 'FINAL ANSWER: [ANSWER]' with no additional text, explanations, or comments. 
+
+            - If the answer is a number, use numerals (e.g., '42', not 'forty-two'), without commas or units (e.g., no '$', '%') unless explicitly requested.
+            - If the answer is a string, use no articles ('a', 'the'), no abbreviations (e.g., 'New York', not 'NY'), and write digits as text (e.g., 'one', not '1') unless specified.
+            - For comma-separated lists, apply the above rules to each element based on whether it's a number or string.
+            - Answer as literally as possible, making minimal assumptions and adhering to the question's narrowest interpretation.
+            - For videos, analyze the entire content but extract only the precise answer to the query, ignoring irrelevant details.
+            - For Wikipedia or search tools, distill results to the minimal correct answer, ignoring extraneous content.
+            - If proving something, compute step-by-step internally but output only the final result in the required format.
+            - If tool outputs are verbose, extract only the essential answer that satisfies the question.
+            - Under no circumstances include explanations, intermediate steps, or text outside the 'FINAL ANSWER: [ANSWER]' format.
+
+            Example:
+            Question: What is 2 + 2?
+            Response: FINAL ANSWER: 4
 
-            Only write the answer in that exact format. Do not explain anything. Do not include any other text.
+            Your response must always be:
+            FINAL ANSWER: [ANSWER]
             """
         )
 
@@ -183,7 +227,40 @@ class BasicAgent:
             final_str = result["final_answer"].strip()
         else:
             final_str = str(result).strip()
-            
-        return final_str
 
+        return f"FINAL ANSWER: {final_str}"
+
+    def evaluate_random_questions(self, csv_path: str = "gaia_qa.csv", sample_size: int = 3, show_steps: bool = True):
+        df = pd.read_csv(csv_path)
+        if not {"question", "answer"}.issubset(df.columns):
+            print("CSV must contain 'question' and 'answer' columns.")
+            print("Found columns:", df.columns.tolist())
+            return
+        samples = df.sample(n=sample_size)
+        for _, row in samples.iterrows():
+            question = row["question"].strip()
+            expected = f"FINAL ANSWER: {str(row['answer']).strip()}"
+            result = self(question).strip()
+            if show_steps:
+                print("---")
+                print("Question:", question)
+                print("Expected:", expected)
+                print("Agent:", result)
+                print("Correct:", expected == result)
+            else:
+                print(f"Q: {question}\nE: {expected}\nA: {result}\n✓: {expected == result}\n")
+
+if __name__ == "__main__":
+    args = sys.argv[1:]
+    if not args or args[0] in {"-h", "--help"}:
+        print("Usage: python agent.py [question | dev]")
+        print(" - Provide a question to get a GAIA-style answer.")
+        print(" - Use 'dev' to evaluate 3 random GAIA questions from gaia_qa.csv.")
+        sys.exit(0)
 
+    q = " ".join(args)
+    agent = BasicAgent()
+    if q == "dev":
+        agent.evaluate_random_questions()
+    else:
+        print(agent(q))