Update app.py

app.py

@@ -3,9 +3,11 @@ import re
 import json
 import io
 import time
+import base64
 import traceback
 import contextlib
 import tempfile
+from urllib.parse import urlparse, parse_qs
 
 import gradio as gr
 import requests
@@ -14,7 +16,6 @@ import pandas as pd
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 # Primary model first, then fallbacks used when rate-limited.
-# llama-3.1-8b-instant has much higher daily token quota on the free tier.
 GROQ_MODELS = [
     m.strip()
     for m in os.getenv(
@@ -23,16 +24,48 @@ GROQ_MODELS = [
     ).split(",")
     if m.strip()
 ]
-MAX_TOOL_ITERATIONS = 6
-TOOL_RESULT_MAX_CHARS = 3000
+# Vision-capable Groq model (free tier).
+GROQ_VISION_MODEL = os.getenv("GROQ_VISION_MODEL", "meta-llama/llama-4-scout-17b-16e-instruct")
+# Groq Whisper model for audio.
+GROQ_WHISPER_MODEL = os.getenv("GROQ_WHISPER_MODEL", "whisper-large-v3-turbo")
+
+MAX_TOOL_ITERATIONS = 8
+TOOL_RESULT_MAX_CHARS = 3500
 ANSWER_CACHE_PATH = os.getenv("ANSWER_CACHE_PATH", "answers_cache.json")
 
+# Track downloaded task files so vision/audio tools can re-use them by task_id.
+_TASK_FILE_CACHE: dict[str, dict] = {}
+
 
 # ---------------------------------------------------------------------------
 # Tool implementations
 # ---------------------------------------------------------------------------
 def tool_web_search(query: str, max_results: int = 5) -> str:
-    """DuckDuckGo text search. Returns a short list of titles + snippets + URLs."""
+    """Web search. Tries Tavily first (if TAVILY_API_KEY set), falls back to DuckDuckGo."""
+    tavily_key = os.getenv("TAVILY_API_KEY")
+    if tavily_key:
+        try:
+            from tavily import TavilyClient
+            client = TavilyClient(api_key=tavily_key)
+            res = client.search(
+                query=query,
+                max_results=max_results,
+                search_depth="basic",
+                include_answer=True,
+            )
+            lines = []
+            if res.get("answer"):
+                lines.append(f"Answer: {res['answer']}")
+            for r in res.get("results", [])[:max_results]:
+                lines.append(
+                    f"- {r.get('title', '')}\n  {r.get('url', '')}\n  {r.get('content', '')[:400]}"
+                )
+            if lines:
+                return "\n".join(lines)
+        except Exception as e:
+            print(f"tavily search failed, falling back to DDG: {e}")
+
+    # Fallback: DuckDuckGo
     try:
         from duckduckgo_search import DDGS
         results = []
@@ -48,7 +81,7 @@ def tool_web_search(query: str, max_results: int = 5) -> str:
         return f"web_search error: {e}"
 
 
-def tool_fetch_url(url: str, max_chars: int = 3000) -> str:
+def tool_fetch_url(url: str, max_chars: int = 3500) -> str:
     """Fetch a URL and return readable text (HTML stripped)."""
     try:
         from bs4 import BeautifulSoup
@@ -110,6 +143,116 @@ def tool_python(code: str) -> str:
         return f"python error: {e}\n{traceback.format_exc(limit=2)}"
 
 
+def _extract_youtube_id(url: str) -> str | None:
+    try:
+        u = urlparse(url)
+        if "youtu.be" in u.netloc:
+            return u.path.lstrip("/").split("/")[0] or None
+        if "youtube.com" in u.netloc:
+            if u.path == "/watch":
+                return parse_qs(u.query).get("v", [None])[0]
+            if u.path.startswith("/embed/") or u.path.startswith("/shorts/"):
+                return u.path.split("/")[2]
+    except Exception:
+        pass
+    return None
+
+
+def tool_youtube_transcript(url: str, max_chars: int = 3500) -> str:
+    """Fetch the transcript of a YouTube video by URL or ID."""
+    try:
+        from youtube_transcript_api import YouTubeTranscriptApi
+        vid = _extract_youtube_id(url) or url.strip()
+        try:
+            data = YouTubeTranscriptApi.get_transcript(vid, languages=["en", "en-US", "en-GB"])
+        except Exception:
+            # Try any available language
+            tlist = YouTubeTranscriptApi.list_transcripts(vid)
+            t = next(iter(tlist), None)
+            data = t.fetch() if t else []
+        text = " ".join(seg.get("text", "") for seg in data).strip()
+        text = re.sub(r"\s+", " ", text)
+        if len(text) > max_chars:
+            text = text[:max_chars] + " ...[truncated]"
+        return text or "(empty transcript)"
+    except Exception as e:
+        return f"youtube_transcript error: {e}"
+
+
+def tool_transcribe_audio(task_id: str) -> str:
+    """Transcribe an audio file attached to a GAIA task using Groq Whisper."""
+    try:
+        from groq import Groq
+        # Make sure the file is downloaded.
+        info = _TASK_FILE_CACHE.get(task_id)
+        if not info:
+            tool_get_task_file(task_id)  # populates cache
+            info = _TASK_FILE_CACHE.get(task_id)
+        if not info or not os.path.exists(info.get("path", "")):
+            return "transcribe_audio error: no local file for task"
+
+        client = Groq(api_key=os.getenv("GROQ_API_KEY"))
+        with open(info["path"], "rb") as f:
+            tr = client.audio.transcriptions.create(
+                file=(os.path.basename(info["path"]), f.read()),
+                model=GROQ_WHISPER_MODEL,
+                response_format="text",
+            )
+        text = tr if isinstance(tr, str) else getattr(tr, "text", str(tr))
+        text = text.strip()
+        if len(text) > 4000:
+            text = text[:4000] + " ...[truncated]"
+        return text or "(empty transcript)"
+    except Exception as e:
+        return f"transcribe_audio error: {e}"
+
+
+def tool_view_image(task_id: str, question: str = "") -> str:
+    """Describe / answer a question about an image attached to a GAIA task using Groq vision."""
+    try:
+        from groq import Groq
+        info = _TASK_FILE_CACHE.get(task_id)
+        if not info:
+            tool_get_task_file(task_id)
+            info = _TASK_FILE_CACHE.get(task_id)
+        if not info or not os.path.exists(info.get("path", "")):
+            return "view_image error: no local file for task"
+
+        suffix = os.path.splitext(info["path"])[1].lower().lstrip(".")
+        if suffix == "jpg":
+            suffix = "jpeg"
+        if suffix not in {"png", "jpeg", "gif", "webp"}:
+            return f"view_image error: unsupported image type .{suffix}"
+
+        with open(info["path"], "rb") as f:
+            b64 = base64.b64encode(f.read()).decode("ascii")
+        data_url = f"data:image/{suffix};base64,{b64}"
+
+        prompt = (
+            question.strip()
+            or "Describe this image in detail, including any text, numbers, or symbols visible."
+        )
+
+        client = Groq(api_key=os.getenv("GROQ_API_KEY"))
+        resp = client.chat.completions.create(
+            model=GROQ_VISION_MODEL,
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": prompt},
+                        {"type": "image_url", "image_url": {"url": data_url}},
+                    ],
+                }
+            ],
+            temperature=0.0,
+            max_tokens=800,
+        )
+        return (resp.choices[0].message.content or "").strip()
+    except Exception as e:
+        return f"view_image error: {e}"
+
+
 def tool_get_task_file(task_id: str, api_url: str = DEFAULT_API_URL) -> str:
     """Download the file attached to a task and return a text preview."""
     try:
