src/models/config.py

"""Loads models.json and resolves client model names to target lists."""   # module docstring
from __future__ import annotations   # defer annotation evaluation (PEP 563)

import json   # stdlib json for parsing the models.json config file
import logging   # logging for fetch/load events
import os   # read MODELS_URL + AES_KEY env vars (runtime fetch + decrypt)
import urllib.request   # fetch models.json from a secret URL if configured
from pathlib import Path   # Path for locating the data directory
from typing import Dict, List, Optional, Tuple   # type aliases for signatures

from cryptography.fernet import Fernet   # AES-128-CBC decrypt of the fetched token blob

logger = logging.getLogger("router.config")   # named logger for this module


def _fetch_models(url: str, aes_key: str) -> dict:   # download + decrypt models.json at startup
    """Download + Fernet-decrypt models.json from a URL (e.g. a raw gist).

    The gist stores the encrypted token (Fernet.encrypt output), NOT raw JSON.
    We fetch the bytes, decrypt with AES_KEY, then parse to a dict.
    """   # docstring
    logger.info("fetching models.json from %s", url[:60])   # log the source (truncated)
    req = urllib.request.Request(url, headers={"User-Agent": "router"})   # build the request
    with urllib.request.urlopen(req, timeout=30) as r:   # open the stream
        token = r.read()   # read the raw encrypted token bytes
    plaintext = Fernet(aes_key.encode()).decrypt(token).decode("utf-8")   # AES-128-CBC decrypt -> JSON text
    return json.loads(plaintext)   # parse to a Python dict


class ModelConfig:   # holds the loaded models.json and resolves client names
    def __init__(self, data_dir: Path):   # data_dir is the folder holding models.json
        self.models: Dict[str, dict] = {}   # exact name -> {"targets": [...]} (raw JSON)
        self._lookup: Dict[str, str] = {}   # lowercased name -> exact name (case-insensitive search)
        self._load(data_dir / "models.json")   # parse the config immediately

    def _load(self, path: Path) -> None:   # read + parse the JSON config
        # If MODELS_URL + AES_KEY are set, fetch the encrypted config from the
        # gist and decrypt at runtime (keeps keys out of the Docker image).
        # Otherwise fall back to the local data/models.json (dev/test only).
        url = os.environ.get("MODELS_URL", "").strip()   # secret gist raw URL, optional
        aes_key = os.environ.get("AES_KEY", "").strip()   # Fernet key for decryption, optional
        if url and aes_key:   # both secrets present -> runtime fetch + decrypt
            raw = _fetch_models(url, aes_key)   # download + decrypt + parse
        else:
            raw = json.loads(path.read_text(encoding="utf-8"))   # local file fallback (utf-8 for locale safety)
        self.models = raw   # store the full dict keyed by model name
        for name in raw:   # build the case-insensitive lookup index
            self._lookup[name.lower()] = name   # map lowercase -> exact case

    def resolve(self, requested: Optional[str]) -> Tuple[Optional[str], Optional[List[dict]]]:   # client name -> (exact_name, targets)
        """Map a client model name to (exact_name, targets). Returns (None, None) if unknown."""   # docstring
        exact = self._lookup.get((requested or "").strip().lower())   # case-insensitive lookup; empty-safe
        if exact is None:   # unknown model name
            return None, None   # caller returns 404
        return exact, self.models[exact].get("targets", [])   # return the exact name + its target list

    def list_models(self) -> List[dict]:   # build the data array for GET /v1/models
        out = []   # the list to return
        for name, cfg in self.models.items():   # iterate in insertion order
            out.append({   # build one OpenAI-shaped model descriptor
                "id": name,   # the client-facing model name
                "object": "model",   # OpenAI type tag
                "targets": len(cfg.get("targets", [])),   # how many targets serve this model
            })
        return out   # the full list

    def total_targets(self) -> int:   # sum of all targets across all models (for status display)
        return sum(len(c.get("targets", [])) for c in self.models.values())   # flatten + count