Update agent.py

agent.py

@@ -1,4 +1,4 @@
-# agent_v30.py
+# agent_v31.py
 import os
 import re
 import io
@@ -24,6 +24,12 @@ class GaiaAgent:
         except Exception:
             return None, None
 
+    def get_web_info(self, query):
+        try:
+            return self.search_tool.run(query)
+        except Exception:
+            return "[NO WEB INFO FOUND]"
+
     def ask(self, prompt, model="gpt-4-turbo"):
         response = self.client.chat.completions.create(
             model=model,
@@ -35,12 +41,6 @@ class GaiaAgent:
         )
         return response.choices[0].message.content.strip()
 
-    def get_web_info(self, query):
-        try:
-            return self.search_tool.run(query)
-        except Exception:
-            return "[NO WEB INFO FOUND]"
-
     def ask_audio(self, audio_bytes, question):
         path = "/tmp/audio.mp3"
         with open(path, "wb") as f:
@@ -69,17 +69,36 @@ class GaiaAgent:
             df.columns = [col.lower() for col in df.columns]
             if 'category' in df.columns and 'sales' in df.columns:
                 df['sales'] = pd.to_numeric(df['sales'], errors='coerce')
-                food_df = df[df['category'].str.lower().str.contains('food')]
+                food_df = df[df['category'].str.lower() == 'food']
                 total = food_df['sales'].sum()
                 return f"${total:.2f}" if not pd.isna(total) else "$0.00"
         except Exception:
             pass
         return "$0.00"
 
+    def extract_commutative_set(self, question):
+        try:
+            rows = re.findall(r"\|([a-e])\|([a-e\|]+)\|", question)
+            table = {}
+            for row in rows:
+                key, values = row
+                table[key] = values.strip('|').split('|')
+            elements = list(table.keys())
+            non_comm = set()
+            for i, x in enumerate(elements):
+                for j, y in enumerate(elements):
+                    if x != y:
+                        a = table[x][j]
+                        b = table[y][i]
+                        if a != b:
+                            non_comm.update([x, y])
+            return ", ".join(sorted(non_comm))
+        except:
+            return ""
+
     def extract_answer(self, raw, question):
         q = question.lower()
         raw = raw.strip().strip("\"'").strip()
-        raw = re.sub(r"^[-•\s]*", "", raw)
 
         if "studio albums" in q:
             try:
@@ -92,11 +111,6 @@ class GaiaAgent:
             match = re.search(r"\b([KQBNR]?[a-h]?[1-8]?x?[a-h][1-8][+#]?)\b", raw)
             return match.group(1) if match else raw
 
-        if "vegetables" in q or "ingredients" in q:
-            unwanted = {"pure", "extract", "granulated", "sugar", "juice", "vanilla", "ripe", "fresh", "whole", "bean", "pinch", "cups", "salt", "water"}
-            terms = [t.lower() for t in re.findall(r"[a-zA-Z]+(?: [a-zA-Z]+)?", raw)]
-            return ", ".join(sorted(set(t for t in terms if t.split()[0] not in unwanted)))
-
         if "usd with two decimal places" in q:
             match = re.search(r"\$?([0-9]+(?:\.[0-9]{1,2})?)", raw)
             return f"${float(match.group(1)):.2f}" if match else "$0.00"
@@ -120,6 +134,12 @@ class GaiaAgent:
             match = re.search(r"80NSSC[0-9A-Z]{6,7}", raw)
             return match.group(0) if match else raw
 
+        if "vegetables" in q or "ingredients" in q:
+            stopwords = set(["pure", "extract", "granulated", "sugar", "juice", "vanilla", "ripe", "fresh", "whole", "bean", "pinch", "cups", "salt", "water"])
+            tokens = [t.lower() for t in re.findall(r"[a-zA-Z]+", raw)]
+            clean = [t for t in tokens if t not in stopwords and len(t) > 2]
+            return ", ".join(sorted(set(clean)))
+
         return raw
 
     def __call__(self, question, task_id=None):
@@ -131,15 +151,16 @@ class GaiaAgent:
             if "youtube.com" in question:
                 video_id = re.search(r"v=([\w-]+)", question)
                 if video_id:
-                    summary = self.get_web_info(f"transcript or analysis of YouTube video {video_id.group(1)}")
-                    return self.ask(f"Video summary: {summary}\n\n{question}")
+                    summary = self.get_web_info(f"youtube video transcript {video_id.group(1)}")
+                    return self.ask(f"Transcript: {summary}\n\n{question}")
 
             if "malko competition" in question.lower():
-                search = self.get_web_info("malko competition winners after 1977 yugoslavia site:wikipedia.org")
-                return self.ask(f"Web result: {search}\n\n{question}")
+                search = self.get_web_info("malko competition winner yugoslavia after 1977 site:wikipedia.org")
+                return self.ask(f"Using the search result:\n{search}\n\n{question}")
 
             if "commutative" in question:
-                return self.ask(f"Based on this table, which elements show the operation is not commutative?\n{question}\nList them comma-separated, alphabetically.")
+                result = self.extract_commutative_set(question)
+                return result
 
             if file_bytes and "image" in ctype:
                 raw = self.ask_image(file_bytes, question)