Mujoco_WASM/decision_server/model_catalog.py

"""Bounded public model metadata, never an arbitrary URL proxy or paid request."""

import asyncio
import re
import time

import aiohttp

from .protocol import loads
from .web_config import DEEPSEEK_MODELS


class ModelCatalog:
    def __init__(self):
        self.models = {}
        self.checked_at = 0
        self.available = False
        self.lock = asyncio.Lock()

    async def refresh(self, session):
        if self.lock.locked():
            return  # Don't accumulate an unbounded queue while the upstream is unavailable.
        async with self.lock:
            if time.monotonic() - self.checked_at < 300:
                return
            self.checked_at = time.monotonic()
            try:
                async with session.get(
                    "https://openrouter.ai/api/v1/models",
                    allow_redirects=False,
                    timeout=aiohttp.ClientTimeout(total=10),
                ) as response:
                    if response.status != 200:
                        raise ValueError("catalog_unavailable")
                    raw = bytearray()
                    async for chunk in response.content.iter_chunked(65536):
                        raw.extend(chunk)
                        if len(raw) > 8 * 1024 * 1024:
                            raise ValueError("catalog_too_large")
                    result = loads(raw.decode("utf-8"))
                    models = {}
                    for item in result.get("data", []):
                        ident = item.get("id")
                        if (
                            isinstance(ident, str)
                            and re.fullmatch(r"[A-Za-z0-9_./:-]{1,128}", ident)
                            and "structured_outputs" in item.get("supported_parameters", [])
                        ):
                            models[ident] = str(item.get("name", ident))[:160]
                        if (
                            len(models) >= 256
                        ):  # Keep the public response within the browser byte limit.
                            break
                    if not models:
                        raise ValueError("no_structured_models")
                    self.models, self.available = models, True
            except Exception:
                # Never use upstream text in a response; keep only a bounded known-good catalog.
                self.available = False

    def public(self):
        return {
            "models": [{"provider": "deepseek", "id": m, "name": m} for m in DEEPSEEK_MODELS]
            + [
                {"provider": "openrouter", "id": m, "name": n}
                for m, n in sorted(self.models.items())
            ],
            "openrouterAvailable": self.available,
            "cached": bool(self.models) and not self.available,
        }