app.py

"""GAIA course submission UI: safe one-question OpenAI tests and explicit submit."""
from __future__ import annotations

import json
import os
import re
import threading
from pathlib import Path

import gradio as gr
import pandas as pd
import requests
import spaces

from agent import (AttachmentUnavailable, UnsupportedMedia, UnresolvedAnswer, BudgetError, BudgetLedger, GaiaAgent, MODEL, SCORING_API,
                   STORAGE, MAX_ESTIMATED_TOTAL_USD, MAX_ESTIMATED_QUESTION_USD, MAX_MODEL_CALLS_PER_QUESTION)

ANSWERS_DIR = STORAGE / "answers"
ANSWERS_DIR.mkdir(parents=True, exist_ok=True)
# Gradio event queue has concurrency 1; the lock guards accidental parallel requests.
RUN_LOCK = threading.Lock()
FILES_LOCK = threading.RLock()
COLS = ["#", "Question", "Cached answer", "Status"]


def _skips_path(username: str) -> Path:
    return ANSWERS_DIR / (username + "_skips.json")


def _load_skips(username: str) -> dict:
    with FILES_LOCK:
        try:
            data = json.loads(_skips_path(username).read_text())
            return data if isinstance(data, dict) else {}
        except (FileNotFoundError, ValueError, OSError):
            return {}


def _save_skips(username: str, data: dict) -> None:
    with FILES_LOCK:
        dest = _skips_path(username)
        tmp = dest.with_suffix(".tmp")
        tmp.write_text(json.dumps(data, ensure_ascii=False, indent=2))
        tmp.replace(dest)



def _user(profile) -> str:
    username = getattr(profile, "username", "") if profile else ""
    if not username:
        raise RuntimeError("Sign in with Hugging Face first")
    return re.sub(r"[^A-Za-z0-9_-]", "", username)


def _path(username: str) -> Path:
    return ANSWERS_DIR / (username + ".json")


def _load(username: str) -> dict:
    with FILES_LOCK:
        try:
            raw = json.loads(_path(username).read_text())
            return raw if isinstance(raw, dict) else {}
        except (FileNotFoundError, ValueError, OSError):
            return {}


def _save(username: str, data: dict) -> None:
    with FILES_LOCK:
        dest = _path(username)
        tmp = dest.with_suffix(".tmp")
        tmp.write_text(json.dumps(data, ensure_ascii=False, indent=2))
        tmp.replace(dest)


def _questions() -> list:
    r = requests.get(f"{SCORING_API}/questions", timeout=25)
    r.raise_for_status()
    data = r.json()
    if not isinstance(data, list) or not data or any(not q.get("task_id") or not q.get("question") for q in data):
        raise RuntimeError("Scoring API returned invalid question data")
    return data


def _table(questions: list, answers: dict, skips: dict | None = None) -> pd.DataFrame:
    skips = skips or {}
    return pd.DataFrame([{"#":i, "Question":q["question"],
                          "Cached answer":answers.get(q["task_id"], ""),
                          "Status": ("Cached" if q["task_id"] in answers
                                     else "Skipped: " + skips[q["task_id"]] if q["task_id"] in skips
                                     else "Not run")}
                         for i,q in enumerate(questions,1)], columns=COLS)


def _usage() -> str:
    return (f"Estimated OpenAI usage by this Space: ${BudgetLedger.total():.4f} / "
            f"${MAX_ESTIMATED_TOTAL_USD:.2f} initial budget. "
            "Not a billing cap; verify actual usage in OpenAI Platform.")


def check_setup(profile: gr.OAuthProfile | None):
    """Free checks: does not create an OpenAI client or make a model request."""
    try:
        user = _user(profile)
        if not os.environ.get("OPENAI_API_KEY"):
            raise RuntimeError("OPENAI_API_KEY missing: add it in Settings > Secrets")
        questions = _questions()
        return (f"Setup OK: {len(questions)} GAIA questions found, OpenAI key is configured, "
                f"HF username: {user}. NO paid OpenAI call was made.\n{_usage()}",
                _table(questions, _load(user), _load_skips(user)))
    except Exception as e:
        return f"Setup check: {e} (no paid OpenAI call made)", pd.DataFrame(columns=COLS)


