phase1/scripts/start_private_space.py

"""Fetch private serving artifacts, then start the public FastAPI Space."""

from __future__ import annotations

import os
from pathlib import Path

from huggingface_hub import HfApi, snapshot_download


def _product_env(suffix: str, default: str = "") -> str:
    """Prefer Moonley configuration while accepting one-release legacy keys."""
    return os.environ.get(f"MOONLEY_{suffix}", os.environ.get(f"THEMIS_{suffix}", default))


def _require_private_dataset(repo_id: str, token: str, label: str) -> None:
    """Refuse to boot if a serving-artifact repository is publicly visible."""
    info = HfApi(token=token).repo_info(repo_id=repo_id, repo_type="dataset")
    if not bool(getattr(info, "private", False)):
        raise SystemExit(f"{label} repository must be private: {repo_id}")


def _download(repo_id: str, revision: str, target: Path, token: str, label: str) -> None:
    target.mkdir(parents=True, exist_ok=True)
    print(f"[private-release] fetching {label} at pinned revision {revision}", flush=True)
    snapshot_download(
        repo_id=repo_id,
        repo_type="dataset",
        revision=revision,
        local_dir=str(target),
        token=token,
    )


def main() -> None:
    token = os.environ.get("HF_TOKEN", "").strip()
    if not token:
        raise SystemExit("HF_TOKEN is required to load private Moonley artifacts")

    release_repo = _product_env("RELEASE_REPO", "vg15o2/themis-indian-kanoon-qwen-v1").strip()
    release_revision = _product_env(
        "RELEASE_REVISION", "12f58201987cc8ec7697010754ab75765c5f5a24"
    ).strip()
    data_dir = Path(_product_env("DATA", "/tmp/moonley_release")).resolve()
    _require_private_dataset(release_repo, token, "legal corpus")
    _download(release_repo, release_revision, data_dir, token, "legal corpus")
    if not (data_dir / "release_manifest.json").is_file():
        raise SystemExit("private serving release is missing release_manifest.json")

    statute_repo = _product_env("STATUTE_REPO", "vg15o2/themis-statutes-v1").strip()
    statute_revision = _product_env(
        "STATUTE_REVISION", "ebf66528e417358a09903d95f6718ccfeb94a426"
    ).strip()
    statute_dir = Path(_product_env("STATUTE_CHROMA", "/tmp/moonley_statutes")).resolve()
    _require_private_dataset(statute_repo, token, "statute embeddings")
    _download(statute_repo, statute_revision, statute_dir, token, "exact statute library")
    if not (statute_dir / "chroma.sqlite3").is_file():
        raise SystemExit("private statute release is missing chroma.sqlite3")
    if not any(statute_dir.rglob("data_level0.bin")):
        raise SystemExit("private statute release is missing its vector segment")

    template_repo = _product_env(
        "DRAFTING_TEMPLATE_REPO", "vg15o2/themis-drafting-templates-v1"
    ).strip()
    template_revision = _product_env(
        "DRAFTING_TEMPLATE_REVISION",
        "2b036d4bef7ebe7a3b3bb8d0a094d7fefdd8c4d2",
    ).strip()
    drafting_dir = Path(_product_env("DRAFTING_DIR", "/app/phase1/drafting")).resolve()
    _require_private_dataset(template_repo, token, "drafting templates")
    _download(
        template_repo,
        template_revision,
        drafting_dir,
        token,
        "drafting templates",
    )
    expected_templates = {
        "article_32_petition.pdf",
        "civil_appeal.pdf",
        "curative_petition.pdf",
        "slp_civil_full.pdf",
        "slp_criminal_full.pdf",
        "slp_outline.pdf",
    }
    missing_templates = sorted(
        name for name in expected_templates if not (drafting_dir / "templates" / name).is_file()
    )
    if missing_templates:
        raise SystemExit(
            "private drafting release is missing: " + ", ".join(missing_templates)
        )

    # The API process does not need Hub credentials after the snapshots are local.
    os.environ.pop("HF_TOKEN", None)
    print("[private-release] downloads complete; starting Moonley API", flush=True)
    os.execvp(
        "uvicorn",
        [
            "uvicorn",
            "serve_agent:app",
            "--app-dir",
            "phase1/scripts",
            "--host",
            "0.0.0.0",
            "--port",
            "7860",
        ],
    )


if __name__ == "__main__":
    main()