scripts/cache_examples.py

"""Pre-generate example cache files for judge-proofing.

Run this script once on a machine with GPU access (HF Space or local with CUDA)
AFTER the Space has warmed up so the model is resident.  The output JSON files
are written to data/example_cache/<key>.json and are committed to the repo so
that example-button clicks are instant even on a cold ZeroGPU worker.

Usage (from the project root):
    python scripts/cache_examples.py

Each file contains a serialised HandoffOutput (Pydantic .model_dump()).  The app
loads these at click time via app._load_cached_example() and renders them through
the normal render_handoff_html / render_safety_html path — validated, not raw.

NOTE: Do NOT run this script inside a ZeroGPU Space (SPACES_ZERO_GPU is set).
ZeroGPU only allocates GPU within a live request context (@spaces.GPU window),
so model.generate() calls made at the top level will run on CPU and be very slow
or time out.  Run on a regular GPU box (A10/A100 Space, local CUDA machine, or
a standard HF Space with a persistent GPU).
"""
from __future__ import annotations

import json
import os
import sys
from pathlib import Path

# Make sure the project root is on the path when run as a script.
sys.path.insert(0, str(Path(__file__).parent.parent))

# Parse this before importing app.py. Otherwise --no-model still triggers the
# import-time 4B model load, defeating the purpose of the flag.
USE_MODEL = "--no-model" not in sys.argv
if not USE_MODEL:
    os.environ["DENTAL_SOAP_USE_MODEL"] = "0"

from app import build_outputs, EXAMPLES  # noqa: E402  (after sys.path fix)
from examples import STEP2_CHECKS, STEP3_CHECKS, STEP4_CHECKS  # noqa: E402
from schema import HandoffOutput  # noqa: E402

CACHE_DIR = Path(__file__).parent.parent / "data" / "example_cache"
CACHE_DIR.mkdir(parents=True, exist_ok=True)


def _split_checks(checks: list[str]):
    return (
        [c for c in checks if c in STEP2_CHECKS],
        [c for c in checks if c in STEP3_CHECKS],
        [c for c in checks if c in STEP4_CHECKS],
    )


def cache_example(key: str, use_model: bool = True) -> None:
    ex = EXAMPLES[key]
    checks_dental, checks_jaw, checks_body = _split_checks(ex["checks"])

    print(f"Generating cache for '{key}' (use_model={use_model}) ...")
    _safety, _handoff, raw_json, status, _pdf, _email, _workflow = build_outputs(
        ex["story"],
        ex.get("name", ""),
        ex["age"],
        ex["language"],
        ex["chief_concern"],
        ex["duration"],
        ex["tooth_or_area"],
        ex["recent_dental_work"],
        ex["pain_score"],
        ex["meds"],
        ex["allergies"],
        ex["goals"],
        checks_dental,
        checks_jaw,
        checks_body,
        use_model,
    )
    print(f"  status: {status}")

    if use_model and not status.startswith("AI model path used (ZeroGPU):"):
        raise RuntimeError(
            f"Refusing to cache '{key}': model-backed generation was requested "
            f"but the app returned: {status}"
        )

    output = HandoffOutput.model_validate_json(raw_json)
    out_path = CACHE_DIR / f"{key}.json"
    tmp_path = out_path.with_suffix(".json.tmp")
    tmp_path.write_text(
        json.dumps(output.model_dump(), ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    tmp_path.replace(out_path)
    print(f"  written: {out_path}")


if __name__ == "__main__":
    # Guard: ZeroGPU only allocates GPU inside a request context, so running the
    # model generate() calls at script top-level will silently fall back to CPU
    # and produce unusably slow or broken output.  Bail out early with a clear
    # message so the operator knows to re-run on a non-ZeroGPU GPU machine.
    if os.getenv("SPACE_ID") and os.getenv("SPACES_ZERO_GPU"):
        sys.exit(
            "[cache_examples] ERROR: this script is running inside a ZeroGPU Space "
            "(SPACE_ID and SPACES_ZERO_GPU are both set).  ZeroGPU only allocates GPU "
            "within a live @spaces.GPU request window — model calls here run on CPU "
            "and will be extremely slow or time out.  Re-run on a regular GPU box "
            "(a standard HF Space with a persistent GPU, a local CUDA machine, or an "
            "A10/A100 HF Space)."
        )

    # Pass --no-model to generate from the deterministic template path only.
    for example_key in EXAMPLES:
        cache_example(example_key, use_model=USE_MODEL)
    print("Done. Commit data/example_cache/*.json to the repo.")