def solve_one(n, profile: gr.OAuthProfile | None):
    try:
        user = _user(profile)
        if not os.environ.get("OPENAI_API_KEY"):
            raise RuntimeError("OPENAI_API_KEY secret missing")
        questions = _questions()
        num = int(n)
        if not (1 <= num <= len(questions)):
            raise ValueError(f"Question number must be between 1 and {len(questions)}")
        answers = _load(user)
        q = questions[num-1]
        if q["task_id"] in answers or q["task_id"] in _load_skips(user):
            yield f"Question {num} is cached or marked skipped. Use Clear selected Q# to retry.\n{_usage()}", _table(questions,answers, _load_skips(user))
            return
        if not RUN_LOCK.acquire(blocking=False):
            yield "Another evaluation is running. No OpenAI request was sent.", _table(questions,answers, _load_skips(user))
            return
        try:
            yield f"Solving question {num} with {MODEL} (maximum {MAX_MODEL_CALLS_PER_QUESTION} model calls). Not submitting yet...", _table(questions,answers, _load_skips(user))
            try:
                answer, calls = GaiaAgent().run(q["question"], q["task_id"], q.get("file_name") or "")
            except (AttachmentUnavailable, UnsupportedMedia, UnresolvedAnswer) as exc:
                skips = _load_skips(user)
                skips[q["task_id"]] = str(exc)[:200]
                _save_skips(user, skips)
                yield f"Unresolved/Skipped Q{num}: {exc}\n{_usage()}", _table(questions,answers,skips)
                return
            answers[q["task_id"]] = answer
            _save(user,answers)
            yield (f"Question {num} saved using {calls} OpenAI calls. "
                   f"{len(answers)}/{len(questions)} cached.\n{_usage()}"), _table(questions,answers, _load_skips(user))
        finally:
            RUN_LOCK.release()
    except Exception as exc:
        yield f"Stopped: {type(exc).__name__}: {str(exc)[:450]}\n{_usage()}", pd.DataFrame(columns=COLS)


def run_batch(approved: bool, max_questions: int, profile: gr.OAuthProfile | None):
    """Run up to N billable questions, skip missing attachments without stalling.

    No automatic model retries; all successful answers are saved immediately.
    """
    if not approved:
        yield "Tick the OpenAI API consent box first. No paid requests made.", pd.DataFrame(columns=COLS)
        return
    try:
        user = _user(profile)
        if not os.environ.get("OPENAI_API_KEY"):
            raise RuntimeError("OPENAI_API_KEY secret missing")
        limit = int(max_questions)
        if limit not in (3, 5, 10, 20):
            raise ValueError("Choose 3, 5, 10 or 20")
        questions = _questions()
        answers, skips = _load(user), _load_skips(user)
        if not RUN_LOCK.acquire(blocking=False):
            yield "Another evaluation is running; no paid requests made.", _table(questions,answers,skips)
            return
        try:
            paid_questions = 0
            newly_skipped = 0
            for i,q in enumerate(questions,1):
                task_id = q["task_id"]
                if task_id in answers or task_id in skips:
                    continue
                if paid_questions >= limit:
                    break
                file_name = (q.get("file_name") or "").lower()
                if file_name.endswith((".m4a", ".mp3", ".wav", ".mp4", ".webm")):
                    skips[task_id] = "Audio/video attachment unsupported (no OpenAI charge)"
                    _save_skips(user, skips)
                    newly_skipped += 1
                    yield f"Skipped unsupported audio/video Q{i}.\n{_usage()}", _table(questions,answers,skips)
                    continue
                yield (f"Starting Q{i}. {paid_questions}/{limit} supported questions processed in this batch. "
                       f"No answers are submitted automatically.\n{_usage()}"), _table(questions,answers,skips)
                try:
                    answer,calls = GaiaAgent().run(q["question"], task_id, q.get("file_name") or "")
                except (AttachmentUnavailable, UnsupportedMedia, UnresolvedAnswer) as exc:
                    skips[task_id] = str(exc)[:200]
                    _save_skips(user, skips)
                    newly_skipped += 1
                    yield f"Unresolved/Skipped Q{i}: {exc}\n{_usage()}", _table(questions,answers,skips)
                    continue
                except BudgetError as exc:
                    yield f"Budget guard stopped the batch before Q{i}: {exc}\n{_usage()}", _table(questions,answers,skips)
                    return
                except Exception as exc:
                    yield (f"Stopped at Q{i}: {type(exc).__name__}: {str(exc)[:300]}. "
                           f"Already saved answers remain intact.\n{_usage()}"), _table(questions,answers,skips)
                    return
                answers[task_id] = answer
                _save(user,answers)
                paid_questions += 1
                yield (f"Saved Q{i} ({calls} model calls). {len(answers)}/{len(questions)} answers saved; "
                       f"{len(skips)} questions skipped.\n{_usage()}"), _table(questions,answers,skips)
            yield (f"Batch finished. Saved {len(answers)}/{len(questions)}; skipped {len(skips)} "
                   f"({newly_skipped} new skips in this batch). Download a backup now.\n{_usage()}"), _table(questions,answers,skips)
        finally:
            RUN_LOCK.release()
    except Exception as exc:
        yield f"Stopped: {type(exc).__name__}: {str(exc)[:400]}", pd.DataFrame(columns=COLS)


