store.py

"""Shared job store for the papers-reproducibility demo.

Both the Gradio Space (app.py) and the local agent runner (local_runner.py)
read and write the same Hugging Face Dataset repo, since they run on
different machines. The repo holds:

    state.json              a JSON list of job records (the "database")
    reports/<job_id>.html   self-contained HTML report for a finished job
    traces/<job_id>.jsonl   append-only trace events for a job

This is a demo-grade store: every write downloads state.json, mutates it in
Python, and re-uploads the whole file. There is no locking, so concurrent
writers can race — acceptable for one admin and light request traffic, but
worth knowing if this grows into something bigger.
"""

from __future__ import annotations

import json
import os
import uuid
from datetime import datetime, timezone

from huggingface_hub import HfApi
from huggingface_hub.utils import EntryNotFoundError, HfHubHTTPError

try:
    from dotenv import load_dotenv

    load_dotenv()
except ImportError:
    pass

STATE_FILENAME = "state.json"
STATUSES = ("pending", "approved", "rejected", "running", "completed", "failed")


def _repo_id() -> str:
    repo_id = os.environ.get("HF_STORE_REPO")
    if not repo_id:
        raise RuntimeError(
            "HF_STORE_REPO is not set. Point it at the shared dataset repo, "
            "e.g. 'your-username/papers-repro-store'."
        )
    return repo_id


def _token() -> str | None:
    return os.environ.get("HF_TOKEN")


def _api() -> HfApi:
    return HfApi(token=_token())


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _download_state() -> list[dict]:
    from huggingface_hub import hf_hub_download

    try:
        path = hf_hub_download(
            repo_id=_repo_id(),
            repo_type="dataset",
            filename=STATE_FILENAME,
            token=_token(),
            force_download=True,
        )
    except (EntryNotFoundError, HfHubHTTPError):
        return []
    with open(path, "r") as f:
        return json.load(f)


def _upload_state(jobs: list[dict]) -> None:
    _api().upload_file(
        path_or_fileobj=json.dumps(jobs, indent=2).encode("utf-8"),
        path_in_repo=STATE_FILENAME,
        repo_id=_repo_id(),
        repo_type="dataset",
        commit_message="Update job state",
    )


def list_jobs() -> list[dict]:
    """All jobs, newest request first."""
    jobs = _download_state()
    return sorted(jobs, key=lambda j: j.get("requested_at", ""), reverse=True)


def get_job(job_id: str) -> dict | None:
    for job in _download_state():
        if job["id"] == job_id:
            return job
    return None


def existing_arxiv_ids() -> set[str]:
    """arxiv ids already tracked (any status) — lets the daily scan skip papers it has seen before."""
    return {job["arxiv_id"] for job in _download_state() if job.get("arxiv_id")}


def _base_job(
    *,
    title: str,
    paper_url: str,
    code_url: str,
    data_url: str,
    mode: str,
    notes: str,
    requested_by: str,
    source: str,
) -> dict:
    return {
        "id": str(uuid.uuid4()),
        "title": title.strip(),
        "paper_url": paper_url.strip(),
        "code_url": code_url.strip(),
        "data_url": data_url.strip(),
        "mode": mode if mode in ("author", "replicator") else "replicator",
        "notes": notes.strip(),
        "requested_by": requested_by.strip(),
        "source": source,
        "status": "pending",
        "requested_at": _now(),
        "approved_at": None,
        "started_at": None,
        "finished_at": None,
        "error": None,
        "report_path": None,
        "trace_path": None,
        # daily-scan-only fields (None for manual requests)
        "arxiv_id": None,
        "upvotes": None,
        "repro_score": None,
        "repro_summary": None,
        "decided_by": None,
    }


def create_request(
    title: str,
    paper_url: str,
    code_url: str,
    data_url: str = "",
    mode: str = "replicator",
    notes: str = "",
    requested_by: str = "",
) -> dict:
    job = _base_job(
        title=title,
        paper_url=paper_url,
        code_url=code_url,
        data_url=data_url,
        mode=mode,
        notes=notes,
        requested_by=requested_by,
        source="manual",
    )
    jobs = _download_state()
    jobs.append(job)
    _upload_state(jobs)
    return job


