Update agent.py

agent.py

@@ -2,9 +2,8 @@ import os
 import io
 import re
 import requests
-import mimetypes
-import subprocess
 import tempfile
+import subprocess
 from openai import OpenAI
 from duckduckgo_search import DDGS
 from PIL import Image
@@ -34,15 +33,15 @@ def run_web_search(query, max_results=3):
     try:
         ddgs = DDGS()
         results = ddgs.text(query)
+        bodies = []
         for i, r in enumerate(results):
             if i >= max_results:
                 break
-            # Prefer summary/body if available
             if r.get('body'):
-                return r['body']
+                bodies.append(r['body'])
             elif r.get('title'):
-                return r['title']
-        return ""
+                bodies.append(r['title'])
+        return "\n".join(bodies)
     except Exception:
         return ""
 
@@ -134,63 +133,38 @@ def guess_youtube_link(question):
     return None
 
 def format_gaia_answer(answer, question=None):
-    """Enforces strict GAIA benchmark answer formatting rules."""
     if not answer or not isinstance(answer, str):
         return ""
-
-    # Remove apologies and boilerplate
     answer = re.sub(r"(?i)i'?m sorry[,\.]?|i cannot|i can't|unable to|please provide.*|information not available|I can't assist.*|I'm unable.*|process the file directly", "", answer)
-    answer = answer.strip()
-
-    # Remove "Final Answer:" and similar prefixes
     answer = re.sub(r'(?i)final answer:?\s*', '', answer).strip()
-
-    # Remove enclosing quotes/brackets
     if answer.startswith('"') and answer.endswith('"'):
         answer = answer[1:-1]
     if answer.startswith('[') and answer.endswith(']'):
         answer = answer[1:-1]
-
-    # Remove period at end unless part of the answer (like "Indeed.")
     if not re.match(r'^[A-Za-z]+\.$', answer):
         answer = re.sub(r'\.$', '', answer)
 
-    # For specific answer types:
     if question:
-        # Numeric answer only
+        # Pure number answers
         if re.search(r'how many|number of|at bats|total sales|albums|output.*python|highest number', question, re.I):
             num_match = re.search(r'(\$?\d[\d,\.]*)', answer)
             if num_match:
                 return num_match.group(1).replace(',', '')
 
-        # Only first name (e.g. Malko, Magda M)
         if re.search(r'first name', question, re.I):
-            first = answer.strip().split()[0]
-            return first
-
-        # Only surname
+            return answer.strip().split()[0]
         if re.search(r'surname', question, re.I):
-            surname = answer.strip().split()[-1]
-            return surname
-
-        # Only city
+            return answer.strip().split()[-1]
         if re.search(r'city', question, re.I):
-            city = answer.strip().split()[0]
-            return city
-
-        # Only code (Olympics, NASA award)
+            return answer.strip().split()[0]
         if re.search(r'IOC country code|award number|NASA', question, re.I):
             code_match = re.search(r'[A-Z0-9]{3,}', answer)
             if code_match:
                 return code_match.group(0)
-
-        # Only algebraic move (chess)
         if 'algebraic notation' in question or 'chess' in question:
             move_match = re.search(r'[A-Za-z0-9]+[#\+]?$', answer)
             if move_match:
                 return move_match.group(0)
-
-        # Direct quote (Teal'c)
         if "what does teal'c say" in question.lower():
             qmatch = re.search(r'"(Indeed\.)"', answer)
             if qmatch:
@@ -198,27 +172,19 @@ def format_gaia_answer(answer, question=None):
             if "Indeed." in answer:
                 return "Indeed."
             return answer
-
-        # For lists (ingredients, vegetables, page numbers, etc)
         if re.search(r'list|comma.*separated|page numbers', question, re.I):
-            # Extract all possible meaningful phrases
             items = [x.strip('",.').lower() for x in re.split(r'[,\n]', answer) if x.strip()]
-            # Remove likely non-items (like "and", "or", etc.)
             items = [item for item in items if item and not re.match(r'(and|or|to|with|for|a|the)$', item)]
-            # For page numbers, sort as int
             if 'page numbers' in question:
                 nums = [int(x) for x in re.findall(r'\d+', answer)]
                 return ', '.join(str(n) for n in sorted(nums))
-            # For vegetables, ingredients, etc. sort alpha
             if 'ingredient' in question or 'vegetable' in question or 'grocery' in question:
-                # merge multi-word items split by commas (heuristic)
                 merged = []
                 skip = False
                 for i, item in enumerate(items):
                     if skip:
                         skip = False
                         continue
-                    # Try to merge known phrases (e.g., "sweet potatoes", "green beans", etc.)
                     if i+1 < len(items) and item in ['sweet', 'green', 'lemon', 'ripe', 'whole', 'fresh']:
                         merged.append(f"{item} {items[i+1]}")
                         skip = True
@@ -227,13 +193,10 @@ def format_gaia_answer(answer, question=None):
                 merged = sorted(set(merged))
                 return ', '.join(merged)
             return ', '.join(items)
-
-        # Only last names for pitchers (before/after)
         if re.search(r'pitcher.*before.*after', question, re.I):
             names = re.findall(r'\b[A-Z][a-z]+', answer)
             return ', '.join(names[:2])
 
-    # Generic fallback
     return answer.strip().rstrip('.').strip()
 
 class GaiaAgent:
@@ -249,10 +212,10 @@ class GaiaAgent:
             "Always output the answer only—no explanations, no extra text."
         )
 
-    def answer_with_tools(self, question, task_id):
+    def __call__(self, question: str, task_id: str = None) -> str:
         file_text = ""
         prompt_parts = [self.instructions]
-        # 1. File handling (image, Excel, CSV, PDF, text, audio)
+        # 1. File (image, Excel, etc)
         if task_id:
             file_bytes, content_type = fetch_file(task_id)
             if file_bytes and content_type:
@@ -265,7 +228,7 @@ class GaiaAgent:
             transcript = transcribe_youtube_audio(youtube_url)
             if transcript:
                 prompt_parts.append(f"Here is the transcript of the video:\n{transcript}\n")
-        # 3. Web search fallback if not enough info
+        # 3. Web search for open facts
         search_needed = not file_text and not youtube_url
         search_keywords = [
             "who", "what", "when", "where", "name", "number", "how many",
@@ -275,13 +238,8 @@ class GaiaAgent:
             search_results = run_web_search(question)
             if search_results:
                 prompt_parts.append(f"Here are relevant web search results:\n{search_results}\n")
-        # 4. Compose prompt
         prompt_parts.append(f"Question: {question}\nAnswer strictly and concisely.")
         prompt = "\n".join(prompt_parts)
-        return prompt
-
-    def __call__(self, question: str, task_id: str = None) -> str:
-        prompt = self.answer_with_tools(question, task_id)
         response = self.client.chat.completions.create(
             model="gpt-4o",
             messages=[
@@ -293,7 +251,6 @@ class GaiaAgent:
         )
         raw_output = safe_strip(response.choices[0].message.content)
         formatted = format_gaia_answer(raw_output, question)
-        # Retry with web search if result is empty or likely incorrect for key factual types
         if not formatted or formatted.lower() in ('', 'unknown', 'none', 'n/a') or 'apolog' in formatted.lower():
             web_info = run_web_search(question)
             if web_info: