app.py

import os
import re
import json
import time
import pathlib

import gradio as gr
import requests
import pandas as pd

from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    LiteLLMModel,
    VisitWebpageTool,
    tool,
)

# ZeroGPU Spaces refuse to start without a @spaces.GPU function. Harmless stub.
# This agent only needs plain CPU hardware, so on CPU basic you can delete this block.
try:
    import spaces

    @spaces.GPU(duration=1)
    def zero_gpu_stub():
        return None
except ImportError:
    pass

# (Keep Constants as is)
# --- Constants ---
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"

# --- Config (override via Space *Variables*; GEMINI_API_KEY must be a Space *secret*) ---
GEMINI_MODEL = os.getenv("GEMINI_MODEL", "gemini-flash-latest").removeprefix("gemini/")
TEMPERATURE = os.getenv("TEMPERATURE")                       # unset = provider default
MIN_CALL_GAP = float(os.getenv("MIN_CALL_GAP", "12"))        # seconds between model calls (12s ~ 5 RPM)
SKIP_MODEL_CHECK = os.getenv("SKIP_MODEL_CHECK", "0") == "1" # skip the startup "does this model exist" check
MAX_STEPS = int(os.getenv("MAX_STEPS", "8"))                 # more steps = more requests
MAX_PAGE_CHARS = int(os.getenv("MAX_PAGE_CHARS", "15000"))
MAX_FILE_CHARS = int(os.getenv("MAX_FILE_CHARS", "15000"))
SKIP_MEDIA = os.getenv("SKIP_MEDIA", "1") == "1"             # skip audio/image/video/YouTube questions
DRY_RUN_LIMIT = int(os.getenv("DRY_RUN_LIMIT", "0"))         # >0: run only N questions, do NOT submit

CACHE_PATH = pathlib.Path("answers_cache.json")
FILES_DIR = pathlib.Path("/tmp/gaia_files")
FILES_DIR.mkdir(parents=True, exist_ok=True)

MEDIA_EXT = (".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp",
             ".mp3", ".wav", ".m4a", ".ogg", ".mp4", ".mov", ".avi")

ANSWER_RULES = """You are a research agent solving GAIA benchmark questions. Work step by step: write Python code, use your tools, verify, then answer.

FINAL ANSWER RULES (scored by EXACT string match):
- Call final_answer() exactly once, with ONLY the answer as a short string or number. Never a sentence, never an explanation, never the prefix "FINAL ANSWER:".
- Number: digits only. No thousands separators, no units, no currency or % symbols unless the question asks for them.
- Text: as few words as possible; no leading articles (a/an/the); no abbreviations unless the question asks for them.
- List: comma-separated values, in the order the question asks for.
- Obey any format the question specifies (rounding, alphabetical order, case, "first name only", etc.).

WORKING RULES:
- Do all arithmetic, counting, sorting and string manipulation in Python code, never in your head.
- Search snippets are not evidence. Use web_search to find pages, then visit_webpage on the best 1-2 pages and read the actual data. Prefer Wikipedia and primary sources.
- If a tool errors or returns nothing, change the query or the source instead of repeating the same call.
- If the question contains reversed or odd-looking text, decode it with Python first.
- You have a small step budget. If you cannot fully verify the answer, still call final_answer() with your best guess. Never finish without calling final_answer().
"""


# --- Errors that should stop the WHOLE run (not just one question) ---
class RunAborted(RuntimeError):
    pass


class QuotaExhausted(RunAborted):
    pass


class FatalModelError(RunAborted):
    pass


# --- Error classification ---
FATAL_MARKERS = ("notfounderror", "no longer available", "is not found", "not_found",
                 "authenticationerror", "api key not valid", "api_key_invalid",
                 "permission_denied", "permissiondeniederror", "invalid api key")
TRANSIENT_MARKERS = ("429", "ratelimit", "rate limit", "resource_exhausted", "502", "503", "504",
                     "overloaded", "unavailable", "timeout", "timed out",
                     "internalservererror", "connection")


