Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import re
 import gradio as gr
 import requests
 import pandas as pd
@@ -6,183 +7,176 @@ import tempfile
 
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
-SYSTEM_PROMPT = """You are an expert question-answering agent participating in the GAIA benchmark.
+# ── file helpers ──────────────────────────────────────────────────────────────
 
-CRITICAL RULES:
-1. Your final response must be ONLY the answer — no explanation, no preamble, no "The answer is", no "FINAL ANSWER:"
-2. Be as concise as possible: just the name, number, word, or phrase requested
-3. For numbers: use digits (e.g., "42" not "forty-two") unless the question asks for words
-4. For lists: separate items with commas
-5. Always search the web for factual questions you're not 100% sure about
-6. For files (Excel, Python, etc.): download them and analyze carefully
-7. Think carefully before answering — exact match grading is used
-"""
+def download_task_file(task_id: str):
+    url = f"{DEFAULT_API_URL}/files/{task_id}"
+    try:
+        resp = requests.get(url, timeout=30)
+        if resp.status_code != 200:
+            return None, None
+        cd = resp.headers.get("content-disposition", "")
+        fname = "task_file"
+        if "filename=" in cd:
+            fname = cd.split("filename=")[-1].strip().strip('"')
+        ext = os.path.splitext(fname)[-1] or ".bin"
+        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=ext)
+        tmp.write(resp.content)
+        tmp.close()
+        return tmp.name, fname
+    except Exception as e:
+        print(f"File download error: {e}")
+        return None, None
 
-class BasicAgent:
-    def __init__(self):
-        print("Initializing GAIA Agent...")
-        from smolagents import CodeAgent, DuckDuckGoSearchTool, InferenceClientModel, tool
-        import urllib.request
-
-        # Tool to download task files from the API
-        @tool
-        def download_task_file(task_id: str) -> str:
-            """
-            Downloads a file associated with a GAIA task and saves it locally.
-            Returns the local file path where it was saved.
-            Args:
-                task_id: The GAIA task ID whose file you want to download
-            """
-            url = f"{DEFAULT_API_URL}/files/{task_id}"
-            try:
-                resp = requests.get(url, timeout=30)
-                if resp.status_code == 200:
-                    # Try to get filename from headers
-                    cd = resp.headers.get("content-disposition", "")
-                    fname = "task_file"
-                    if "filename=" in cd:
-                        fname = cd.split("filename=")[-1].strip().strip('"')
-                    
-                    ext = os.path.splitext(fname)[-1] or ".bin"
-                    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=ext)
-                    tmp.write(resp.content)
-                    tmp.close()
-                    return f"File saved to: {tmp.name} (original name: {fname})"
-                else:
-                    return f"No file found for task_id={task_id} (status {resp.status_code})"
-            except Exception as e:
-                return f"Error downloading file: {e}"
-
-        @tool
-        def read_excel_file(file_path: str) -> str:
-            """
-            Reads an Excel file and returns its contents as a string summary.
-            Args:
-                file_path: Local path to the Excel file
-            """
-            try:
-                df = pd.read_excel(file_path)
-                return f"Shape: {df.shape}\nColumns: {list(df.columns)}\n\nFirst rows:\n{df.head(20).to_string()}\n\nFull data:\n{df.to_string()}"
-            except Exception as e:
-                return f"Error reading Excel: {e}"
-
-        @tool
-        def read_python_file(file_path: str) -> str:
-            """
-            Reads a Python file and returns its contents.
-            Args:
-                file_path: Local path to the Python file
-            """
-            try:
-                with open(file_path, "r") as f:
-                    return f.read()
-            except Exception as e:
-                return f"Error reading file: {e}"
-
-        @tool  
-        def fetch_webpage(url: str) -> str:
-            """
-            Fetches the text content of a webpage.
-            Args:
-                url: The URL to fetch
-            """
-            try:
-                headers = {"User-Agent": "Mozilla/5.0"}
-                resp = requests.get(url, headers=headers, timeout=15)
-                resp.raise_for_status()
-                # Basic HTML stripping
-                text = resp.text
-                import re
-                text = re.sub(r'<[^>]+>', ' ', text)
-                text = re.sub(r'\s+', ' ', text).strip()
-                return text[:5000]
-            except Exception as e:
-                return f"Error fetching webpage: {e}"
-
-        # Try stronger models in order of preference
-        model = None
-        for model_id in [
-            "Qwen/Qwen2.5-72B-Instruct",
-            "meta-llama/Llama-3.3-70B-Instruct",
-            "mistralai/Mixtral-8x7B-Instruct-v0.1",
-            "meta-llama/Meta-Llama-3-8B-Instruct",
-        ]:
-            try:
-                from smolagents import InferenceClientModel
-                model = InferenceClientModel(model_id=model_id)
-                # Quick test
-                print(f"Using model: {model_id}")
-                break
-            except Exception as e:
-                print(f"Model {model_id} failed: {e}, trying next...")
-
-        if model is None:
-            raise RuntimeError("No model could be loaded!")
-
-        self.agent = CodeAgent(
-            tools=[
-                DuckDuckGoSearchTool(),
-                download_task_file,
-                read_excel_file,
-                read_python_file,
-                fetch_webpage,
-            ],
-            model=model,
-            max_steps=10,
-            verbosity_level=1,
-        )
-        print("Agent ready.")
 
