Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 import re
+import base64
 import traceback
 import requests
 import gradio as gr
@@ -9,16 +10,14 @@ from google.genai import types
 
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
+# ── Helpers ───────────────────────────────────────────────────────────
 
-# ── Answer cleaner ────────────────────────────────────────────────────
 def clean_answer(text: str) -> str:
     if not text:
         return "N/A"
     text = text.strip()
-    # Strip code fences
     text = re.sub(r"^```[a-zA-Z]*\s*", "", text)
     text = re.sub(r"\s*```$", "", text)
-    # Strip label prefixes
     for prefix in [
         "final answer:", "answer:", "the answer is:",
         "the final answer is:", "result:",
@@ -26,114 +25,197 @@ def clean_answer(text: str) -> str:
         if text.lower().startswith(prefix):
             text = text[len(prefix):].strip()
             break
-    return " ".join(text.split())[:200]
+    return " ".join(text.split())[:300]
 
 
 def extract_text(response) -> str:
     """
-    Pull the final answer text out of a GenerateContentResponse.
-    When google_search is used, Gemini emits:
-      [search_tool_use part] → [search_tool_result part] → [text part]
-    We want the LAST text part.
+    Collect ALL text and code-execution output parts, return the last one.
+    Handles: plain text, google_search grounding, code_execution output.
     """
-    text_parts = []
+    parts_text = []
     try:
         for candidate in response.candidates:
             for part in candidate.content.parts:
+                # Plain text part
                 if hasattr(part, "text") and part.text and part.text.strip():
-                    text_parts.append(part.text.strip())
+                    parts_text.append(part.text.strip())
+                # Code execution result
+                if hasattr(part, "code_execution_result") and part.code_execution_result:
+                    out = getattr(part.code_execution_result, "output", "")
+                    if out and str(out).strip():
+                        parts_text.append(str(out).strip())
     except Exception as e:
         print(f"  [extract_text error] {e}")
+    if parts_text:
+        return parts_text[-1]
+    try:
+        t = response.text
+        return t.strip() if t else ""
+    except Exception:
+        return ""
+
+
+def fetch_file(task_id: str) -> tuple[bytes | None, str]:
+    """
+    Download the file attached to a GAIA task.
+    Returns (file_bytes, filename) or (None, "").
+    """
+    url = f"{DEFAULT_API_URL}/files/{task_id}"
+    try:
+        r = requests.get(url, timeout=20)
+        if r.status_code == 200:
+            # Try to get filename from Content-Disposition header
+            cd = r.headers.get("Content-Disposition", "")
+            fname = ""
+            if "filename=" in cd:
+                fname = cd.split("filename=")[-1].strip().strip('"')
+            if not fname:
+                fname = f"attachment_{task_id}"
+            print(f"  [file] Downloaded {len(r.content)} bytes  ({fname})")
+            return r.content, fname
+    except Exception as e:
+        print(f"  [file] Download failed: {e}")
+    return None, ""
 
-    if text_parts:
-        return text_parts[-1]   # last text part = final answer after search
 
-    # Last resort: .text property
+def build_contents(question: str, file_bytes: bytes | None, fname: str) -> list:
+    """
+    Build the `contents` list for generate_content(), optionally including a file.
+    Supports: images, PDFs, plain-text, CSV, audio.
+    """
+    if file_bytes is None:
+        return [question]
+
+    ext = fname.rsplit(".", 1)[-1].lower() if "." in fname else ""
+
+    # --- Image ---
+    if ext in ("png", "jpg", "jpeg", "gif", "webp", "bmp"):
+        mime = {"png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg",
+                "gif": "image/gif", "webp": "image/webp", "bmp": "image/bmp"}.get(ext, "image/png")
+        return [
+            types.Part.from_bytes(data=file_bytes, mime_type=mime),
+            question,
+        ]
+
+    # --- PDF ---
+    if ext == "pdf":
+        return [
+            types.Part.from_bytes(data=file_bytes, mime_type="application/pdf"),
+            question,
+        ]
+
+    # --- Audio ---
+    if ext in ("mp3", "wav", "ogg", "flac", "m4a"):
+        mime = {"mp3": "audio/mpeg", "wav": "audio/wav", "ogg": "audio/ogg",
+                "flac": "audio/flac", "m4a": "audio/mp4"}.get(ext, "audio/mpeg")
+        return [
+            types.Part.from_bytes(data=file_bytes, mime_type=mime),
+            question,
+        ]
+
+    # --- Plain text / CSV / code — embed as text context ---
     try:
-        return response.text.strip()
+        text_content = file_bytes.decode("utf-8", errors="replace")
+        augmented = (
+            f"The following file ({fname}) is attached:\n\n"
+            f"```\n{text_content[:8000]}\n```\n\n"
+            f"{question}"
+        )
+        return [augmented]
     except Exception:
-        return ""
+        pass
+
+    # Last resort: base64-encode unknown binary
+    b64 = base64.b64encode(file_bytes).decode()
+    return [f"[File {fname} base64]: {b64[:500]}…\n\n{question}"]
 
 
 # ── Agent ─────────────────────────────────────────────────────────────
+
 class BasicAgent:
 
     SYSTEM = """You are an expert assistant solving GAIA benchmark questions.
 
-STRICT OUTPUT FORMAT:
-- Return ONLY the final answer. No explanation. No preamble. No punctuation tail.
-- Numbers: digits only, no commas (e.g. 1234567 not 1,234,567). Omit trailing .0.
-- Lists: comma-separated, in the order the question requests (default: alphabetical).
-- Yes/no questions: exactly "yes" or "no" (lowercase).
-- Dates: use the format the question implies (e.g. "January 5, 1990" or "1990-01-05").
-- Names: full name unless the question asks for first/last only.
+RULES:
+- Use Google Search for any factual, date, count, or real-world question.
+- If a file is attached, read it carefully before answering.
+- Think step-by-step internally, but output ONLY the final answer.
 
-Use Google Search whenever you need facts, dates, counts, or external data.
-Think carefully step by step, then output ONLY the answer."""
+STRICT OUTPUT FORMAT:
+- Return ONLY the final answer. No explanation, no preamble, no trailing punctuation.
+- Numbers: digits only, no commas (1234567 not 1,234,567). Drop trailing .0 from whole numbers.
+- Lists: comma-separated, alphabetical order unless stated otherwise.
+- Yes/no: exactly "yes" or "no" (lowercase).
+- Names: full name unless the question asks for first or last name only.
+- Dates: match the format implied by the question."""
 
     def __init__(self):
         api_key = os.getenv("GEMINI_API_KEY")
         if not api_key:
             raise EnvironmentError("GEMINI_API_KEY not set.")
-
         self.client   = genai.Client(api_key=api_key)
-        self.model_id = "gemini-2.0-flash"          # stable + has search grounding
+        self.model_id = "gemini-2.0-flash"
 
-        # Search-enabled config (separate from code-exec to avoid conflicts)
         self.search_config = types.GenerateContentConfig(
             system_instruction=self.SYSTEM,
             tools=[types.Tool(google_search=types.GoogleSearch())],
-            temperature=0,      # deterministic
+            temperature=0,
         )
-
-        # Code-execution config (for maths / data questions)
         self.code_config = types.GenerateContentConfig(
             system_instruction=self.SYSTEM,
             tools=[types.Tool(code_execution=types.ToolCodeExecution())],
             temperature=0,
         )
-
-        # Plain config for the verify pass
         self.plain_config = types.GenerateContentConfig(
-            system_instruction=(
-                "You are a strict answer formatter. "
-                "Return ONLY the final answer, no explanation."
-            ),
+            system_instruction=self.SYSTEM,
             temperature=0,
         )
+        print(f"✅ BasicAgent ready ({self.model_id})")
 
-        print(f"✅ BasicAgent ready  (model={self.model_id})")
-
-    # ------------------------------------------------------------------
-    def __call__(self, question: str) -> str:
+    def __call__(self, question: str, task_id: str = "") -> str:
         try:
-            return self._run(question)
+            return self._run(question, task_id)
         except Exception as exc:
             print(f"  [FATAL] {exc}\n{traceback.format_exc()}")
             return "N/A"
 
-    def _run(self, question: str) -> str:
-        q_lower = question.lower()
+    def _run(self, question: str, task_id: str) -> str:
+        # Try to download attached file
+        file_bytes, fname = (None, "")
+        if task_id:
+            file_bytes, fname = fetch_file(task_id)
 
-        # Route to code-exec for pure maths / counting questions
-        needs_code = any(kw in q_lower for kw in [
-            "calculate", "compute", "how many", "sum of", "average",
-            "percentage", "multiply", "divide", "convert",
+        contents = build_contents(question, file_bytes, fname)
+        q_lower  = question.lower()
+
+        # Choose config
+        has_file    = file_bytes is not None
+        needs_code  = any(kw in q_lower for kw in [
+            "calculate", "compute", "sum of", "average", "percentage",
+            "multiply", "divide", "convert", "how many",
         ])
-        config = self.code_config if needs_code else self.search_config
 
-        resp  = self.client.models.generate_content(
-            model=self.model_id, contents=question, config=config,
+        if has_file:
+            # Files: use plain config (search can't see the file)
+            config = self.plain_config
+        elif needs_code:
+            config = self.code_config
+        else:
+            config = self.search_config
+
+        resp = self.client.models.generate_content(
+            model=self.model_id, contents=contents, config=config,
         )
-        raw   = extract_text(resp)
-        ans   = clean_answer(raw)
-        print(f"  [raw]  {raw[:120]!r}")
+        raw = extract_text(resp)
+        ans = clean_answer(raw)
+        print(f"  [raw]  {raw[:150]!r}")
         print(f"  [ans]  {ans!r}")
 
+        # Fallback: plain call if empty
         if not ans or ans == "N/A":
-            # Fallback: try plain (no tools)
             resp2 = self.client.models.generate_content(
-                model=self.model_id, contents=question, config=self.plain_config,
+                model=self.model_id, contents=contents, config=self.plain_config,
             )
             ans = clean_answer(extract_text(resp2))
             print(f"  [fallback] {ans!r}")
@@ -141,69 +223,87 @@ Think carefully step by step, then output ONLY the answer."""
         return ans or "N/A"
 
 
-# ── Debug helper – fetch & show first N answers without submitting ────
-def debug_first_n(n: int = 3):
-    """Call this from __main__ to see raw vs cleaned answers."""
-    agent = BasicAgent()
-    resp  = requests.get(f"{DEFAULT_API_URL}/questions", timeout=15)
-    questions = resp.json()[:n]
-    for item in questions:
-        print("\n" + "="*60)
-        print(f"Q: {item['question']}")
-        ans = agent(item["question"])
-        print(f"FINAL SUBMITTED: {ans!r}")
-
-
-# ── Gradio runner ─────────────────────────────────────────────────────
-def run_and_submit_all(profile: gr.OAuthProfile | None):
-    if not profile:
-        return "Please log in to Hugging Face first.", None
-
-    username  = profile.username
-    space_id  = os.getenv("SPACE_ID", "")
-    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
+# ── Shared run logic ──────────────────────────────────────────────────
 
-    try:
-        agent = BasicAgent()
-    except Exception as e:
-        return f"Error initialising agent: {e}", None
-
-    # Fetch questions
+def _fetch_questions() -> list | str:
     try:
         r = requests.get(f"{DEFAULT_API_URL}/questions", timeout=15)
         r.raise_for_status()
-        questions_data = r.json()
-        print(f"Fetched {len(questions_data)} questions.")
+        data = r.json()
+        print(f"Fetched {len(data)} questions.")
+        return data
     except Exception as e:
-        return f"Error fetching questions: {e}", None
+        return f"Error fetching questions: {e}"
 
-    results_log, answers_payload = [], []
 
+def _run_agent_on_questions(agent, questions_data: list) -> tuple[list, list]:
+    results_log, answers_payload = [], []
     for item in questions_data:
-        task_id = item.get("task_id")
-        q_text  = item.get("question")
+        task_id = item.get("task_id", "")
+        q_text  = item.get("question", "")
         if not task_id or q_text is None:
             continue
-
         print(f"\n── {task_id} ──")
         print(f"Q: {q_text[:150]}")
         try:
-            answer = agent(q_text)
+            answer = agent(q_text, task_id)
         except Exception as e:
             answer = "N/A"
             print(f"  ERROR: {e}")
-
+        print(f"→ {answer!r}")
         answers_payload.append({"task_id": task_id, "submitted_answer": answer})
         results_log.append({
-            "Task ID": task_id,
-            "Question": q_text[:120],
+            "Task ID":          task_id,
+            "Question":         q_text[:120],
             "Submitted Answer": answer,
+            "Has File":         "yes" if item.get("file_name") else "no",
         })
+    return results_log, answers_payload
+
+
+# ── Dry Run (no submit) ───────────────────────────────────────────────
+
+def dry_run(profile: gr.OAuthProfile | None):
+    if not profile:
+        return "Please log in first.", None
+    try:
+        agent = BasicAgent()
+    except Exception as e:
+        return f"Agent init error: {e}", None
+
+    data = _fetch_questions()
+    if isinstance(data, str):
+        return data, None
+
+    sample = data[:3]   # first 3 questions only
+    logs, _ = _run_agent_on_questions(agent, sample)
+    return f"Dry run complete — showing {len(logs)} answers (not submitted).", pd.DataFrame(logs)
+
+
+# ── Full Run + Submit ─────────────────────────────────────────────────
+
+def run_and_submit_all(profile: gr.OAuthProfile | None):
+    if not profile:
+        return "Please log in to Hugging Face first.", None
+
+    username   = profile.username
+    space_id   = os.getenv("SPACE_ID", "")
+    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
+
+    try:
+        agent = BasicAgent()
+    except Exception as e:
+        return f"Error initialising agent: {e}", None
+
+    data = _fetch_questions()
+    if isinstance(data, str):
+        return data, None
+
+    results_log, answers_payload = _run_agent_on_questions(agent, data)
 
     if not answers_payload:
         return "No answers produced.", pd.DataFrame(results_log)
 
-    # Submit
     try:
         r = requests.post(
             f"{DEFAULT_API_URL}/submit",
@@ -229,26 +329,29 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
 
 
 # ── UI ────────────────────────────────────────────────────────────────
+
 with gr.Blocks() as demo:
-    gr.Markdown("# GAIA Agent — Gemini 2.0 Flash + Search")
+    gr.Markdown("# GAIA Agent — Gemini 2.0 Flash + Search + File Support")
     gr.Markdown(
-        "1. Set `GEMINI_API_KEY` in Space secrets.\n"
-        "2. Log in below.\n"
-        "3. Click **Run Evaluation**."
+        "**Setup:** Add `GEMINI_API_KEY` to Space secrets, then log in.\n\n"
+        "- **Dry Run** — answers the first 3 questions and shows results *without* submitting.\n"
+        "- **Full Run** — answers all questions and submits to the leaderboard."
     )
     gr.LoginButton()
-    btn     = gr.Button("▶ Run Evaluation & Submit All Answers", variant="primary")
-    status  = gr.Textbox(label="Status", lines=6, interactive=False)
-    table   = gr.DataFrame(label="Answers", wrap=True)
-    btn.click(fn=run_and_submit_all, outputs=[status, table])
+
+    with gr.Row():
+        dry_btn  = gr.Button("🔍 Dry Run (first 3, no submit)", variant="secondary")
+        full_btn = gr.Button("▶ Full Run & Submit All",         variant="primary")
+
+    status = gr.Textbox(label="Status", lines=6, interactive=False)
+    table  = gr.DataFrame(label="Answers", wrap=True)
+
+    dry_btn.click( fn=dry_run,            outputs=[status, table])
+    full_btn.click(fn=run_and_submit_all, outputs=[status, table])
 
 
 if __name__ == "__main__":
     print("\n── Env check ──")
     for v in ("SPACE_HOST", "SPACE_ID", "GEMINI_API_KEY"):
         print(f"  {v}: {'SET' if os.getenv(v) else 'MISSING'}")
-
-    # Uncomment to preview first 3 answers before submitting:
-    # debug_first_n(3)
-
     demo.launch(debug=True, share=False)