def _is_daily_quota_error(e: Exception) -> bool:
    flat = re.sub(r"[^a-z]", "", str(e).lower())
    return "perday" in flat          # e.g. GenerateRequestsPerDayPerProjectPerModel-FreeTier


def _is_fatal(e: Exception) -> bool:
    s = str(e).lower()
    return any(k in s for k in FATAL_MARKERS)


def _is_transient(e: Exception) -> bool:
    s = str(e).lower()
    return any(k in s for k in TRANSIENT_MARKERS)


def _retry_delay(e: Exception, default: float) -> float:
    m = re.search(r"retry[^0-9]{0,30}(\d+(?:\.\d+)?)\s*s", str(e), re.I)
    return min(float(m.group(1)) + 1, 90.0) if m else default


# --- Helpers ---
def clean_answer(raw) -> str:
    """Normalise whatever the agent returned into a bare, exact-match-friendly string."""
    if isinstance(raw, (list, tuple)):
        raw = ", ".join(str(x) for x in raw)
    if isinstance(raw, float) and raw.is_integer():
        raw = int(raw)
    text = str(raw).strip()
    text = re.sub(r"(?i)^\s*final answer\s*[:\-]\s*", "", text)
    text = text.strip().strip("`").strip()
    if len(text) >= 2 and text[0] == text[-1] and text[0] in "\"'":
        text = text[1:-1].strip()
    text = " ".join(text.split())                       # collapse newlines / double spaces
    text = text.rstrip(".").strip()
    if re.fullmatch(r"-?\d{1,3}(,\d{3})+(\.\d+)?", text):   # 1,234 -> 1234
        text = text.replace(",", "")
    if re.fullmatch(r"-?\d+\.0+", text):                    # 12.0 -> 12
        text = text.split(".")[0]
    return text


def needs_media(question: str, file_name: str) -> bool:
    if file_name and file_name.lower().endswith(MEDIA_EXT):
        return True
    return bool(re.search(r"youtube\.com|youtu\.be", question, re.I))


def load_cache() -> dict:
    try:
        return json.loads(CACHE_PATH.read_text())
    except Exception:
        return {}


def save_cache(cache: dict) -> None:
    try:
        CACHE_PATH.write_text(json.dumps(cache, indent=2))
    except Exception as e:
        print(f"Could not write cache: {e}")


def download_attachment(api_url: str, task_id: str, file_name: str) -> str | None:
    dest = FILES_DIR / f"{task_id}__{file_name}"
    if dest.exists():
        return str(dest)
    try:
        r = requests.get(f"{api_url}/files/{task_id}", timeout=60)
        r.raise_for_status()
        dest.write_bytes(r.content)
        return str(dest)
    except Exception as e:
        print(f"Attachment download failed for {task_id}: {e}")
        return None


def check_gemini_model(model_id: str, api_key: str) -> None:
    """Fail fast with a readable message if the key or model name is bad.
    (These are metadata calls; they don't use generation quota.)"""
    base = "https://generativelanguage.googleapis.com/v1beta"
    headers = {"x-goog-api-key": api_key}
    try:
        r = requests.get(f"{base}/models/{model_id}", headers=headers, timeout=20)
    except Exception as e:
        print(f"Model preflight skipped (network error): {e}")
        return
    if r.status_code == 200:
        print(f"Model preflight OK: {model_id}")
        return
    if r.status_code in (400, 401, 403):
        raise RuntimeError(f"Gemini rejected your API key (HTTP {r.status_code}): {r.text[:200]}")

    names = []
    try:
        lr = requests.get(f"{base}/models?pageSize=200", headers=headers, timeout=20)
        for m in lr.json().get("models", []):
            if "generateContent" in m.get("supportedGenerationMethods", []) and "flash" in m["name"]:
                names.append(m["name"].replace("models/", ""))
    except Exception:
        pass
    if model_id in names:
        return
    raise RuntimeError(
        f"Model '{model_id}' is not available (HTTP {r.status_code}). Set the GEMINI_MODEL Space "
        f"variable to one of: {', '.join(names[:12]) or '(could not list models)'}. "
        "Set SKIP_MODEL_CHECK=1 to bypass this check."
    )