-    def __call__(self, question: str, task_id: str = "") -> str:
-        print(f"\n{'='*60}")
-        print(f"Task ID: {task_id}")
-        print(f"Question: {question[:200]}")
-        print('='*60)
+def read_file_contents(local_path: str, fname: str) -> str:
+    ext = os.path.splitext(fname)[-1].lower()
+    try:
+        if ext in (".xlsx", ".xls"):
+            df = pd.read_excel(local_path)
+            return f"Excel file — shape {df.shape}\nColumns: {list(df.columns)}\n\n{df.to_string()}"
+        elif ext == ".csv":
+            df = pd.read_csv(local_path)
+            return f"CSV — shape {df.shape}\nColumns: {list(df.columns)}\n\n{df.to_string()}"
+        elif ext in (".py", ".txt", ".md", ".json"):
+            with open(local_path) as f:
+                return f.read()
+        else:
+            return f"Binary file: {fname}"
+    except Exception as e:
+        return f"Error reading file: {e}"
 
-        try:
-            # Include task_id in prompt so agent can download files if needed
-            task_context = f"[Task ID for file download: {task_id}]\n\n" if task_id else ""
-            
-            prompt = (
-                f"{task_context}"
-                f"Question: {question}\n\n"
-                f"Instructions: Answer with ONLY the final answer. "
-                f"No explanation. No prefix like 'The answer is'. "
-                f"Just the bare answer that would pass exact-match grading. "
-                f"If the question involves a file, use the download_task_file tool with the task ID above."
-            )
 
-            result = self.agent.run(prompt)
+def search_web(query: str) -> str:
+    try:
+        from duckduckgo_search import DDGS
+        with DDGS() as ddgs:
+            results = list(ddgs.text(query, max_results=5))
+        if not results:
+            return "No results found."
+        out = []
+        for r in results:
+            out.append(f"Title: {r.get('title','')}\nSnippet: {r.get('body','')}")
+        return "\n\n".join(out)
+    except Exception as e:
+        return f"Search error: {e}"
+
 
-            if result is None:
-                return ""
+def clean_answer(text: str) -> str:
+    text = text.strip()
+    for prefix in [
+        "FINAL ANSWER:", "Final Answer:", "Answer:",
+        "The answer is:", "The answer is",
+        "**Answer:**", "**Final Answer:**",
+    ]:
+        if text.lower().startswith(prefix.lower()):
+            text = text[len(prefix):].strip()
+    text = text.split("\n")[0].strip().strip('"').strip("'").strip("*").strip()
+    return text
 
-            answer = str(result).strip()
 