def view_cached(profile: gr.OAuthProfile | None):
    try:
        user = _user(profile)
        questions = _questions()
        answers = _load(user)
        return f"{len(answers)}/{len(questions)} saved.\n{_usage()}",_table(questions,answers, _load_skips(user))
    except Exception as exc:
        return f"Error: {exc}", pd.DataFrame(columns=COLS)


def clear_one(n, confirmed: bool, profile: gr.OAuthProfile | None):
    """Explicitly clear exactly one answer so it may be rerun if needed."""
    try:
        user = _user(profile)
        questions = _questions()
        if not confirmed:
            return "Check confirmation before clearing an answer. No paid requests made.", _table(questions,_load(user), _load_skips(user))
        num = int(n)
        if not (1 <= num <= len(questions)):
            raise ValueError("Question number out of range")
        answers = _load(user)
        task_id = questions[num-1]["task_id"]
        removed = answers.pop(task_id, None)
        skips = _load_skips(user)
        skipped = skips.pop(task_id, None)
        _save(user, answers)
        _save_skips(user, skips)
        return (f"Cleared Q{num} cached answer/skip. Now eligible to retry. No API charge."
                if removed is not None or skipped is not None
                else f"Q{num} was neither cached nor skipped. No changes."), _table(questions,answers,skips)
    except Exception as exc:
        return f"Clear failed: {exc}", pd.DataFrame(columns=COLS)


def backup(profile: gr.OAuthProfile | None):
    """Download after each few questions to survive a restart on nonpersistent Spaces."""
    try:
        user = _user(profile)
        data = _load(user)
        if not data:
            return None, "No cached answers to back up"
        source = _path(user)
        return str(source), f"Backup ready ({len(data)} answers). Save it to your computer."
    except Exception as exc:
        return None, f"Backup failed: {exc}"


def restore(uploaded_path, profile: gr.OAuthProfile | None):
    try:
        user = _user(profile)
        if not uploaded_path:
            raise ValueError("Upload a previously downloaded JSON backup")
        questions = _questions()
        allowed = {q["task_id"] for q in questions}
        path = Path(uploaded_path)
        if path.stat().st_size > 500000:
            raise ValueError("Backup file too large")
        incoming = json.loads(path.read_text())
        if not isinstance(incoming, dict):
            raise ValueError("Backup must be a JSON object")
        cleaned = {k:v for k,v in incoming.items() if k in allowed and isinstance(v,str) and len(v) < 2000}
        answers = _load(user)
        for k,v in cleaned.items():
            answers.setdefault(k,v) # Never overwrite an existing answer silently
        _save(user,answers)
        return f"Restored {len(cleaned)} candidate answers; {len(answers)}/{len(questions)} cached. No API charges.", _table(questions,answers, _load_skips(user))
    except Exception as exc:
        return f"Restore failed: {exc}", pd.DataFrame(columns=COLS)


def submit(allow_blanks: bool, profile: gr.OAuthProfile | None):
    """Explicit action; submission itself uses the free course scoring API."""
    try:
        user = _user(profile)
        questions = _questions()
        answers = _load(user)
        missing = [i for i,q in enumerate(questions,1) if q["task_id"] not in answers]
        if missing and not allow_blanks:
            return (f"Not submitted: missing Q{missing}. To submit now with empty answers "
                    "for unsupported questions, tick the confirmation checkbox. No API spend."), _table(questions,answers, _load_skips(user))
        space_id = os.environ.get("SPACE_ID")
        if not space_id:
            raise RuntimeError("SPACE_ID not set. Must submit from your public Hugging Face Space")
        payload = {"username":user,
                   "agent_code":f"https://huggingface.co/spaces/{space_id}/tree/main",
                   "answers":[{"task_id":q["task_id"],"submitted_answer":answers.get(q["task_id"], "")} for q in questions]}
        r = requests.post(f"{SCORING_API}/submit", json=payload, timeout=60)
        r.raise_for_status()
        score = r.json()
        return (f"Official submission: score {score.get('score', 'n/a')}%, "
                f"{score.get('correct_count','?')}/{score.get('total_attempted','?')} correct "
                f"({len(missing)} unanswered submitted as blank). "
                f"{str(score.get('message',''))[:400]}"), _table(questions,answers, _load_skips(user))
    except Exception as exc:
        return f"Submission failed: {exc}", pd.DataFrame(columns=COLS)


