Update agent.py

agent.py

@@ -26,13 +26,28 @@ class GaiaAgent:
         res = self.client.chat.completions.create(
             model=model,
             messages=[
-                {"role": "system", "content": "You are a precise assistant. Think step by step and return only the exact answer."},
-                {"role": "user", "content": prompt + "\n\nReturn only the final answer. Do not explain. Format it exactly as expected."}
+                {"role": "system", "content": "You are a precise assistant. Think step by step and return only the final answer in the correct format."},
+                {"role": "user", "content": prompt + "\n\nFinal Answer:"}
             ],
             temperature=0.0,
         )
         return self.clean(res.choices[0].message.content)
 
+    def ask_image(self, image_bytes: bytes, question: str) -> str:
+        b64 = base64.b64encode(image_bytes).decode()
+        messages = [
+            {"role": "system", "content": "You are a visual assistant. Only return the final answer to the question."},
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": question},
+                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}}
+                ]
+            }
+        ]
+        res = self.client.chat.completions.create(model="gpt-4o", messages=messages)
+        return self.clean(res.choices[0].message.content)
+
     def q_excel_sales(self, file: bytes, question: str) -> str:
         try:
             df = pd.read_excel(io.BytesIO(file), engine="openpyxl")
@@ -54,27 +69,34 @@ class GaiaAgent:
         prompt = f"Transcript: {content}\n\nQuestion: {question}"
         return self.ask(prompt)
 
+    def extract_youtube_hint(self, question: str) -> str:
+        match = re.search(r"https://www\.youtube\.com/watch\?v=([\w-]+)", question)
+        if match:
+            return f"This task is about a YouTube video (ID: {match.group(1)}). Assume the video visually or audibly answers the question."
+        return ""
+
     def __call__(self, question: str, task_id: str = None) -> str:
-        # File-based branching
-        if task_id:
-            file, content_type = self.fetch_file(task_id)
+        context = ""
 
-            if task_id == "7bd855d8-463d-4ed5-93ca-5fe35145f733" and isinstance(file, bytes):
-                return self.q_excel_sales(file, question)
+        if "youtube.com/watch" in question:
+            context += self.extract_youtube_hint(question) + "\n"
 
-            if task_id in [
-                "99c9cc74-fdc8-46c6-8f8d-3ce2d3bfeea3",
-                "1f975693-876d-457b-a649-393859e79bf3"
-            ] and isinstance(file, bytes):
-                return self.q_audio_transcribe(file, question)
+        if task_id:
+            file, content_type = self.fetch_file(task_id)
 
-            if isinstance(file, bytes) and content_type and "text" in content_type:
-                try:
-                    text = file.decode("utf-8", errors="ignore")[:3000]
-                    prompt = f"Document:\n{text}\n\nQuestion: {question}"
-                    return self.ask(prompt)
-                except:
-                    pass
+            if isinstance(file, bytes) and content_type:
+                if "image" in content_type:
+                    return self.ask_image(file, question)
+                if "audio" in content_type or task_id.endswith(".mp3"):
+                    return self.q_audio_transcribe(file, question)
+                if "spreadsheet" in content_type or content_type.endswith("excel") or content_type.endswith("xlsx"):
+                    return self.q_excel_sales(file, question)
+                if "text" in content_type:
+                    try:
+                        text = file.decode("utf-8", errors="ignore")[:3000]
+                        context += f"File Content:\n{text}\n"
+                    except Exception:
+                        pass
 
-        # Fallback
-        return self.ask(f"Question: {question}")
+        prompt = f"{context}\nQuestion: {question}"
+        return self.ask(prompt)