Update agent.py

agent.py

@@ -2,6 +2,8 @@ import os
 import requests
 import re
 import base64
+import pandas as pd
+import io
 from openai import OpenAI
 
 class GaiaAgent:
@@ -39,11 +41,22 @@ class GaiaAgent:
         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": "text", "text": "Please analyze the image and answer the question accurately."},
                 {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_b64}"}}
             ]
         }
 
+    def handle_excel_sales_question(self, excel_bytes: bytes, question: str) -> str:
+        try:
+            df = pd.read_excel(io.BytesIO(excel_bytes))
+            if 'category' in df.columns and 'sales' in df.columns:
+                food_only = df[df['category'].str.lower() == 'food']
+                total = food_only['sales'].sum()
+                return f"${total:.2f}"
+            return "[SKIPPED: Required columns not found in Excel]"
+        except Exception as e:
+            return f"[Excel processing error: {e}]"
+
     def __call__(self, question: str, task_id: str = None) -> str:
         messages = [{"role": "system", "content": self.instructions}]
 
@@ -73,6 +86,9 @@ class GaiaAgent:
             elif isinstance(content_type, str) and "audio" in content_type:
                 messages.append({"role": "user", "content": f"[Audio content detected]\n\nQuestion: {question}"})
 
+            elif isinstance(content_type, str) and "spreadsheet" in content_type or content_type.endswith("excel") or content_type.endswith("xlsx"):
+                return self.handle_excel_sales_question(file_data, question)
+
         video_context = self.extract_youtube_context(question)
         if video_context:
             messages.append({"role": "user", "content": f"{video_context}\n\nQuestion: {question}"})
@@ -88,4 +104,3 @@ class GaiaAgent:
             return response.choices[0].message.content.strip()
         except Exception as e:
             return f"[Answer error: {e}]"
-