@@ -126,7 +269,17 @@ def tool_get_task_file(task_id: str, api_url: str = DEFAULT_API_URL) -> str:
         tmp.write(resp.content)
         tmp.close()
 
-        info = f"File: {fname}\nContent-Type: {ctype}\nSaved to: {tmp.name}\nSize: {len(resp.content)} bytes\n"
+        _TASK_FILE_CACHE[task_id] = {
+            "path": tmp.name,
+            "name": fname,
+            "ctype": ctype,
+            "size": len(resp.content),
+        }
+
+        info = (
+            f"File: {fname}\nContent-Type: {ctype}\nSaved to: {tmp.name}\n"
+            f"Size: {len(resp.content)} bytes\n"
+        )
 
         # Try to give a readable preview
         if suffix in {".txt", ".md", ".csv", ".json", ".py", ".tsv", ".log", ".xml", ".html"}:
@@ -152,11 +305,11 @@ def tool_get_task_file(task_id: str, api_url: str = DEFAULT_API_URL) -> str:
             except Exception as e:
                 return info + f"\n(pdf parse error: {e})"
 
-        if suffix in {".mp3", ".wav", ".m4a", ".ogg"}:
-            return info + "\n(audio file; no transcription tool available)"
+        if suffix in {".mp3", ".wav", ".m4a", ".ogg", ".flac", ".webm"}:
+            return info + "\nThis is an audio file. Call transcribe_audio to read it."
 
         if suffix in {".png", ".jpg", ".jpeg", ".gif", ".webp"}:
-            return info + "\n(image file; no vision tool available)"
+            return info + "\nThis is an image. Call view_image to inspect it."
 
         return info + "\n(binary file; no preview)"
     except Exception as e:
@@ -171,7 +324,7 @@ TOOLS_SPEC = [
         "type": "function",
         "function": {
             "name": "web_search",
-            "description": "Search the web with DuckDuckGo. Returns titles, URLs, and snippets.",
+            "description": "Search the web (Tavily preferred, DuckDuckGo fallback). Returns titles, URLs, snippets.",
             "parameters": {
                 "type": "object",
                 "properties": {
@@ -191,7 +344,7 @@ TOOLS_SPEC = [
                 "type": "object",
                 "properties": {
                     "url": {"type": "string"},
-                    "max_chars": {"type": "integer", "default": 3000},
+                    "max_chars": {"type": "integer", "default": 3500},
                 },
                 "required": ["url"],
             },
@@ -216,7 +369,7 @@ TOOLS_SPEC = [
         "type": "function",
         "function": {
             "name": "python",
-            "description": "Execute a short Python snippet for math, string parsing, or data work. Use print() or assign to `result`.",
+            "description": "Execute a short Python snippet for math, dates, parsing CSV, list/string work. Use print() or assign to `result`.",
             "parameters": {
                 "type": "object",
                 "properties": {"code": {"type": "string"}},
@@ -228,7 +381,7 @@ TOOLS_SPEC = [
         "type": "function",
         "function": {
             "name": "get_task_file",
-            "description": "Download the file attached to a GAIA task by task_id and return a text preview.",
+            "description": "Download the file attached to a GAIA task by task_id and return a text preview. Always call this first if the question references an attached file.",
             "parameters": {
                 "type": "object",
                 "properties": {"task_id": {"type": "string"}},
@@ -236,35 +389,86 @@ TOOLS_SPEC = [
             },
         },
     },
+    {
+        "type": "function",
+        "function": {
+            "name": "transcribe_audio",
+            "description": "Transcribe an attached audio file (mp3/wav/m4a/ogg/flac) for the given task_id.",
+            "parameters": {
+                "type": "object",
+                "properties": {"task_id": {"type": "string"}},
+                "required": ["task_id"],
+            },
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "view_image",
+            "description": "Inspect an attached image (png/jpg/gif/webp) using a vision model. Pass a focused question for best results.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "task_id": {"type": "string"},
+                    "question": {"type": "string"},
+                },
+                "required": ["task_id"],
+            },
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "youtube_transcript",
+            "description": "Fetch the transcript text of a YouTube video given its URL or ID.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "url": {"type": "string"},
+                    "max_chars": {"type": "integer", "default": 3500},
+                },
+                "required": ["url"],
+            },
+        },
+    },
 ]
 
 TOOL_FUNCTIONS = {
     "web_search": lambda args: tool_web_search(args["query"], int(args.get("max_results", 5))),
-    "fetch_url": lambda args: tool_fetch_url(args["url"], int(args.get("max_chars", 3000))),
+    "fetch_url": lambda args: tool_fetch_url(args["url"], int(args.get("max_chars", 3500))),
     "wikipedia": lambda args: tool_wikipedia(args["query"], int(args.get("sentences", 4))),
     "python": lambda args: tool_python(args["code"]),
     "get_task_file": lambda args: tool_get_task_file(args["task_id"]),
+    "transcribe_audio": lambda args: tool_transcribe_audio(args["task_id"]),
+    "view_image": lambda args: tool_view_image(args["task_id"], args.get("question", "")),
+    "youtube_transcript": lambda args: tool_youtube_transcript(
+        args["url"], int(args.get("max_chars", 3500))
+    ),
 }
 
 
 SYSTEM_PROMPT = """You are a careful research agent answering GAIA benchmark questions.
 
-You have tools: web_search, fetch_url, wikipedia, python, get_task_file.
+Tools: web_search, fetch_url, wikipedia, python, get_task_file, transcribe_audio, view_image, youtube_transcript.
 
 Workflow:
-- If the question references an attached file, image, audio, code, or table, call get_task_file with the task_id first.
-- Use web_search then fetch_url to verify facts from primary sources.
-- Use wikipedia for well-known entities or historical facts.
-- Use python for arithmetic, date math, string/list manipulation, or parsing CSV data.
-- Cross-check before answering. Avoid guessing.
-
-Answer formatting (critical, the grader does EXACT string match):
-- Reply with ONLY the answer. No preamble, no explanation, no quotes, no trailing period unless part of the answer.
-- Do NOT include the words "FINAL ANSWER" or any label.
-- Numbers: digits only, no commas, no units, no $ sign, unless the question asks for the unit.
-- Strings: no leading articles ("the", "a") unless required, no abbreviations, write digits as digits.
+- If the question references an attached file/image/audio/code/table, call get_task_file(task_id) first.
+  - For audio (mp3/wav/m4a/ogg), then call transcribe_audio(task_id).
+  - For images (png/jpg/gif/webp), then call view_image(task_id, question="...").
+- If the question references a YouTube link, call youtube_transcript(url).
+- Use web_search then fetch_url to verify facts from primary sources. Prefer official sites and Wikipedia.
+- Use wikipedia for well-known entities, places, and historical facts.
+- Use python for arithmetic, date math, sorting, set operations, parsing strings/CSV. Do NOT eyeball math.
+- Cross-check before answering. Don't guess. If two sources disagree, prefer the most authoritative.
+
+Answer formatting (CRITICAL — grader does an exact-match-style comparison):
+- Reply with ONLY the answer. No preamble, no explanation, no quotes, no trailing period.
+- Do NOT include the words "FINAL ANSWER", "Answer:", or any label.
+- Numbers: digits only, no commas, no units, no $ sign — UNLESS the question asks for the unit.
+- Strings: no leading articles ("the", "a") unless required; spell out, no abbreviations; write digits as digits.
 - Lists: comma-separated, single space after each comma, applying the rules above to each element.
 - If the question asks for a name, give just the name. If it asks "how many", give just the number.
+- If the question asks for a list in alphabetical order, sort it.
 """
 
 
@@ -313,7 +517,6 @@ class GroqAgent:
                     is_429 = "429" in msg or "rate_limit" in msg.lower()
                     is_tpd = "per day" in msg.lower() or "tpd" in msg.lower()
                     if is_429 and is_tpd:
-                        # Daily quota gone — switch model permanently for this run.
                         print(f"[{model}] daily token limit exhausted; switching model.")
                         self.exhausted_models.add(model)
                         break
@@ -323,14 +526,12 @@ class GroqAgent:
                         print(f"[{model}] 429 rate limit; sleeping {wait}s (attempt {attempt + 1}/3)")
                         time.sleep(wait)
                         continue
-                    # Non-429 error: don't retry on the same model.
                     print(f"[{model}] API error: {e}")
                     break
         raise RuntimeError(f"All Groq models failed. Last error: {last_error}")
 
     @staticmethod
     def _parse_retry_seconds(error_msg: str) -> float:
-        # Examples in Groq error: "Please try again in 7m18.912s." or "in 12.3s"
         m = re.search(r"in\s+(?:(\d+)m)?([\d.]+)s", error_msg)
         if not m:
             return 5.0
@@ -360,7 +561,7 @@ class GroqAgent:
 
             if not tool_calls:
                 answer = (msg.content or "").strip()
-                return self._postprocess_answer(answer)
+                return self._postprocess_answer(answer, question)
 
             messages.append(
                 {
@@ -387,7 +588,7 @@ class GroqAgent:
                 except json.JSONDecodeError:
                     args = {}
                 fn = TOOL_FUNCTIONS.get(name)
-                print(f"[tool] {name}({args})")
+                print(f"[tool] {name}({str(args)[:200]})")
                 if fn is None:
                     result = f"unknown tool: {name}"
                 else:
@@ -419,19 +620,46 @@ class GroqAgent:
         )
         try:
             resp = self._chat(messages, use_tools=False, max_tokens=256)
-            return self._postprocess_answer((resp.choices[0].message.content or "").strip())
+            return self._postprocess_answer(
+                (resp.choices[0].message.content or "").strip(), question
+            )
         except Exception as e:
             return f"AGENT ERROR: {e}"
 
     @staticmethod
-    def _postprocess_answer(text: str) -> str:
+    def _postprocess_answer(text: str, question: str = "") -> str:
         if not text:
             return ""
         text = text.strip()
-        text = re.sub(r"^(final answer|answer)\s*:\s*", "", text, flags=re.IGNORECASE)
-        if len(text) >= 2 and text[0] == text[-1] and text[0] in {'"', "'", "`"}:
+
+        # Drop common labels.
+        text = re.sub(
+            r"^(final\s*answer|answer|the\s*answer\s*is)\s*[:\-]?\s*",
+            "",
+            text,
+            flags=re.IGNORECASE,
+        )
+        # If model wrapped answer in code fence or quotes.
+        text = text.strip("`")
+        if len(text) >= 2 and text[0] == text[-1] and text[0] in {'"', "'"}:
             text = text[1:-1].strip()
-        return text
+
+        # If the question is "how many" / numeric and the model returned a sentence,
+        # try to extract a single number.
+        q_lower = question.lower()
+        wants_number = bool(
+            re.search(r"\bhow many\b|\bhow much\b|\bwhat number\b|\bcount\b", q_lower)
+        )
+        if wants_number and not re.fullmatch(r"-?\d+(\.\d+)?", text):
+            m = re.search(r"-?\d+(?:\.\d+)?", text.replace(",", ""))
+            if m:
+                text = m.group(0)
+
+        # Strip a single trailing period if the text isn't a list/sentence with internal periods.
+        if text.endswith(".") and text.count(".") == 1 and " " not in text[-3:]:
+            text = text[:-1]
+
+        return text.strip()
 
 
 # ---------------------------------------------------------------------------
@@ -531,7 +759,6 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
     }
     print(f"Submitting {len(answers_payload)} answers for user '{username}'...")
 
-    # Retry submission a few times — the leaderboard's HF dataset write is flaky.
     last_error = None
     for attempt in range(3):
         try:
@@ -583,12 +810,14 @@ with gr.Blocks() as demo:
     gr.Markdown(
         """
         **Setup**
-        1. Add a Space secret named `GROQ_API_KEY` with your Groq API key.
-        2. Optional: set `GROQ_MODELS` (comma-separated, default `llama-3.3-70b-versatile,llama-3.1-8b-instant`).
-        3. Log in to Hugging Face below and click **Run Evaluation & Submit All Answers**.
-
-        Tools available to the agent: `web_search`, `fetch_url`, `wikipedia`, `python`, `get_task_file`.
-        Answers are cached locally, so a failed submission can be retried without re-running the agent.
+        1. Add a Space secret named `GROQ_API_KEY` with your Groq API key (free at console.groq.com).
+        2. *Optional but recommended:* add `TAVILY_API_KEY` (free tier at tavily.com) for better search.
+        3. Optional env vars: `GROQ_MODELS`, `GROQ_VISION_MODEL`, `GROQ_WHISPER_MODEL`.
+        4. Log in to Hugging Face below and click **Run Evaluation & Submit All Answers**.
+
+        Tools: `web_search`, `fetch_url`, `wikipedia`, `python`, `get_task_file`,
+        `transcribe_audio`, `view_image`, `youtube_transcript`.
+        Answers are cached locally so a failed submission can be retried without re-running the agent.
         """
     )
 
@@ -618,6 +847,8 @@ if __name__ == "__main__":
 
     if not os.getenv("GROQ_API_KEY"):
         print("⚠️  GROQ_API_KEY is not set. Set it before running evaluation.")
+    if not os.getenv("TAVILY_API_KEY"):
+        print("ℹ️  TAVILY_API_KEY not set — search will use DuckDuckGo (less reliable).")
 
     print("-" * (60 + len(" App Starting ")) + "\n")
     demo.launch(debug=True, share=False)