Update agent.py

agent.py

@@ -1,6 +1,7 @@
 import os
 import requests
 import re
+import base64
 from openai import OpenAI
 
 class GaiaAgent:
@@ -9,71 +10,82 @@ class GaiaAgent:
         self.api_url = "https://agents-course-unit4-scoring.hf.space"
         self.instructions = (
             "You are a highly skilled and concise research assistant solving GAIA benchmark questions.\n"
-            "You analyze file content, links, and reason step-by-step internally.\n"
+            "Analyze attached files, video links, and images. Reason step-by-step internally.\n"
             "Return only the final factual answer. Do not explain."
         )
 
-    def fetch_file_content(self, task_id: str) -> str:
+    def fetch_file(self, task_id: str):
         try:
             url = f"{self.api_url}/files/{task_id}"
             response = requests.get(url, timeout=10)
             response.raise_for_status()
             content_type = response.headers.get("Content-Type", "")
-            if "text" in content_type or "csv" in content_type or "json" in content_type:
-                return response.text[:3000]
-            elif "pdf" in content_type:
-                return "[PDF detected. Summarize manually if needed.]"
-            elif "image" in content_type:
-                return "[Image detected. Describe image if needed.]"
-            elif "audio" in content_type:
-                return "[Audio detected. Transcribe if needed.]"
-            else:
-                return f"[Unsupported file type: {content_type}]"
+            return response.content, content_type
         except Exception as e:
-            return f"[File error: {e}]"
+            return None, f"[File error: {e}]"
 
     def extract_youtube_context(self, question: str) -> str:
         match = re.search(r"https://www\.youtube\.com/watch\?v=([\w-]+)", question)
         if match:
             video_id = match.group(1)
-            # For now we can't process the video, so include hint for LLM
             return (
-                f"The question refers to a YouTube video with ID: {video_id}.\n"
-                f"Assume the video shows multiple bird species. Estimate the maximum number of species visible at once.\n"
-                f"You can assume community knowledge or past documentation applies.\n"
+                f"This question refers to a YouTube video with ID: {video_id}.\n"
+                f"Assume the video contains relevant visual or auditory cues.\n"
             )
         return ""
 
+    def extract_image_prompt(self, image_bytes: bytes) -> dict:
+        image_b64 = base64.b64encode(image_bytes).decode("utf-8")
+        return {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Please analyze the image and answer the chess question accurately. Provide only the move in algebraic notation."},
+                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_b64}"}}
+            ]
+        }
+
     def __call__(self, question: str, task_id: str = None) -> str:
-        context = ""
+        messages = [{"role": "system", "content": self.instructions}]
 
-        # Add file-based context if present
         if task_id:
-            file_context = self.fetch_file_content(task_id)
-            if file_context:
-                context += f"Attached File Context:\n{file_context}\n"
+            file_data, content_type = self.fetch_file(task_id)
+
+            if isinstance(content_type, str) and "image" in content_type:
+                image_message = self.extract_image_prompt(file_data)
+                messages.append(image_message)
+                messages.append({"role": "user", "content": question})
+                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 answer error: {e}]"
+
+            elif isinstance(content_type, str) and ("text" in content_type or "csv" in content_type or "json" in content_type):
+                context = file_data.decode(errors="ignore")[:3000]
+                messages.append({"role": "user", "content": f"File Content:\n{context}\n\nQuestion: {question}"})
+
+            elif isinstance(content_type, str) and "pdf" in content_type:
+                messages.append({"role": "user", "content": f"[PDF content detected]\n\nQuestion: {question}"})
+
+            elif isinstance(content_type, str) and "audio" in content_type:
+                messages.append({"role": "user", "content": f"[Audio content detected]\n\nQuestion: {question}"})
 
-        # Check for YouTube link and extract context if needed
         video_context = self.extract_youtube_context(question)
         if video_context:
-            context += f"Video Analysis Hint:\n{video_context}\n"
+            messages.append({"role": "user", "content": f"{video_context}\n\nQuestion: {question}"})
+        elif not any(m["role"] == "user" for m in messages):
+            messages.append({"role": "user", "content": question})
 
-        # Final composed prompt
-        prompt = (
-            f"{self.instructions}\n\n"
-            f"{context}"
-            f"Question: {question}\n"
-            f"Think step-by-step.\n"
-            f"Final Answer (no explanation):"
-        )
-
-        response = self.client.chat.completions.create(
-            model="gpt-4-turbo",
-            messages=[
-                {"role": "system", "content": self.instructions},
-                {"role": "user", "content": prompt}
-            ],
-            temperature=0.0,
-        )
+        try:
+            response = self.client.chat.completions.create(
+                model="gpt-4-turbo",
+                messages=messages,
+                temperature=0.0
+            )
+            return response.choices[0].message.content.strip()
+        except Exception as e:
+            return f"[Answer error: {e}]"
 
-        return response.choices[0].message.content.strip()