Update agent.py

agent.py

@@ -5,97 +5,129 @@ 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"
+}
+CSV_TASKS = {
+    "7bd855d8-463d-4ed5-93ca-5fe35145f733"
+}
+
 class GaiaAgent:
     def __init__(self):
         self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
-        self.instructions = (
-            "You are a multimodal GAIA assistant capable of understanding text, images, audio, and code. "
-            "Use file context if provided, think step by step, and respond with the exact answer only."
-        )
         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."
 
-    def fetch_file(self, task_id: str) -> (str, bytes, str):
+    def fetch_file(self, task_id):
         try:
             url = f"{self.api_url}/files/{task_id}"
-            response = requests.get(url, timeout=15)
-            response.raise_for_status()
-            content_type = response.headers.get("Content-Type", "")
-            return url, response.content, content_type
+            r = requests.get(url, timeout=15)
+            r.raise_for_status()
+            return r.content, r.headers.get("Content-Type", "")
         except Exception as e:
-            return None, None, f"[Fetch error: {e}]"
-
-    def __call__(self, question: str, task_id: str = None) -> str:
-        image = None
-        audio = None
-        tool_context = ""
+            return None, f"[FILE ERROR: {e}]"
 
-        if task_id:
-            url, file_bytes, file_type = self.fetch_file(task_id)
-            if file_bytes is None:
-                tool_context = file_type  # error message
-            elif "image" in file_type:
-                image = base64.b64encode(file_bytes).decode("utf-8")
-            elif "audio" in file_type:
-                audio = file_bytes
-            elif file_type.endswith("python"):
-                try:
-                    exec_env = {}
-                    exec(file_bytes.decode("utf-8"), {}, exec_env)
-                    result = exec_env.get("result", "[Executed. Check code return value manually if needed.]")
-                    tool_context = f"Python result: {result}"
-                except Exception as e:
-                    tool_context = f"[Python execution error: {e}]"
-            elif "text" in file_type or "csv" in file_type:
-                tool_context = file_bytes.decode("utf-8")[:2000]
-            elif "pdf" in file_type:
-                tool_context = "[PDF file detected. OCR not yet implemented.]"
+    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": f"{tool_context}\n\nQUESTION: {question}\nANSWER:"}
+            {
+                "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):
         try:
-            if image:
-                response = self.client.chat.completions.create(
-                    model="gpt-4o",
-                    messages=[
-                        {"role": "system", "content": self.instructions},
-                        {
-                            "role": "user",
-                            "content": [
-                                {"type": "text", "text": question},
-                                {
-                                    "type": "image_url",
-                                    "image_url": {
-                                        "url": f"data:image/png;base64,{image}",
-                                        "detail": "auto"
-                                    }
-                                }
-                            ]
-                        }
-                    ]
-                )
-            elif audio:
-                transcript = self.client.audio.transcriptions.create(
-                    model="whisper-1",
-                    file=io.BytesIO(audio),
-                    response_format="text"
-                )
-                messages.append({"role": "user", "content": f"Transcript: {transcript.strip()}"})
-                response = self.client.chat.completions.create(
-                    model="gpt-4-turbo",
-                    messages=messages,
-                    temperature=0.0
-                )
-            else:
-                response = self.client.chat.completions.create(
-                    model="gpt-4-turbo",
-                    messages=messages,
-                    temperature=0.0
-                )
+            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()
+            return f"${total:.2f}"
+        except Exception as e:
+            return f"[CSV ERROR: {e}]"
 
-            return response.choices[0].message.content.strip()
+    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
+
+        # 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
+        if task_id in CSV_TASKS:
+            file, err = self.fetch_file(task_id)
+            if file:
+                return self.handle_csv(file, question)
+            return err
+
+        # fallback to LLM only
+        return self.ask_llm(question)
+
+    def ask_llm(self, prompt: str) -> str:
+        try:
+            response = self.client.chat.completions.create(
+                model="gpt-4-turbo",
+                messages=[
+                    {"role": "system", "content": self.instructions},
+                    {"role": "user", "content": prompt.strip()}
+                ],
+                temperature=0.0,
+            )
+            return response.choices[0].message.content.strip()
         except Exception as e:
-            return f"[Agent error: {e}]"
+            return f"[LLM ERROR: {e}]"