# --- Tools ---
class RobustDuckDuckGoSearchTool(DuckDuckGoSearchTool):
    """DuckDuckGoSearchTool + retries (DDG often rate-limits cloud IPs, including HF Spaces)."""

    def forward(self, query: str) -> str:
        last_err = None
        for attempt in range(3):
            try:
                return super().forward(query)
            except Exception as e:
                last_err = e
                if "no results" in str(e).lower():
                    return f"No results for '{query}'. Try shorter or different keywords."
                time.sleep(3 * (attempt + 1))
        return (f"web_search failed after retries ({last_err}). "
                "Try different keywords, or use visit_webpage on a known URL.")


@tool
def read_file(path: str) -> str:
    """Read a locally saved question attachment (txt, py, csv, json, md, xlsx, pdf) and return its content as text.

    Args:
        path: Local file path that was given in the task prompt.
    """
    p = pathlib.Path(path)
    if not p.exists():
        return f"File not found: {path}"
    ext = p.suffix.lower()
    try:
        if ext in (".xlsx", ".xlsm", ".xls"):
            sheets = pd.read_excel(p, sheet_name=None, header=None)
            chunks = []
            for name, df in sheets.items():
                chunks.append(f"## Sheet: {name} ({len(df)} rows x {df.shape[1]} cols)\n"
                              f"{df.to_csv(index=False, header=False)}")
            text = "\n".join(chunks)
        elif ext == ".pdf":
            from pypdf import PdfReader
            text = "\n".join((pg.extract_text() or "") for pg in PdfReader(str(p)).pages)
        elif ext in MEDIA_EXT:
            return "This is an audio/image/video file; it cannot be read as text."
        else:
            text = p.read_text(errors="replace")
    except Exception as e:
        return f"Could not read {path}: {e}"
    if len(text) > MAX_FILE_CHARS:
        text = text[:MAX_FILE_CHARS] + f"\n...[truncated at {MAX_FILE_CHARS} chars]"
    return text


# --- Model ---
class _Guard:
    """Mixin for any smolagents model: spaces out calls (free-tier RPM), retries per-minute 429s,
    and records daily-quota / fatal errors so the runner can stop instead of failing silently.
    (smolagents swallows model errors inside its step loop, so flags are the only reliable signal.)"""

    quota_exhausted = False
    fatal_error = None
    _last_call = 0.0

    def generate(self, *args, **kwargs):
        if self.quota_exhausted:
            raise RuntimeError("Free daily quota exhausted; not calling the API again.")
        if self.fatal_error:
            raise RuntimeError(self.fatal_error)
        for attempt in range(5):
            gap = MIN_CALL_GAP - (time.time() - self._last_call)
            if gap > 0:
                time.sleep(gap)
            self._last_call = time.time()
            try:
                return super().generate(*args, **kwargs)
            except Exception as e:
                if _is_daily_quota_error(e):
                    self.quota_exhausted = True
                    raise
                if _is_fatal(e):
                    self.fatal_error = str(e)[:300]
                    raise
                if attempt < 4 and _is_transient(e):
                    delay = _retry_delay(e, default=10.0 * (attempt + 1))
                    print(f"Transient model error, retrying in {delay:.0f}s: {str(e)[:120]}")
                    time.sleep(delay)
                    continue
                raise


class GuardedLiteLLMModel(_Guard, LiteLLMModel):
    pass


