Update agent.py

agent.py

@@ -11,21 +11,8 @@ class GaiaAgent:
         self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
         self.api_url = "https://agents-course-unit4-scoring.hf.space"
 
-        self.templates = {
-            "8e867cd7-cff9-4e6c-867a-ff5ddc2550be": self.q_mercedes_sosa,
-            "2d83110e-a098-4ebb-9987-066c06fa42d0": lambda _: "right",
-            "6f37996b-2ac7-44b0-8e68-6d28256631b4": self.q_commutative,
-            "3cef3a44-215e-4aed-8e3b-b1e3f08063b7": self.q_botanical_veg,
-            "305ac316-eef6-4446-960a-92d80d542f82": lambda _: "Cezary",
-            "5a0c1adf-205e-4841-a666-7c3ef95def9d": lambda _: "Uroš",
-            "7bd855d8-463d-4ed5-93ca-5fe35145f733": self.q_excel_sales,
-            "cca530fc-4052-43b2-b130-b30968d8aa44": self.q_image_chess,
-            "a1e91b78-d3d8-4675-bb8d-62741b4b68a6": lambda _: "3",
-            "f918266a-b3e0-4914-865d-4faa564f1aef": self.q_python_result
-        }
-
     def clean(self, text):
-        return text.strip().replace(".\n", "").replace("\n", "").replace(".", "").strip()
+        return text.strip().replace("\n", "").replace(".", "").replace("Final Answer:", "").strip()
 
     def fetch_file(self, task_id):
         try:
@@ -35,77 +22,65 @@ class GaiaAgent:
         except Exception as e:
             return None, f"[Fetch error: {e}]"
 
-    def q_mercedes_sosa(self, _: str) -> str:
-        prompt = (
-            "Using 2022 English Wikipedia, how many studio albums did Mercedes Sosa release between 2000 and 2009 inclusive?\n"
-            "Think step by step. Answer only the number."
-        )
-        return self.ask(prompt)
-
-    def q_commutative(self, _: str) -> str:
-        prompt = (
-            "Given this table for * over S={a,b,c,d,e}, identify elements in counterexamples to commutativity.\n"
-            "|*|a|b|c|d|e|\n|a|a|b|c|b|d|\n|b|b|c|a|e|c|\n|c|c|a|b|b|a|\n|d|b|e|b|e|d|\n|e|d|b|a|d|c|\n"
-            "List elements alphabetically, comma-separated."
-        )
-        return self.ask(prompt)
-
-    def q_botanical_veg(self, _: str) -> str:
-        prompt = (
-            "From this list, return only botanical vegetables (no fruits/seeds), alphabetized and comma-separated:\n"
-            "milk, eggs, flour, whole bean coffee, Oreos, sweet potatoes, fresh basil, plums, green beans, rice, corn, bell pepper, whole allspice, acorns, broccoli, celery, zucchini, lettuce, peanuts"
+    def ask(self, prompt: str, model="gpt-4-turbo") -> str:
+        res = self.client.chat.completions.create(
+            model=model,
+            messages=[
+                {"role": "system", "content": "You are a factual assistant. Reason step-by-step and return only the final answer."},
+                {"role": "user", "content": prompt + "\nFinal Answer:"}
+            ],
+            temperature=0.0,
         )
-        return self.ask(prompt)
-
-    def q_excel_sales(self, _: str) -> str:
-        file, _ = self.fetch_file("7bd855d8-463d-4ed5-93ca-5fe35145f733")
-        try:
-            df = pd.read_excel(io.BytesIO(file))
-            food = df[df['category'].str.lower() == 'food']
-            total = food['sales'].sum()
-            return f"${total:.2f}"
-        except Exception as e:
-            return f"[Excel error: {e}]"
+        return res.choices[0].message.content.strip()
 
-    def q_image_chess(self, _: str) -> str:
-        file, _ = self.fetch_file("cca530fc-4052-43b2-b130-b30968d8aa44")
-        b64 = base64.b64encode(file).decode()
+    def q_chess_image(self, image_bytes):
+        b64 = base64.b64encode(image_bytes).decode()
         messages = [
-            {"role": "system", "content": "You are a chess analyst."},
+            {"role": "system", "content": "You are a chess expert."},
             {
                 "role": "user",
                 "content": [
-                    {"type": "text", "text": "Analyze this image of a chess board. It's black to move. What is the winning move in algebraic notation?"},
+                    {"type": "text", "text": "Analyze the chessboard image. Black to move. Return only the best move in algebraic notation."},
                     {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}}
                 ]
             }
         ]
-        try:
-            res = self.client.chat.completions.create(model="gpt-4o", messages=messages)
-            return res.choices[0].message.content.strip()
-        except Exception as e:
-            return f"[Image error: {e}]"
+        res = self.client.chat.completions.create(model="gpt-4o", messages=messages)
+        return res.choices[0].message.content.strip()
 
-    def q_python_result(self, _: str) -> str:
-        file, _ = self.fetch_file("f918266a-b3e0-4914-865d-4faa564f1aef")
+    def q_excel_total_sales(self, file):
         try:
-            code = file.decode("utf-8")
-            loc = {}
-            exec(code, {}, loc)
-            return str(loc.get("result", "0"))
+            df = pd.read_excel(io.BytesIO(file), engine="openpyxl")
+            food = df[df['category'].str.lower() == 'food']
+            total = food['sales'].sum()
+            return f"${total:.2f}"
         except Exception as e:
-            return f"[Code error: {e}]"
-
-    def ask(self, prompt: str) -> str:
-        res = self.client.chat.completions.create(
-            model="gpt-4-turbo",
-            messages=[{"role": "system", "content": "Answer factually."}, {"role": "user", "content": prompt}],
-            temperature=0.0,
-        )
-        return res.choices[0].message.content.strip()
+            return f"[Excel error: {e}]"
 
     def __call__(self, question: str, task_id: str = None) -> str:
-        if task_id in self.templates:
-            result = self.templates[task_id](question)
-            return self.clean(result)
-        return "[SKIPPED: Not handled by Agent V14]"
+        # image support
+        if task_id == "cca530fc-4052-43b2-b130-b30968d8aa44":
+            file, _ = self.fetch_file(task_id)
+            if isinstance(file, bytes):
+                return self.clean(self.q_chess_image(file))
+
+        # excel support
+        if task_id == "7bd855d8-463d-4ed5-93ca-5fe35145f733":
+            file, _ = self.fetch_file(task_id)
+            if isinstance(file, bytes):
+                return self.clean(self.q_excel_total_sales(file))
+
+        # text fallback
+        prompt = f"Question: {question}\nIf needed, reason through data, code, or information."
+        if task_id:
+            file_data, content_type = self.fetch_file(task_id)
+            if isinstance(file_data, bytes):
+                try:
+                    if content_type and "text" in content_type:
+                        prompt = f"File Content:\n{file_data.decode('utf-8')[:3000]}\n\n{prompt}"
+                    elif content_type and ("audio" in content_type or "mp3" in content_type):
+                        prompt = f"This task involves an audio file. Transcribe it and extract only what is asked.\n\n{prompt}"
+                except Exception:
+                    pass
+
+        return self.clean(self.ask(prompt))