partner_registry.py

"""Loader for the downstream partner's model registry.

The partner ships a JSON list of model entries, each with an `id`, `tier`,
`scores` (per-category 1-10), and `cost` (1-10). This file does not ship the
registry data itself - it is loaded at runtime from a path supplied via the
PARTNER_REGISTRY_PATH environment variable, kept outside source control.
"""

from __future__ import annotations

import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Optional


PARTNER_SCORE_KEYS: tuple[str, ...] = (
    "coding",
    "math",
    "research",
    "creative",
    "chat",
    "roleplay",
    "ideas",
)

TIERS: tuple[str, ...] = ("lite", "standard", "pro", "ultra")


@dataclass(frozen=True)
class PartnerModel:
    id: str
    tier: str
    is_open_router: bool
    strengths: tuple[str, ...]
    scores: dict[str, int]
    cost: int

    def fits_tier(self, tier_set: set[str]) -> bool:
        return self.tier in tier_set


@dataclass
class PartnerRegistry:
    models: list[PartnerModel]

    def all(self) -> list[PartnerModel]:
        return list(self.models)

    def by_tier(self, *tiers: str) -> list[PartnerModel]:
        keep = set(tiers)
        return [m for m in self.models if m.tier in keep]

    def get(self, model_id: str) -> Optional[PartnerModel]:
        for m in self.models:
            if m.id == model_id:
                return m
        return None

    def __len__(self) -> int:
        return len(self.models)


def _coerce(entry: dict) -> PartnerModel:
    scores = {k: int(entry.get("scores", {}).get(k, 0)) for k in PARTNER_SCORE_KEYS}
    return PartnerModel(
        id=str(entry["id"]),
        tier=str(entry.get("tier", "standard")).lower(),
        is_open_router=bool(entry.get("isOpenRouter", False)),
        strengths=tuple(entry.get("strengths", [])),
        scores=scores,
        cost=int(entry.get("cost", 5)),
    )


def load_registry(path: str | Path | None = None) -> PartnerRegistry:
    """Loads from one of three sources, in priority order:

    1. The `path` argument, if supplied.
    2. The PARTNER_REGISTRY_JSON env var containing the raw JSON content (used
       in deployments where a file is awkward to ship, e.g. HF Space secrets).
    3. The PARTNER_REGISTRY_PATH env var pointing at a JSON file on disk.
    """
    raw_text: Optional[str] = None
    source = "argument"

    if path is None:
        inline = os.environ.get("PARTNER_REGISTRY_JSON")
        if inline:
            raw_text = inline
            source = "env:PARTNER_REGISTRY_JSON"
        else:
            env_path = os.environ.get("PARTNER_REGISTRY_PATH")
            if env_path:
                path = env_path
                source = "env:PARTNER_REGISTRY_PATH"

    if raw_text is None:
        if path is None:
            raise RuntimeError(
                "no registry source supplied (set PARTNER_REGISTRY_JSON or PARTNER_REGISTRY_PATH)"
            )
        p = Path(path)
        if not p.exists():
            raise FileNotFoundError(f"partner registry JSON not found at {p}")
        raw_text = p.read_text(encoding="utf-8")

    raw = json.loads(raw_text)
    if not isinstance(raw, list):
        raise ValueError(f"partner registry from {source} must be a top-level list")
    models = [_coerce(e) for e in raw]
    if not models:
        raise ValueError(f"partner registry from {source} is empty")
    return PartnerRegistry(models=models)