Update agent.py

agent.py

@@ -1,10 +1,10 @@
 import os
 import io
+import re
 import requests
 import mimetypes
 import subprocess
 import tempfile
-import re
 from openai import OpenAI
 from duckduckgo_search import DDGS
 from PIL import Image
@@ -30,108 +30,6 @@ def safe_strip(text):
         text = text.decode(errors="ignore")
     return str(text).replace("\r", "").strip()
 
-def format_gaia_answer(answer, question=None):
-    """
-    Enforces strict GAIA benchmark answer formatting rules.
-    - Strips explanations, apologies, quotes, brackets, units, periods.
-    - For lists: comma-separated, no quotes, no brackets, alphabetized if asked.
-    - For numbers: digits only (unless $ required).
-    - For names: no title, no extra text.
-    - For code: just the output.
-    - Optionally takes question for context-sensitive formatting.
-    """
-    if not answer or not isinstance(answer, str):
-        return ""
-
-    # Remove apologies/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.*", "", answer)
-    answer = answer.strip()
-
-    # Remove "Final Answer:" and similar prefixes
-    answer = re.sub(r'(?i)final answer:?\s*', '', answer).strip()
-
-    # Remove enclosing quotes/brackets
-    answer = answer.strip()
-    if answer.startswith('"') and answer.endswith('"'):
-        answer = answer[1:-1]
-    if answer.startswith('[') and answer.endswith(']'):
-        answer = answer[1:-1]
-
-    # Remove periods at end, unless required (like Teal'c "Indeed.")
-    # Exception: If the answer is just 'Indeed.' or similar, keep it.
-    if not re.match(r'^[A-Za-z]+\.$', answer):
-        answer = re.sub(r'\.$', '', answer)
-
-    # Remove extra text before/after answer for known Q types
-    # Numbers only
-    if question:
-        if re.search(r'how many|number of|at bats|total sales|albums|output.*python', question, re.I):
-            num_match = re.search(r'(\$?\d[\d,\.]*)', answer)
-            if num_match:
-                return num_match.group(1).replace(',', '')
-
-        # Only the first name (Malko, Magda M)
-        if re.search(r'first name', question, re.I):
-            first = answer.strip().split()[0]
-            return first
-
-        # Only the surname (LibreText vet)
-        if re.search(r'surname', question, re.I):
-            surname = answer.strip().split()[-1]
-            return surname
-
-        # Only the city (Vietnamese specimens)
-        if re.search(r'city', question, re.I):
-            city = answer.strip().split()[0]
-            return city
-
-        # Only the code (Olympics, NASA award)
-        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():
-            # Try to extract quoted phrase or just Indeed.
-            qmatch = re.search(r'"(Indeed\.)"', answer)
-            if qmatch:
-                return qmatch.group(1)
-            # Fallback: find Indeed.
-            if "Indeed." in answer:
-                return "Indeed."
-            return answer
-
-        # For lists: comma separated, strip spaces, no quotes/brackets, alpha order if needed
-        if re.search(r'list|comma.*separated|page numbers', question, re.I):
-            # extract all words/numbers, remove measurements
-            items = re.findall(r'\b[A-Za-z0-9\-\']+\b', answer)
-            # Special: 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))
-            # Special: ingredients/veggies/fruits, sort alpha
-            if 'ingredients' in question or 'vegetables' in question or 'grocery' in question:
-                # Lowercase, no duplicates, alpha order
-                items = [x.lower() for x in items]
-                items = sorted(set(items))
-                return ', '.join(items)
-            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: remove any trailing period, strip whitespace
-    return answer.strip().rstrip('.').strip()
-
 def run_web_search(query, max_results=3):
     try:
         ddgs = DDGS()
@@ -139,6 +37,7 @@ def run_web_search(query, max_results=3):
         for i, r in enumerate(results):
             if i >= max_results:
                 break
+            # Prefer summary/body if available
             if r.get('body'):
                 return r['body']
             elif r.get('title'):
@@ -197,9 +96,6 @@ def transcribe_audio(audio_bytes):
         return ""
 
 def transcribe_youtube_audio(youtube_url):