# ZeroGPU hosting requires a registered @spaces.GPU Gradio handler. It is deliberately
# never clicked; OpenAI requests always run on CPU via remote API.
@spaces.GPU(duration=1)
def _unused_zero_gpu_probe():
    return "registered"


with gr.Blocks(title="GAIA Agent - budget-safe OpenAI") as demo:
    gr.Markdown("# 🧠 GAIA Agent — GPT-6 Sol precision edition")
    gr.Markdown(f"**Model:** `{MODEL}` · **Initial local spending guard:** ${MAX_ESTIMATED_TOTAL_USD:.2f} "
                f"· **Per-question estimate guard:** ${MAX_ESTIMATED_QUESTION_USD:.2f} · "
                f"**Maximum {MAX_MODEL_CALLS_PER_QUESTION} model requests/question**. No automatic retries.\n\n"
                "Estimated guards are **not** hard limits on OpenAI billing. Check your actual usage dashboard.")
    gr.LoginButton()
    invisible = gr.Button("ZeroGPU startup registration",visible=False)
    invisible_result = gr.Textbox(visible=False)
    invisible.click(_unused_zero_gpu_probe,outputs=[invisible_result])
    setup = gr.Button("1. Check setup (FREE: no OpenAI call)")
    with gr.Row():
        number = gr.Number(label="Question number", value=1, minimum=1, maximum=20, precision=0)
        test = gr.Button("2. Test ONE question", variant="primary")
    clear_confirm = gr.Checkbox(label="I confirm clearing the selected cached answer (free)", value=False)
    clear_button = gr.Button("Clear selected Q# answer to allow a retry (free)")
    consent = gr.Checkbox(label="I understand this makes paid OpenAI API requests", value=False)
    batch_size = gr.Dropdown([3, 5, 10, 20], value=3, label="Maximum supported questions to solve in one batch", interactive=True)
    batch = gr.Button("3. Run selected batch (OpenAI API)")
    with gr.Row():
        see = gr.Button("Show saved answers")
        send = gr.Button("Submit answers to GAIA scoring (free)")
    allow_blanks = gr.Checkbox(label="Allow official submission with blank answers for missing questions", value=False)
    status = gr.Textbox(label="Progress and cost estimate", lines=5, interactive=False)
    table = gr.Dataframe(label="Questions and cached answers", wrap=True, interactive=False)
    gr.Markdown("### Backup (recommended)")
    with gr.Row():
        export = gr.Button("Download answer backup")
        exported = gr.File(label="Download your saved JSON answers", interactive=False)
    backup_status = gr.Textbox(label="Backup status", interactive=False)
    with gr.Row():
        uploaded = gr.File(label="Upload an earlier answer backup (.json)", file_types=[".json"], type="filepath")
        restore_button = gr.Button("Restore answers (free)")
    gr.Markdown("**Before updating your Space or running a large batch:** click Download answer backup. "
                "Unavailable GAIA attachments (HTTP 404) are marked Skipped, not answered blindly. "
                "The batch stops at any model/API error or local budget limit. "
                "Your HF Space remains PUBLIC, but your API key must be stored only as a **Secret**. "
                "Audio transcription is disabled by default to avoid untracked charges. "
                "The answer cache and cost ledger persist only if /data persistent storage is mounted; otherwise export backups regularly.")
    setup.click(check_setup, outputs=[status,table])
    test.click(solve_one,inputs=[number],outputs=[status,table])
    clear_button.click(clear_one,inputs=[number,clear_confirm],outputs=[status,table])
    batch.click(run_batch,inputs=[consent,batch_size],outputs=[status,table])
    see.click(view_cached,outputs=[status,table])
    send.click(submit,inputs=[allow_blanks],outputs=[status,table])
    export.click(backup,outputs=[exported,backup_status])
    restore_button.click(restore,inputs=[uploaded],outputs=[status,table])

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch()