Blame for app/curate_api.py

"""Curate a corpus: sample previews of the clean, calibrate, create, compose workflow, plus the recipe to run it in full."""
from __future__ import annotations

import json
import os
import threading
import time
import uuid

from fastapi import APIRouter, HTTPException
from fastapi.responses import Response
from pydantic import BaseModel

from datainventor.clean import BENCHMARKS, DEFAULTS as CLEAN_RULES
from datainventor.create import DOC_TYPES, STYLES, chunks, rephrase, styles_for
from datainventor.llm import EndpointError
from datainventor.preview import Preview, run_preview
from datainventor.recipe import DEFAULT, dump, validate
from datainventor.sources import KINDS, inspect

router = APIRouter(prefix="/api/curate")
PER_SOURCE = int(os.environ.get("CURATE_PER_SOURCE", "200"))
GRADE_N = int(os.environ.get("CURATE_GRADE", "48"))
REPHRASE_N = int(os.environ.get("CURATE_REPHRASE", "8"))
MAX_SOURCES, MAX_TARGETS = 4, 2
INSTALL = "pip install 'datainventor[full] @ https://huggingface.co/spaces/codelion/data-inventor/resolve/main/dist/datainventor-0.2.2-py3-none-any.whl'"
PREVIEWS: dict[str, Preview] = {}
TRIES: dict[str, list] = {}
GUARD = threading.Lock()


def _ep():
    from main import make_endpoint  # the host's endpoint; visitors never supply one
    return make_endpoint()


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


@router.get("/defaults")
def defaults():
    from main import MAX_ROWS
    return {"recipe": DEFAULT, "benchmarks": list(BENCHMARKS), "kinds": list(KINDS), "clean_rules": CLEAN_RULES,
            "styles": {k: v[0] for k, v in STYLES.items()}, "styles_by_kind": {k: styles_for(k) for k in KINDS},
            "doc_types": list(DOC_TYPES), "per_source": PER_SOURCE, "max_sources": MAX_SOURCES, "max_targets": MAX_TARGETS,
            "max_docs": MAX_ROWS, "install": INSTALL}


@router.get("/inspect")
def do_inspect(dataset: str):
    ds = dataset.strip().removeprefix("https://huggingface.co/datasets/").strip("/")
    if not ds or ds.count("/") != 1 or len(ds) > 120:
        raise HTTPException(400, "Use a Hugging Face dataset id like HuggingFaceFW/fineweb-edu.")
    try:
        return inspect(ds, _token())
    except ValueError as e:
        raise HTTPException(404, str(e))
    except Exception:
        raise HTTPException(404, "The dataset viewer could not open that dataset. It must be public and have the viewer enabled.")


class RecipeIn(BaseModel):
    recipe: dict


def _checked(rec: dict) -> dict:
    for s in rec.get("sources", []) + rec.get("targets", {}).get("datasets", []):
        if s.get("files"):
            raise HTTPException(400, "Local files only work in the runner on your own machines. Use Hugging Face datasets here.")
    if len(rec.get("sources", [])) > MAX_SOURCES or len(rec.get("targets", {}).get("datasets", [])) > MAX_TARGETS:
        raise HTTPException(400, f"The hosted preview takes up to {MAX_SOURCES} sources and {MAX_TARGETS} target datasets.")
    try:
        return validate(rec)
    except ValueError as e:
        raise HTTPException(400, str(e))


@router.post("/recipe")
def recipe_file(body: RecipeIn):
    """The recipe as YAML, for the runner."""
    r = _checked(body.recipe)
    return Response(dump(r), media_type="application/x-yaml", headers={"Content-Disposition": 'attachment; filename="recipe.yaml"'})


@router.post("/previews")
def start(body: RecipeIn):
    r = _checked(body.recipe)
    ep = _ep()
    with GUARD:
        running = sum(1 for p in PREVIEWS.values() if p.state in ("queued", "running"))
        if running >= 2:
            raise HTTPException(429, "Two previews are already running. Wait for one to finish.")
        pid = uuid.uuid4().hex[:10]
        p = Preview(r)
        PREVIEWS[pid] = p
        for old in sorted(PREVIEWS, key=lambda k: PREVIEWS[k].started)[:-20]:
            PREVIEWS.pop(old, None)
    from main import MAX_ROWS
    threading.Thread(target=run_preview, args=(p, ep, _token()),
                     kwargs={"per_source": PER_SOURCE, "grade_n": GRADE_N, "rephrase_n": REPHRASE_N, "invent_n": 4, "max_docs": MAX_ROWS},
                     daemon=True).start()
    return {"id": pid}


def _get(pid: str) -> Preview:
    if pid not in PREVIEWS:
        raise HTTPException(404, "That preview is no longer here. Run it again.")
    return PREVIEWS[pid]


@router.get("/previews/{pid}")
def status(pid: str):
    p = _get(pid)
    out = {"id": pid, "state": p.state, "stage": p.stage, "message": p.message, "elapsed": round((p.finished or time.time()) - p.started, 1)}
    if p.state in ("done", "stopped"):
        out["report"] = p.report
    return out


@router.post("/previews/{pid}/stop")
def stop(pid: str):
    _get(pid).stop.set()
    return {"ok": True}


@router.get("/previews/{pid}/download")
def download(pid: str):
    p = _get(pid)
    if not p.docs:
        raise HTTPException(404, "There are no documents to download yet.")
    body = "\n".join(json.dumps(d, ensure_ascii=False) for d in p.docs) + "\n"
    return Response(body, media_type="application/jsonl", headers={"Content-Disposition": f'attachment; filename="curated-sample-{pid}.jsonl"'})


class TryIn(BaseModel):
    text: str
    kind: str = "web"
    styles: list[str] = []


@router.post("/try")
def try_styles(body: TryIn):
    """Rewrite one document in up to four styles, side by side."""
    text = body.text.strip()
    if len(text.split()) < 60:
        raise HTTPException(400, "Paste or pick a document of at least 60 words.")
    styles = [s for s in (body.styles or styles_for(body.kind)) if s in STYLES][:4]
    if not styles:
        raise HTTPException(400, "Choose at least one style.")
    now = time.time()
    with GUARD:
        recent = [t for t in TRIES.get("all", []) if now - t < 60]
        if len(recent) >= 6:
            raise HTTPException(429, "Too many rewrites in the last minute. Try again shortly.")
        TRIES["all"] = recent + [now]
    ep = _ep()
    piece = chunks(text)[0]
    from concurrent.futures import ThreadPoolExecutor
    try:
        with ThreadPoolExecutor(len(styles)) as ex:
            outs = list(ex.map(lambda s: rephrase(ep, piece, s), styles))
    except EndpointError as e:
        raise HTTPException(502, str(e))
    return {"source": piece, "results": [{"style": o["style"], "name": STYLES[o["style"]][0], "text": o["text"], "ok": o["ok"],
                                          "why": o.get("why", ""), "fidelity": o.get("fidelity")} for o in outs]}