-    """
-    Download audio from YouTube, transcribe using whisper
-    """
     if not whisper:
         return ""
     try:
@@ -218,32 +114,128 @@ def transcribe_youtube_audio(youtube_url):
         return ""
 
 def extract_file_text(file_bytes, content_type, task_id=""):
-    # Images
     if "image" in content_type:
         return ocr_image(file_bytes)
-    # Excel
     if "spreadsheet" in content_type or "excel" in content_type or task_id.endswith(".xlsx"):
         return read_excel(file_bytes)
-    # PDF
     if "pdf" in content_type or task_id.endswith(".pdf"):
         return read_pdf(file_bytes)
-    # Audio
     if "audio" in content_type or task_id.endswith(".mp3") or task_id.endswith(".wav"):
         return transcribe_audio(file_bytes)
-    # Text, CSV, JSON
     if "text" in content_type or "csv" in content_type or "json" in content_type or task_id.endswith(".csv") or task_id.endswith(".json") or task_id.endswith(".txt"):
         return safe_strip(file_bytes[:10000])
     return ""
 
 def guess_youtube_link(question):
-    # If the question mentions YouTube or a video link, try to extract it
-    import re
     matches = re.findall(r"(https?://[^\s]+)", question)
     for url in matches:
         if "youtube.com" in url or "youtu.be" in url:
             return url
     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
+        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
+        if re.search(r'surname', question, re.I):
+            surname = answer.strip().split()[-1]
+            return surname
+
+        # Only city
+        if re.search(r'city', question, re.I):
+            city = answer.strip().split()[0]
+            return city
+
+        # Only code (Olympics, NASA award)
+        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:
+                return qmatch.group(1)
+            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
+                    else:
+                        merged.append(item)
+                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:
     def __init__(self):
         self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
@@ -257,10 +249,8 @@ class GaiaAgent:
             "Always output the answer only—no explanations, no extra text."
         )
 
-    def __call__(self, question: str, task_id: str = None) -> str:
+    def answer_with_tools(self, question, task_id):
         file_text = ""
-        web_context = ""
-        video_transcript = ""
         prompt_parts = [self.instructions]
         # 1. File handling (image, Excel, CSV, PDF, text, audio)
         if task_id:
@@ -269,25 +259,29 @@ class GaiaAgent:
                 file_text = extract_file_text(file_bytes, content_type, task_id)
                 if file_text:
                     prompt_parts.append(f"Here is the extracted file content:\n{file_text}\n")
-        # 2. YouTube/video handling (by URL in question)
+        # 2. YouTube/video
         youtube_url = guess_youtube_link(question)
         if youtube_url:
             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 for open-world/factoid questions or if no file info
+        # 3. Web search fallback if not enough info
+        search_needed = not file_text and not youtube_url
         search_keywords = [
             "who", "what", "when", "where", "name", "number", "how many",
             "first", "last", "award", "recipient", "code", "surname", "year", "album", "actor", "winner"
         ]
-        if (not file_text and not youtube_url) or any(kw in question.lower() for kw in search_keywords):
+        if search_needed or any(kw in question.lower() for kw in search_keywords):
             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)
-        # 5. Call LLM
+        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=[
@@ -298,8 +292,27 @@ class GaiaAgent:
             max_tokens=512,
         )
         raw_output = safe_strip(response.choices[0].message.content)
-        # 6. Format the answer strictly per benchmark rules
-        return format_gaia_answer(raw_output, question)
+        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:
+                prompt2 = (
+                    f"{self.instructions}\n\n"
+                    f"Here are relevant web search results:\n{web_info}\n"
+                    f"Question: {question}\nAnswer strictly and concisely."
+                )
+                response2 = self.client.chat.completions.create(
+                    model="gpt-4o",
+                    messages=[
+                        {"role": "system", "content": self.instructions},
+                        {"role": "user", "content": prompt2}
+                    ],
+                    temperature=0.0,
+                    max_tokens=256,
+                )
+                formatted = format_gaia_answer(safe_strip(response2.choices[0].message.content), question)
+        return formatted
 
 def answer_question(question, task_id=None):
     agent = GaiaAgent()