# --- Basic Agent Definition ---
# ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
class BasicAgent:
    """Same name and call interface as the template: agent(question) -> str."""

    def __init__(self):
        api_key = os.getenv("GEMINI_API_KEY")
        if not api_key:
            raise RuntimeError("GEMINI_API_KEY is not set. Add it as a Space *secret*.")
        if not SKIP_MODEL_CHECK:
            check_gemini_model(GEMINI_MODEL, api_key)

        model_kwargs = dict(model_id=f"gemini/{GEMINI_MODEL}", api_key=api_key)
        if TEMPERATURE:
            model_kwargs["temperature"] = float(TEMPERATURE)
        self.model = GuardedLiteLLMModel(**model_kwargs)

        visit = VisitWebpageTool()
        visit.max_output_length = MAX_PAGE_CHARS   # honored by recent versions; harmless otherwise

        base = dict(
            tools=[RobustDuckDuckGoSearchTool(max_results=6, rate_limit=2.0), visit, read_file],
            model=self.model,
            additional_authorized_imports=["pandas", "numpy", "json", "csv", "io",
                                           "string", "fractions", "decimal", "openpyxl"],
            max_steps=MAX_STEPS,
            verbosity_level=1,
        )
        # Answer rules go in the system prompt when this smolagents version supports
        # `instructions`; otherwise they are prepended to every task.
        try:
            self.agent = CodeAgent(**base, instructions=ANSWER_RULES, max_print_outputs_length=4000)
            self.rules_in_system = True
        except TypeError:
            self.agent = CodeAgent(**base)
            self.rules_in_system = False
        print(f"BasicAgent initialized | model=gemini/{GEMINI_MODEL} | max_steps={MAX_STEPS} | "
              f"call_gap={MIN_CALL_GAP}s | rules_in_system={self.rules_in_system}")

    def _abort_if_model_dead(self, exc: Exception | None = None) -> None:
        if self.model.fatal_error:
            raise FatalModelError(f"Gemini call failed and cannot be retried: {self.model.fatal_error}")
        if self.model.quota_exhausted or (exc is not None and _is_daily_quota_error(exc)):
            raise QuotaExhausted("Gemini free-tier DAILY quota is used up (it resets at midnight "
                                 "Pacific time). Re-run later.")

    def _build_task(self, question: str, file_name: str, file_path: str | None) -> str:
        parts = [question.strip()]
        if file_name and file_path:
            parts.append(f"\nAn attachment for this question is saved locally at: {file_path}\n"
                         f"Read it with read_file('{file_path}') (or pandas for spreadsheets).")
        elif file_name:
            parts.append(f"\n(The attachment '{file_name}' could not be downloaded; "
                         "answer from other sources if possible.)")
        task = "\n".join(parts)
        if not self.rules_in_system:
            task = ANSWER_RULES + "\nQUESTION:\n" + task
        return task

    def __call__(self, question: str, file_name: str = "", file_path: str | None = None) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        task = self._build_task(question, file_name, file_path)
        try:
            raw = self.agent.run(task, reset=True)
        except Exception as e:
            self._abort_if_model_dead(e)
            raise
        self._abort_if_model_dead()
        answer = clean_answer(raw)
        if not answer or answer.lower().startswith("error in generating"):
            raise RuntimeError(f"Agent produced no usable answer: {answer[:100]!r}")
        print(f"Agent returning answer: {answer}")
        return answer