def create_candidate(
    *,
    title: str,
    paper_url: str,
    code_url: str,
    data_url: str = "",
    arxiv_id: str = "",
    upvotes: int = 0,
    repro_score: float = 0.0,
    repro_summary: str = "",
    mode: str = "replicator",
    notes: str = "",
) -> dict:
    """Create a job sourced from the daily-papers scan, awaiting Slack accept/reject.

    Lands as an ordinary "pending" job so the existing Admin tab can approve
    or reject it too, in case Slack is unreachable or the message is missed.
    """
    job = _base_job(
        title=title,
        paper_url=paper_url,
        code_url=code_url,
        data_url=data_url,
        mode=mode,
        notes=notes,
        requested_by="daily-scan",
        source="daily_scan",
    )
    job["arxiv_id"] = arxiv_id
    job["upvotes"] = upvotes
    job["repro_score"] = repro_score
    job["repro_summary"] = repro_summary.strip()
    jobs = _download_state()
    jobs.append(job)
    _upload_state(jobs)
    return job


def set_status(job_id: str, status: str, **fields) -> dict:
    if status not in STATUSES:
        raise ValueError(f"Unknown status {status!r}")
    jobs = _download_state()
    updated = None
    for job in jobs:
        if job["id"] == job_id:
            job["status"] = status
            job.update(fields)
            updated = job
            break
    if updated is None:
        raise KeyError(f"No job with id {job_id!r}")
    _upload_state(jobs)
    return updated


def append_trace_events(job_id: str, events: list[dict]) -> str:
    """Append events to traces/<job_id>.jsonl and return the repo path."""
    if not events:
        return f"traces/{job_id}.jsonl"

    from huggingface_hub import hf_hub_download

    path_in_repo = f"traces/{job_id}.jsonl"
    try:
        local_path = hf_hub_download(
            repo_id=_repo_id(),
            repo_type="dataset",
            filename=path_in_repo,
            token=_token(),
            force_download=True,
        )
        with open(local_path, "r") as f:
            existing = f.read()
    except (EntryNotFoundError, HfHubHTTPError):
        existing = ""

    new_lines = "\n".join(json.dumps(e) for e in events)
    content = existing + (new_lines + "\n" if not existing or existing.endswith("\n") else "\n" + new_lines + "\n")

    _api().upload_file(
        path_or_fileobj=content.encode("utf-8"),
        path_in_repo=path_in_repo,
        repo_id=_repo_id(),
        repo_type="dataset",
        commit_message=f"Append {len(events)} trace event(s) for {job_id}",
    )
    return path_in_repo


def read_trace(job_id: str) -> list[dict]:
    from huggingface_hub import hf_hub_download

    try:
        local_path = hf_hub_download(
            repo_id=_repo_id(),
            repo_type="dataset",
            filename=f"traces/{job_id}.jsonl",
            token=_token(),
            force_download=True,
        )
    except (EntryNotFoundError, HfHubHTTPError):
        return []
    events = []
    with open(local_path, "r") as f:
        for line in f:
            line = line.strip()
            if line:
                events.append(json.loads(line))
    return events


def upload_report(job_id: str, html_path: str) -> str:
    path_in_repo = f"reports/{job_id}.html"
    _api().upload_file(
        path_or_fileobj=html_path,
        path_in_repo=path_in_repo,
        repo_id=_repo_id(),
        repo_type="dataset",
        commit_message=f"Upload report for {job_id}",
    )
    return path_in_repo


def read_report(job_id: str) -> str | None:
    from huggingface_hub import hf_hub_download

    try:
        local_path = hf_hub_download(
            repo_id=_repo_id(),
            repo_type="dataset",
            filename=f"reports/{job_id}.html",
            token=_token(),
            force_download=True,
        )
    except (EntryNotFoundError, HfHubHTTPError):
        return None
    with open(local_path, "r") as f:
        return f.read()


def ensure_repo_exists() -> None:
    """Create the dataset repo (private) if it doesn't exist yet."""
    _api().create_repo(repo_id=_repo_id(), repo_type="dataset", private=True, exist_ok=True)