-            # Strip common unwanted prefixes
-            for prefix in [
-                "FINAL ANSWER:", "Final Answer:", "Answer:",
-                "The answer is:", "The answer is",
-                "**Answer:**", "**Final Answer:**",
-            ]:
-                if answer.lower().startswith(prefix.lower()):
-                    answer = answer[len(prefix):].strip()
+# ── Gemini agent ──────────────────────────────────────────────────────────────
+
+class BasicAgent:
+    def __init__(self):
+        self.api_key = os.getenv("GEMINI_API_KEY", "")
+        if not self.api_key:
+            raise RuntimeError(
+                "GEMINI_API_KEY secret is not set!\n"
+                "Get a FREE key at https://aistudio.google.com/app/apikey\n"
+                "Then add it in Space Settings → Secrets → GEMINI_API_KEY"
+            )
+        self.endpoint = (
+            f"https://generativelanguage.googleapis.com/v1beta/models/"
+            f"gemini-1.5-flash:generateContent?key={self.api_key}"
+        )
+        print("Agent ready (Gemini 1.5 Flash — free tier)")
+
+    def _ask_gemini(self, prompt: str) -> str:
+        body = {
+            "contents": [{"parts": [{"text": prompt}]}],
+            "generationConfig": {
+                "temperature": 0.1,
+                "maxOutputTokens": 512,
+            },
+            "systemInstruction": {
+                "parts": [{
+                    "text": (
+                        "You are a precise answer extraction agent for the GAIA benchmark. "
+                        "Reply with ONLY the final answer — no explanation, no preamble, "
+                        "no 'The answer is', no 'FINAL ANSWER'. "
+                        "Give the shortest correct answer: a name, number, word, or short phrase. "
+                        "Exact match grading is used."
+                    )
+                }]
+            }
+        }
+        resp = requests.post(self.endpoint, json=body, timeout=60)
+        resp.raise_for_status()
+        data = resp.json()
+        try:
+            return data["candidates"][0]["content"]["parts"][0]["text"].strip()
+        except (KeyError, IndexError) as e:
+            print(f"Gemini parse error: {e} | Response: {data}")
+            return ""
 
-            # Take only first line
-            answer = answer.split("\n")[0].strip()
-            # Strip surrounding quotes/asterisks
-            answer = answer.strip('"').strip("'").strip("*").strip()
+    def __call__(self, question: str, task_id: str = "") -> str:
+        print(f"\n{'='*60}\nTask: {task_id}\nQ: {question[:200]}")
+
+        # 1. Try to get attached file
+        file_context = ""
+        if task_id:
+            local_path, fname = download_task_file(task_id)
+            if local_path and fname:
+                contents = read_file_contents(local_path, fname)
+                file_context = f"\n\n--- Attached file ({fname}) ---\n{contents[:3000]}\n---\n"
+                print(f"Downloaded: {fname}")
+
+        # 2. Web search for factual questions
+        search_context = ""
+        search_triggers = [
+            "who", "what", "when", "where", "how many", "which", "name",
+            "album", "studio", "olympic", "pitcher", "actor", "article",
+            "specimen", "country", "competition", "veterinarian", "surname",
+            "wikipedia", "published", "season", "walks", "bats", "malko",
+            "vietnamese", "kuznetzov", "tamai", "petersen", "polish",
+        ]
+        if any(t in question.lower() for t in search_triggers) and not file_context:
+            print("Searching web...")
+            search_results = search_web(question[:200])
+            if "error" not in search_results.lower():
+                search_context = f"\n\n--- Web search results ---\n{search_results[:2000]}\n---\n"
+
+        # 3. Build prompt and ask Gemini
+        prompt = (
+            f"Question: {question}"
+            f"{file_context}"
+            f"{search_context}"
+            "\n\nAnswer with ONLY the final answer. No explanation whatsoever."
+        )
 
-            print(f"ANSWER: {answer}")
+        try:
+            answer = self._ask_gemini(prompt)
+
+            # If answer is too long, ask Gemini to extract just the key part
+            if len(answer.split()) > 15:
+                extract_prompt = (
+                    f"Extract ONLY the final short answer (name, number, or brief phrase) "
+                    f"from this text. Nothing else:\n\n{answer}"
+                )
+                answer = self._ask_gemini(extract_prompt)
+
+            answer = clean_answer(answer)
+            print(f"ANSWER: '{answer}'")
             return answer
 
         except Exception as e:
