Update agent.py

agent.py

@@ -1,80 +1,45 @@
 import os
-import io
-import base64
-import requests
-import pandas as pd
 from openai import OpenAI
 
-TEXT_ONLY_TASKS = {
-    "2d83110e-a098-4ebb-9987-066c06fa42d0",  # reversed question
-    "4fc2f1ae-8625-45b5-ab34-ad4433bc21f8",  # wikipedia FA
-    "6f37996b-2ac7-44b0-8e68-6d28256631b4",  # commutative check
-    "3cef3a44-215e-4aed-8e3b-b1e3f08063b7",  # grocery list - vegetables
-    "305ac316-eef6-4446-960a-92d80d542f82",  # actor - Magda M
-    "cf106601-ab4f-4af9-b045-5295fe67b37d",  # least athletes
-    "5a0c1adf-205e-4841-a666-7c3ef95def9d"   # Malko Competition
-}
-
-CSV_TASKS = {
-    "7bd855d8-463d-4ed5-93ca-5fe35145f733"  # Excel - food sales
-}
-
 class GaiaAgent:
     def __init__(self):
         self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
-        self.api_url = "https://agents-course-unit4-scoring.hf.space"
         self.instructions = (
-            "You are a precise assistant solving GAIA benchmark questions. "
-            "Only answer if you are confident you can provide the exact correct result."
+            "You are a research assistant solving GAIA benchmark questions using 2022 English Wikipedia knowledge.\n"
+            "For each question, reason step-by-step and only return the final answer in exact format."
         )
-
-    def fetch_file(self, task_id):
-        try:
-            url = f"{self.api_url}/files/{task_id}"
-            r = requests.get(url, timeout=10)
-            r.raise_for_status()
-            return r.content, r.headers.get("Content-Type", "")
-        except Exception as e:
-            return None, f"[FILE ERROR: {e}]"
-
-    def handle_csv_sales(self, csv_bytes):
-        try:
-            df = pd.read_excel(io.BytesIO(csv_bytes)) if csv_bytes[:4] == b"PK\x03\x04" else pd.read_csv(io.StringIO(csv_bytes.decode()))
-            if 'category' not in df.columns or 'sales' not in df.columns:
-                return "[MISSING COLUMN]"
-            food_df = df[df['category'].str.lower() == 'food']
-            if food_df.empty:
-                return "[NO FOOD ITEMS FOUND]"
-            total = food_df['sales'].sum()
-            return f"${total:.2f}"
-        except Exception as e:
-            return f"[CSV ERROR: {e}]"
+        self.task_templates = {
+            "8e867cd7-cff9-4e6c-867a-ff5ddc2550be": self.q_mercedes_sosa_albums
+        }
 
     def __call__(self, question: str, task_id: str = None) -> str:
-        # 1. Task filtering
-        if task_id not in TEXT_ONLY_TASKS and task_id not in CSV_TASKS:
-            return "[SKIPPED: Task not eligible for high-confidence answer]"
-
-        # 2. CSV handling
-        if task_id in CSV_TASKS:
-            csv_bytes, err = self.fetch_file(task_id)
-            if csv_bytes:
-                result = self.handle_csv_sales(csv_bytes)
-                if result.startswith("["):
-                    return "[SKIPPED: Confidence check failed]"
-                return result
-            return err
+        if task_id in self.task_templates:
+            return self.task_templates[task_id](question)
+        else:
+            return "[SKIPPED: Task not yet implemented in Agent V9]"
+
+    def q_mercedes_sosa_albums(self, question: str) -> str:
+        prompt = (
+            "You are a research assistant using 2022 English Wikipedia knowledge.\n"
+            "\nQUESTION:\n"
+            "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)?\n"
+            "\nFollow this step-by-step reasoning chain:\n"
+            "Step 1: Identify all studio albums by Mercedes Sosa.\n"
+            "Step 2: Filter the albums published between 2000 and 2009 inclusive.\n"
+            "Step 3: Count them.\n"
+            "Step 4: Return only the count as a number (no text, no explanation).\n"
+            "\nANSWER:"
+        )
 
-        # 3. Text questions with high confidence
         try:
             response = self.client.chat.completions.create(
                 model="gpt-4-turbo",
                 messages=[
                     {"role": "system", "content": self.instructions},
-                    {"role": "user", "content": f"QUESTION: {question}\nANSWER (concise):"}
+                    {"role": "user", "content": prompt}
                 ],
                 temperature=0.0
             )
             return response.choices[0].message.content.strip()
         except Exception as e:
-            return f"[LLM ERROR: {e}]"
+            return f"[AGENT ERROR: {e}]"