def run_and_submit_all(use_cache: bool, profile: gr.OAuthProfile | None):
    """
    Fetches all questions, runs the BasicAgent on them, submits all answers,
    and displays the results.
    """
    # --- Determine HF Space Runtime URL and Repo URL ---
    space_id = os.getenv("SPACE_ID") # Get the SPACE_ID for sending link to the code

    if profile:
        username= f"{profile.username}"
        print(f"User logged in: {username}")
    else:
        print("User not logged in.")
        return "Please Login to Hugging Face with the button.", None

    api_url = DEFAULT_API_URL
    questions_url = f"{api_url}/questions"
    submit_url = f"{api_url}/submit"

    # 1. Instantiate Agent ( modify this part to create your agent)
    try:
        agent = BasicAgent()
    except Exception as e:
        print(f"Error instantiating agent: {e}")
        return f"Error initializing agent: {e}", None
    # In the case of an app running as a hugging Face space, this link points toward your codebase ( usefull for others so please keep it public)
    agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main" if space_id else "local-run"
    print(agent_code)

    # 2. Fetch Questions
    print(f"Fetching questions from: {questions_url}")
    try:
        response = requests.get(questions_url, timeout=15)
        response.raise_for_status()
        questions_data = response.json()
        if not questions_data:
             print("Fetched questions list is empty.")
             return "Fetched questions list is empty or invalid format.", None
        print(f"Fetched {len(questions_data)} questions.")
    except requests.exceptions.RequestException as e:
        print(f"Error fetching questions: {e}")
        return f"Error fetching questions: {e}", None
    except requests.exceptions.JSONDecodeError as e:
         print(f"Error decoding JSON response from questions endpoint: {e}")
         print(f"Response text: {response.text[:500]}")
         return f"Error decoding server response for questions: {e}", None
    except Exception as e:
        print(f"An unexpected error occurred fetching questions: {e}")
        return f"An unexpected error occurred fetching questions: {e}", None

    # 3. Run your Agent
    cache = load_cache() if use_cache else {}
    results_log = []
    answers_payload = []
    abort_reason = None
    print(f"Running agent on {len(questions_data)} questions...")
    for idx, item in enumerate(questions_data, 1):
        task_id = item.get("task_id")
        question_text = item.get("question")
        file_name = item.get("file_name") or ""
        if not task_id or question_text is None:
            print(f"Skipping item with missing task_id or question: {item}")
            continue
        if DRY_RUN_LIMIT and idx > DRY_RUN_LIMIT:
            break

        print(f"\n=== Question {idx}/{len(questions_data)} | task {task_id} ===")
        if task_id in cache:
            submitted_answer, note = cache[task_id], "cached"
        elif SKIP_MEDIA and needs_media(question_text, file_name):
            submitted_answer, note = "unknown", "skipped (audio/image/video)"
        else:
            file_path = download_attachment(api_url, task_id, file_name) if file_name else None
            t0 = time.time()
            try:
                submitted_answer = agent(question_text, file_name, file_path)
                cache[task_id] = submitted_answer
                save_cache(cache)
                note = f"{time.time() - t0:.0f}s"
            except RunAborted as e:
                print(f"Run aborted: {e}")
                abort_reason = str(e)
                break
            except Exception as e:
                print(f"Error running agent on task {task_id}: {e}")
                submitted_answer, note = "unknown", f"AGENT ERROR: {str(e)[:100]}"

        answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer or "unknown"})
        results_log.append({"Task ID": task_id, "Question": question_text,
                            "Submitted Answer": submitted_answer, "Note": note})

    if abort_reason:
        return (f"STOPPED after {len(answers_payload)} of {len(questions_data)} questions: {abort_reason}\n"
                "Finished answers are cached and reused on the next run. Nothing was submitted."), \
               pd.DataFrame(results_log)

    if DRY_RUN_LIMIT:
        return (f"DRY RUN: ran the first {len(answers_payload)} questions, nothing submitted. "
                "Check your Gemini usage in Google AI Studio, then unset DRY_RUN_LIMIT."), \
               pd.DataFrame(results_log)

    if not answers_payload:
        print("Agent did not produce any answers to submit.")
        return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)

    # 4. Prepare Submission 
    submission_data = {"username": username.strip(), "agent_code": agent_code, "answers": answers_payload}
    status_update = f"Agent finished. Submitting {len(answers_payload)} answers for user '{username}'..."
    print(status_update)

    # 5. Submit
    print(f"Submitting {len(answers_payload)} answers to: {submit_url}")
    try:
        response = requests.post(submit_url, json=submission_data, timeout=60)
        response.raise_for_status()
        result_data = response.json()
        final_status = (
            f"Submission Successful!\n"
            f"User: {result_data.get('username')}\n"
            f"Overall Score: {result_data.get('score', 'N/A')}% "
            f"({result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')} correct)\n"
            f"Message: {result_data.get('message', 'No message received.')}"
        )
        print("Submission successful.")
        results_df = pd.DataFrame(results_log)
        return final_status, results_df
    except requests.exceptions.HTTPError as e:
        error_detail = f"Server responded with status {e.response.status_code}."
        try:
            error_json = e.response.json()
            error_detail += f" Detail: {error_json.get('detail', e.response.text)}"
        except requests.exceptions.JSONDecodeError:
            error_detail += f" Response: {e.response.text[:500]}"
        status_message = f"Submission Failed: {error_detail}"
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
    except requests.exceptions.Timeout:
        status_message = "Submission Failed: The request timed out."
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
    except requests.exceptions.RequestException as e:
        status_message = f"Submission Failed: Network error - {e}"
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df
    except Exception as e:
        status_message = f"An unexpected error occurred during submission: {e}"
        print(status_message)
        results_df = pd.DataFrame(results_log)
        return status_message, results_df


