Update app.py

app.py

@@ -1,357 +1,113 @@
 import os
 import re
-import base64
-import traceback
-import requests
-import gradio as gr
-import pandas as pd
-from google import genai
-from google.genai import types
-
-DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
-
-# ── Helpers ───────────────────────────────────────────────────────────
-
-def clean_answer(text: str) -> str:
-    if not text:
-        return "N/A"
-    text = text.strip()
-    text = re.sub(r"^```[a-zA-Z]*\s*", "", text)
-    text = re.sub(r"\s*```$", "", text)
-    for prefix in [
-        "final answer:", "answer:", "the answer is:",
-        "the final answer is:", "result:",
-    ]:
-        if text.lower().startswith(prefix):
-            text = text[len(prefix):].strip()
-            break
-    return " ".join(text.split())[:300]
-
-
-def extract_text(response) -> str:
-    """
-    Collect ALL text and code-execution output parts, return the last one.
-    Handles: plain text, google_search grounding, code_execution output.
-    """
-    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():
-                    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, ""
-
-
-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:
-        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:
-        pass
-
-    # Last resort: base64-encode unknown binary
-    b64 = base64.b64encode(file_bytes).decode()
-    return [f"[File {fname} base64]: {b64[:500]}…\n\n{question}"]
-
-
-# ── Agent ─────────────────────────────────────────────────────────────
+import google.generativeai as genai
 
 class BasicAgent:
 
-    SYSTEM = """You are an expert assistant solving GAIA benchmark questions.
+    def __init__(self):
+        api_key = os.getenv("GEMINI_API_KEY")
+        if not api_key:
+            raise ValueError("Missing GEMINI_API_KEY")
+
+        genai.configure(api_key=api_key)
+
+        # ✅ stable model
+        self.model = genai.GenerativeModel("gemini-1.5-flash")
+
+        print("✅ Gemini Agent Ready")
+
+    def clean(self, text: str) -> str:
+        if not text:
+            return "N/A"
+
+        text = text.strip()
+        text = text.replace("Final Answer:", "").strip()
+        text = re.sub(r"\n.*", "", text)  # first line only
+        text = text.strip()
 
-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.
+        return text[:100]
 
-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."""
+    import os
+import re
+import google.generativeai as genai
+
+class BasicAgent:
 
     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"
-
-        self.search_config = types.GenerateContentConfig(
-            system_instruction=self.SYSTEM,
-            tools=[types.Tool(google_search=types.GoogleSearch())],
-            temperature=0,
-        )
-        self.code_config = types.GenerateContentConfig(
-            system_instruction=self.SYSTEM,
-            tools=[types.Tool(code_execution=types.ToolCodeExecution())],
-            temperature=0,
-        )
-        self.plain_config = types.GenerateContentConfig(
-            system_instruction=self.SYSTEM,
-            temperature=0,
-        )
-        print(f"✅ BasicAgent ready ({self.model_id})")
+            raise ValueError("Missing GEMINI_API_KEY")
 
-    def __call__(self, question: str, task_id: str = "") -> str:
-        try:
-            return self._run(question, task_id)
-        except Exception as exc:
-            print(f"  [FATAL] {exc}\n{traceback.format_exc()}")
+        genai.configure(api_key=api_key)
+
+        self.model = genai.GenerativeModel("gemini-1.5-flash")
+
+        print("✅ PASS Agent Ready")
+
+    def clean(self, text: str) -> str:
+        if not text:
             return "N/A"
 
-    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)
-
-        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",
-        ])
-
-        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[:150]!r}")
-        print(f"  [ans]  {ans!r}")
-
-        # Fallback: plain call if empty
-        if not ans or ans == "N/A":
-            resp2 = self.client.models.generate_content(
-                model=self.model_id, contents=contents, config=self.plain_config,
-            )
-            ans = clean_answer(extract_text(resp2))
-            print(f"  [fallback] {ans!r}")
-
-        return ans or "N/A"
-
-
-# ── Shared run logic ──────────────────────────────────────────────────
-
-def _fetch_questions() -> list | str:
-    try:
-        r = requests.get(f"{DEFAULT_API_URL}/questions", timeout=15)
-        r.raise_for_status()
-        data = r.json()
-        print(f"Fetched {len(data)} questions.")
-        return data
-    except Exception as e:
-        return f"Error fetching questions: {e}"
-
-
-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", "")
-        if not task_id or q_text is None:
-            continue
-        print(f"\n── {task_id} ──")
-        print(f"Q: {q_text[:150]}")
+        text = text.strip()
+
+        # remove junk
+        text = text.replace("Final Answer:", "")
+        text = text.replace("Answer:", "")
+        text = text.strip()
+
+        # keep only first line
+        text = text.split("\n")[0]
+
+        # remove sentences
+        text = text.split(".")[0]
+
+        return text.strip()[:100]
+
+    def __call__(self, question: str, task_id: str = "") -> str:
         try:
-            answer = agent(q_text, task_id)
+            # 🔥 STEP 1 — Solve
+            prompt1 = f"""
+Solve the question.
+
+IMPORTANT:
+- Think step by step internally
+- Return ONLY final answer
+- No explanation
+
+Question:
+{question}
+"""
+            r1 = self.model.generate_content(prompt1)
+            ans1 = r1.text if hasattr(r1, "text") else ""
+
+            ans1 = self.clean(ans1)
+
+            # 🔥 STEP 2 — VERIFY (this is the magic)
+            prompt2 = f"""
+Question: {question}
+
+Proposed Answer: {ans1}
+
+Check if this is correct.
+If wrong, fix it.
+
+Return ONLY final answer.
+No explanation.
+"""
+            r2 = self.model.generate_content(prompt2)
+            final = r2.text if hasattr(r2, "text") else ""
+
+            final = self.clean(final)
+
+            # 🔥 fallback if broken
+            if not final or final == "N/A":
+                final = ans1
+
+            print("Q:", question[:80])
+            print("A:", final)
+
+            return final
+
         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],
-            "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)
-
-    try:
-        r = requests.post(
-            f"{DEFAULT_API_URL}/submit",
-            json={"username": username, "agent_code": agent_code, "answers": answers_payload},
-            timeout=120,
-        )
-        r.raise_for_status()
-        res = r.json()
-        status = (
-            f"✅ Submission successful!\n"
-            f"User:  {res.get('username')}\n"
-            f"Score: {res.get('score', 'N/A')}% "
-            f"({res.get('correct_count', '?')}/{res.get('total_attempted', '?')} correct)\n"
-            f"Msg:   {res.get('message', '')}"
-        )
-    except requests.exceptions.HTTPError as e:
-        status = f"❌ HTTP {e.response.status_code}: {e.response.text[:300]}"
-    except Exception as e:
-        status = f"❌ Error: {e}"
-
-    print(status)
-    return status, pd.DataFrame(results_log)
-
-
-# ── UI ────────────────────────────────────────────────────────────────
-
-with gr.Blocks() as demo:
-    gr.Markdown("# GAIA Agent — Gemini 2.0 Flash + Search + File Support")
-    gr.Markdown(
-        "**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()
-
-    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'}")
-    demo.launch(debug=True, share=False)
+            print("ERROR:", e)
+            return "N/A"