scripts/review_candidate.py

from __future__ import annotations

import argparse
import asyncio
import json
import os
import sys
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

import httpx

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from app.catalog import Catalog, HubModel  # noqa: E402
from app.huggingface import HuggingFaceClient  # noqa: E402


RUNTIME_CATALOG_URL = (
    "https://huggingface.co/datasets/airgpt/embodied-model-index/resolve/main/index.json"
)
COMMUNITY_CATALOG_URL = (
    "https://huggingface.co/datasets/airgpt/embodied-model-index/resolve/main/community-models.json"
)


def merge_candidate(
    catalog: Catalog,
    candidate: HubModel,
    *,
    reserved_models: list[HubModel] | None = None,
) -> Catalog:
    existing = [*catalog.models, *(reserved_models or [])]
    if any(model.slug == candidate.slug for model in existing):
        raise ValueError(f"duplicate slug: {candidate.slug}")
    if any(
        model.source.repo_id.casefold() == candidate.source.repo_id.casefold() for model in existing
    ):
        raise ValueError(f"duplicate repository: {candidate.source.repo_id}")
    if candidate.verification.status != "unverified":
        raise ValueError("new candidates must have unverified status")
    generated_at = datetime.now(UTC).replace(microsecond=0).isoformat().replace("+00:00", "Z")
    return catalog.model_copy(
        update={"generated_at": generated_at, "models": [*catalog.models, candidate]}
    )


def load_candidate(path: Path) -> HubModel:
    payload: dict[str, Any] = json.loads(path.read_text(encoding="utf-8"))
    if "candidate" in payload:
        payload = payload["candidate"]
    payload.pop("upstream", None)
    return HubModel.model_validate(payload)


async def build_preview(candidate_path: Path) -> Catalog:
    candidate = load_candidate(candidate_path)
    upstream = await HuggingFaceClient().model(candidate.source.repo_id, refresh=True)
    if upstream.revision != candidate.source.revision:
        raise ValueError(
            "candidate revision is stale: "
            f"expected {candidate.source.revision}, upstream is {upstream.revision}"
        )
    candidate = candidate.model_copy(
        update={
            "source": candidate.source.model_copy(
                update={"gated": upstream.gated, "license": upstream.license}
            )
        }
    )
    async with httpx.AsyncClient(timeout=20, follow_redirects=True) as client:
        community_response = await client.get(COMMUNITY_CATALOG_URL)
        community_response.raise_for_status()
        runtime_response = await client.get(RUNTIME_CATALOG_URL)
        runtime_response.raise_for_status()
    community_catalog = Catalog.model_validate(community_response.json())
    runtime_catalog = Catalog.model_validate(runtime_response.json())
    return merge_candidate(
        community_catalog,
        candidate,
        reserved_models=runtime_catalog.models,
    )


def publish(catalog: Catalog, token: str) -> None:
    try:
        from huggingface_hub import HfApi
    except ImportError as exc:
        raise RuntimeError("install huggingface_hub before publishing") from exc
    HfApi(token=token).upload_file(
        path_or_fileobj=catalog.model_dump_json(indent=2).encode(),
        path_in_repo="community-models.json",
        repo_id="airgpt/embodied-model-index",
        repo_type="dataset",
        commit_message="Add reviewed AirGPT Hub community model",
    )


def main() -> int:
    parser = argparse.ArgumentParser(description="review an AirGPT Hub candidate manifest")
    parser.add_argument("candidate", type=Path)
    parser.add_argument("--publish", action="store_true", help="write the reviewed catalog to HF")
    args = parser.parse_args()
    catalog = asyncio.run(build_preview(args.candidate))
    print(catalog.model_dump_json(indent=2))
    if args.publish:
        token = os.getenv("HF_TOKEN")
        if not token:
            parser.error("HF_TOKEN is required with --publish")
        publish(catalog, token)
    else:
        print("\nPreview only. Re-run with --publish after manual review.", file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())