# --- Build Gradio Interface using Blocks ---
with gr.Blocks() as demo:
    gr.Markdown("# GAIA Agent Evaluation Runner (smolagents CodeAgent + Gemini free tier)")
    gr.Markdown(
        """
        **Instructions:**

        1.  Make sure `GEMINI_API_KEY` is set as a Space **secret** (Settings, then Variables and secrets).
        2.  Log in to your Hugging Face account using the button below. This uses your HF username for submission.
        3.  Click 'Run Evaluation & Submit All Answers' to fetch questions, run your agent, submit answers, and see the score.

        ---
        **Disclaimers:**
        This run is deliberately slow: model calls are spaced out to stay inside the Gemini free-tier rate limit.
        Answers are cached to disk as they finish, so if the daily quota runs out mid-run, re-running skips the finished questions
        (the cache is lost if the Space restarts). Your Space must be public so the course can read your code.
        """
    )

    gr.LoginButton()

    use_cache = gr.Checkbox(value=True, label="Reuse cached answers from earlier runs (saves free quota)")
    run_button = gr.Button("Run Evaluation & Submit All Answers")

    status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
    # Removed max_rows=10 from DataFrame constructor
    results_table = gr.DataFrame(label="Questions and Agent Answers", wrap=True)

    run_button.click(
        fn=run_and_submit_all,
        inputs=[use_cache],
        outputs=[status_output, results_table]
    )

if __name__ == "__main__":
    print("\n" + "-"*30 + " App Starting " + "-"*30)
    # Check for SPACE_HOST and SPACE_ID at startup for information
    space_host_startup = os.getenv("SPACE_HOST")
    space_id_startup = os.getenv("SPACE_ID") # Get SPACE_ID at startup

    if space_host_startup:
        print(f"✅ SPACE_HOST found: {space_host_startup}")
        print(f"   Runtime URL should be: https://{space_host_startup}.hf.space")
    else:
        print("ℹ️  SPACE_HOST environment variable not found (running locally?).")

    if space_id_startup: # Print repo URLs if SPACE_ID is found
        print(f"✅ SPACE_ID found: {space_id_startup}")
        print(f"   Repo URL: https://huggingface.co/spaces/{space_id_startup}")
        print(f"   Repo Tree URL: https://huggingface.co/spaces/{space_id_startup}/tree/main")
    else:
        print("ℹ️  SPACE_ID environment variable not found (running locally?). Repo URL cannot be determined.")

    print("-"*(60 + len(" App Starting ")) + "\n")

    print("Launching Gradio Interface for Basic Agent Evaluation...")
    demo.launch(debug=True, share=False)