-            print(f"Agent error: {e}")
-            import traceback
-            traceback.print_exc()
+            print(f"Gemini error: {e}")
             return ""
 
 
+# ── gradio app ────────────────────────────────────────────────────────────────
+
 def run_and_submit_all(profile: gr.OAuthProfile | None):
     space_id = os.getenv("SPACE_ID", "")
 
@@ -190,18 +184,15 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
         return "⚠️ Please log in with Hugging Face first.", None
 
     username = profile.username
-    print(f"Logged in as: {username}")
-
-    api_url = DEFAULT_API_URL
     agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
 
     try:
         agent = BasicAgent()
-    except Exception as e:
-        return f"❌ Error initializing agent: {e}", None
+    except RuntimeError as e:
+        return f"❌ {e}", None
 
     try:
-        resp = requests.get(f"{api_url}/questions", timeout=15)
+        resp = requests.get(f"{DEFAULT_API_URL}/questions", timeout=15)
         resp.raise_for_status()
         questions_data = resp.json()
         print(f"Fetched {len(questions_data)} questions.")
@@ -217,34 +208,29 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
         if not task_id or not question_text:
             continue
 
-        print(f"\n[{i+1}/{len(questions_data)}] Task: {task_id}")
+        print(f"\n[{i+1}/{len(questions_data)}]")
         try:
             submitted_answer = agent(question_text, task_id=task_id)
         except Exception as e:
             submitted_answer = ""
-            print(f"Error on task {task_id}: {e}")
+            print(f"Error: {e}")
 
-        answers_payload.append({
-            "task_id": task_id,
-            "submitted_answer": submitted_answer
-        })
+        answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
         results_log.append({
             "Task ID": task_id,
-            "Question": question_text[:120] + "..." if len(question_text) > 120 else question_text,
+            "Question": question_text[:120] + ("..." if len(question_text) > 120 else ""),
             "Submitted Answer": submitted_answer,
         })
 
     if not answers_payload:
         return "❌ No answers generated.", None
 
-    submission_data = {
-        "username": username,
-        "agent_code": agent_code,
-        "answers": answers_payload,
-    }
-
     try:
-        resp = requests.post(f"{api_url}/submit", json=submission_data, timeout=120)
+        resp = requests.post(
+            f"{DEFAULT_API_URL}/submit",
+            json={"username": username, "agent_code": agent_code, "answers": answers_payload},
+            timeout=120,
+        )
         resp.raise_for_status()
         result = resp.json()
         status = (
@@ -254,7 +240,6 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
             f"({result.get('correct_count', '?')}/{result.get('total_attempted', '?')} correct)\n"
             f"Message: {result.get('message', '')}"
         )
-        print(status)
         return status, pd.DataFrame(results_log)
     except Exception as e:
         return f"❌ Submission failed: {e}", pd.DataFrame(results_log)
@@ -262,17 +247,17 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
 
 with gr.Blocks() as demo:
     gr.Markdown("# 🤖 GAIA Agent — Final Assignment")
-    gr.Markdown("Log in with your Hugging Face account, then click **Run & Submit** to evaluate your agent.")
-
+    gr.Markdown(
+        "### Setup (one-time, free!)\n"
+        "1. Get a **free** Gemini API key at [aistudio.google.com/app/apikey](https://aistudio.google.com/app/apikey)\n"
+        "2. In your Space → **Settings → Secrets** → add `GEMINI_API_KEY`\n"
+        "3. Log in below and click Run!"
+    )
     gr.LoginButton()
     run_button = gr.Button("🚀 Run Agent & Submit Answers", variant="primary")
     status_output = gr.Textbox(label="Status / Score", lines=6, interactive=False)
     results_table = gr.DataFrame(label="Results", wrap=True)
-
-    run_button.click(
-        fn=run_and_submit_all,
-        outputs=[status_output, results_table],
-    )
+    run_button.click(fn=run_and_submit_all, outputs=[status_output, results_table])
 
 if __name__ == "__main__":
     demo.launch()