Update agent.py

agent.py

@@ -5,128 +5,75 @@ import requests
 import pandas as pd
 from openai import OpenAI
 
-# --- Task classification ---
-AUDIO_TASKS = {
-    "9d191bce-651d-4746-be2d-7ef8ecadb9c2",
-    "99c9cc74-fdc8-46c6-8f8d-3ce2d3bfeea3",
-    "1f975693-876d-457b-a649-393859e79bf3"
-}
-IMAGE_TASKS = {
-    "a1e91b78-d3d8-4675-bb8d-62741b4b68a6",
-    "cca530fc-4052-43b2-b130-b30968d8aa44"
-}
-CODE_TASKS = {
-    "f918266a-b3e0-4914-865d-4faa564f1aef"
+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"
+    "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 helpful assistant solving GAIA benchmark questions using any available tools."
+        self.instructions = (
+            "You are a precise assistant solving GAIA benchmark questions. "
+            "Only answer if you are confident you can provide the exact correct result."
+        )
 
     def fetch_file(self, task_id):
         try:
             url = f"{self.api_url}/files/{task_id}"
-            r = requests.get(url, timeout=15)
+            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_audio(self, audio_bytes):
-        try:
-            transcript = self.client.audio.transcriptions.create(
-                model="whisper-1",
-                file=io.BytesIO(audio_bytes),
-                response_format="text"
-            )
-            return transcript.strip()
-        except Exception as e:
-            return f"[TRANSCRIPTION ERROR: {e}]"
-
-    def handle_image(self, image_bytes, question):
-        b64 = base64.b64encode(image_bytes).decode("utf-8")
-        messages = [
-            {"role": "system", "content": self.instructions},
-            {
-                "role": "user",
-                "content": [
-                    {"type": "text", "text": question},
-                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}}
-                ]
-            }
-        ]
-        try:
-            response = self.client.chat.completions.create(model="gpt-4o", messages=messages)
-            return response.choices[0].message.content.strip()
-        except Exception as e:
-            return f"[IMAGE ERROR: {e}]"
-
-    def handle_csv(self, csv_bytes, question):
+    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()))
-            total = df[df['category'].str.lower() == 'food']['sales'].sum()
+            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}]"
 
-    def handle_code(self, code_bytes):
-        try:
-            exec_env = {}
-            exec(code_bytes.decode("utf-8"), {}, exec_env)
-            return str(exec_env.get("result", "[Executed. Check result variable manually]"))
-        except Exception as e:
-            return f"[EXEC ERROR: {e}]"
-
     def __call__(self, question: str, task_id: str = None) -> str:
-        if not task_id:
-            return self.ask_llm(question)
-
-        # audio
-        if task_id in AUDIO_TASKS:
-            file, err = self.fetch_file(task_id)
-            if file:
-                transcript = self.handle_audio(file)
-                return self.ask_llm(f"Audio transcript: {transcript}\n\nQuestion: {question}")
-            return err
+        # 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]"
 
-        # image
-        if task_id in IMAGE_TASKS:
-            file, err = self.fetch_file(task_id)
-            if file:
-                return self.handle_image(file, question)
-            return err
-
-        # python code
-        if task_id in CODE_TASKS:
-            file, err = self.fetch_file(task_id)
-            if file:
-                return self.handle_code(file)
-            return err
-
-        # CSV/Excel
+        # 2. CSV handling
         if task_id in CSV_TASKS:
-            file, err = self.fetch_file(task_id)
-            if file:
-                return self.handle_csv(file, question)
+            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
 
-        # fallback to LLM only
-        return self.ask_llm(question)
-
-    def ask_llm(self, prompt: str) -> str:
+        # 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": prompt.strip()}
+                    {"role": "user", "content": f"QUESTION: {question}\nANSWER (concise):"}
                 ],
-                temperature=0.0,
+                temperature=0.0
             )
             return response.choices[0].message.content.strip()
         except Exception as e: