Kansho/backend/model_catalog.py

"""Cached model window metadata for journal budget checks.

OpenRouter exposes context_length and top_provider.max_completion_tokens via
GET /api/v1/models. The catalog is cached so journal calls do not refetch on
every generate. Missing metadata fails closed for journal generation.
"""
from __future__ import annotations

import os
import threading
import time
from urllib.parse import urlparse

import httpx

from prompt_budget import (
    ERROR_MODEL_METADATA_UNKNOWN,
    JOURNAL_TARGET_COMPLETION_TOKENS,
    JournalBudgetError,
    ModelWindow,
    _int_env,
)

DEFAULT_TTL_SECONDS = 3600
FAKE_CONTEXT = 32_768
_lock = threading.Lock()
# url -> (fetched_at, {model_id: ModelWindow fields as dict})
_catalog: dict[str, tuple[float, dict[str, dict]]] = {}
_overrides: dict[str, ModelWindow] = {}


def catalog_ttl() -> int:
    raw = (os.environ.get("KANSHO_MODEL_CATALOG_TTL_SECONDS") or "").strip()
    if not raw:
        return DEFAULT_TTL_SECONDS
    try:
        value = int(raw)
    except ValueError:
        return DEFAULT_TTL_SECONDS
    return max(60, value)


def reset_catalog() -> None:
    with _lock:
        _catalog.clear()
        _overrides.clear()


def set_metadata_override(model: str, window: ModelWindow) -> None:
    with _lock:
        _overrides[model] = window


def is_openrouter_url(url: str) -> bool:
    host = (urlparse(url or "").hostname or "").lower()
    return host.endswith("openrouter.ai")


def models_url(chat_url: str) -> str:
    raw = (chat_url or "").rstrip("/")
    if raw.endswith("/chat/completions"):
        return raw[: -len("/chat/completions")] + "/models"
    if is_openrouter_url(raw):
        return "https://openrouter.ai/api/v1/models"
    return raw + "/models" if raw else "https://openrouter.ai/api/v1/models"


def _from_env(model: str, *, provider: str | None, source: str) -> ModelWindow | None:
    context = _int_env("KANSHO_PROVIDER_CONTEXT_LENGTH", 0)
    completion = _int_env("KANSHO_PROVIDER_MAX_COMPLETION_TOKENS", 0)
    if context <= 0:
        return None
    return ModelWindow(
        model=model,
        context_length=context,
        max_completion_tokens=completion or JOURNAL_TARGET_COMPLETION_TOKENS,
        source=source,
        provider=provider,
        cached=False,
    )


def _parse_model_row(row: dict) -> ModelWindow | None:
    model_id = str(row.get("id") or "").strip()
    if not model_id:
        return None
    top = row.get("top_provider") if isinstance(row.get("top_provider"), dict) else {}
    context = row.get("context_length") or top.get("context_length")
    completion = top.get("max_completion_tokens")
    if completion is None:
        limits = row.get("per_request_limits") if isinstance(row.get("per_request_limits"), dict) else {}
        completion = limits.get("completion_tokens") or limits.get("max_tokens")
    try:
        context_n = int(context)
    except (TypeError, ValueError):
        return None
    try:
        completion_n = int(completion) if completion is not None else None
    except (TypeError, ValueError):
        completion_n = None
    if context_n <= 0:
        return None
    return ModelWindow(
        model=model_id,
        context_length=context_n,
        max_completion_tokens=completion_n,
        source="openrouter_models",
        provider="openrouter",
        cached=True,
    )


def _request_models_payload(url: str, key: str) -> list:
    headers = {"Content-Type": "application/json"}
    if key:
        headers["Authorization"] = f"Bearer {key}"
    try:
        response = httpx.get(url, headers=headers, timeout=15.0)
    except httpx.HTTPError as exc:
        raise JournalBudgetError(
            ERROR_MODEL_METADATA_UNKNOWN,
            status_code=503,
            diagnostics={"models_url": url, "reason": "unreachable"},
        ) from exc
    if response.status_code >= 400:
        raise JournalBudgetError(
            ERROR_MODEL_METADATA_UNKNOWN,
            status_code=503,
            diagnostics={"models_url": url, "http_status": response.status_code},
        )
    try:
        payload = response.json()
    except ValueError as exc:
        raise JournalBudgetError(
            ERROR_MODEL_METADATA_UNKNOWN,
            status_code=503,
            diagnostics={"models_url": url, "reason": "invalid_json"},
        ) from exc
    rows = payload.get("data") if isinstance(payload, dict) else payload
    if not isinstance(rows, list):
        raise JournalBudgetError(
            ERROR_MODEL_METADATA_UNKNOWN,
            status_code=503,
            diagnostics={"models_url": url, "reason": "unexpected_shape"},
        )
    return rows


def _fetch_catalog(url: str, key: str) -> dict[str, dict]:
    rows = _request_models_payload(url, key)
    parsed: dict[str, dict] = {}
    for row in rows:
        if not isinstance(row, dict):
            continue
        window = _parse_model_row(row)
        if window:
            parsed[window.model] = {
                "model": window.model,
                "context_length": window.context_length,
                "max_completion_tokens": window.max_completion_tokens,
                "source": window.source,
                "provider": window.provider,
            }
    if not parsed:
        raise JournalBudgetError(
            ERROR_MODEL_METADATA_UNKNOWN,
            status_code=503,
            diagnostics={"models_url": url, "reason": "empty_catalog"},
        )
    return parsed


def _window_from_cache(fields: dict, *, model: str) -> ModelWindow:
    return ModelWindow(
        model=model,
        context_length=int(fields["context_length"]),
        max_completion_tokens=fields.get("max_completion_tokens"),
        source=fields.get("source") or "openrouter_models",
        provider=fields.get("provider"),
        cached=True,
    )


def resolve_generate_metadata(config, *, now: float | None = None) -> ModelWindow:
    """Return a window for the generate provider. Fail closed if unknown."""
    model = getattr(config, "model", None) or ""
    with _lock:
        override = _overrides.get(model)
    if override:
        return override
    if getattr(config, "mode", None) == "fake" or model == "fake":
        return ModelWindow(
            model=model or "fake",
            context_length=FAKE_CONTEXT,
            max_completion_tokens=JOURNAL_TARGET_COMPLETION_TOKENS,
            source="fake",
            provider="fake",
            cached=False,
        )
    env_window = _from_env(model, provider=getattr(config, "name", None), source="env")
    if getattr(config, "local", False):
        if env_window:
            return env_window
        raise JournalBudgetError(
            ERROR_MODEL_METADATA_UNKNOWN,
            status_code=503,
            diagnostics={
                "model": model,
                "provider": getattr(config, "name", None),
                "reason": "local_context_unconfigured",
            },
        )
    url = getattr(config, "url", "") or ""
    if not is_openrouter_url(url):
        if env_window:
            return env_window
        raise JournalBudgetError(
            ERROR_MODEL_METADATA_UNKNOWN,
            status_code=503,
            diagnostics={
                "model": model,
                "provider": getattr(config, "name", None),
                "reason": "non_openrouter_unconfigured",
            },
        )
    catalog_key = models_url(url)
    clock = time.time() if now is None else now
    ttl = catalog_ttl()
    with _lock:
        cached = _catalog.get(catalog_key)
        if cached and clock - cached[0] < ttl and model in cached[1]:
            return _window_from_cache(cached[1][model], model=model)
        stale = cached
    try:
        fetched = _fetch_catalog(catalog_key, getattr(config, "key", "") or "")
    except JournalBudgetError:
        if stale and model in stale[1]:
            window = _window_from_cache(stale[1][model], model=model)
            window.source = "openrouter_models_stale"
            return window
        if env_window:
            env_window.source = "env_fallback"
            return env_window
        raise
    with _lock:
        _catalog[catalog_key] = (clock, fetched)
    if model not in fetched:
        if env_window:
            env_window.source = "env_fallback"
            return env_window
        raise JournalBudgetError(
            ERROR_MODEL_METADATA_UNKNOWN,
            status_code=503,
            diagnostics={"model": model, "reason": "model_not_in_catalog"},
        )
    return _window_from_cache(fetched[model], model=model)


def list_provider_models(config) -> list[dict]:
    """Model ids for the admin picker. Soft-fail; free-text remains valid."""
    if not config or getattr(config, "mode", None) == "fake":
        return []
    url = getattr(config, "url", "") or ""
    if not url:
        return []
    key = getattr(config, "key", "") or ""
    items = _list_openai_models(url, key)
    if items:
        return items
    return _list_ollama_tags(url)


def list_models_for_url(url: str, key: str = "") -> list[dict]:
    from providers import ProviderConfig, is_local_url

    url = (url or "").strip()
    if not url:
        return []
    local = is_local_url(url)
    return list_provider_models(
        ProviderConfig(
            role="generate",
            name="ollama" if local else "openrouter",
            mode="http",
            url=url,
            model="",
            key=key,
            local=local,
            zdr=True,
            no_train=True,
        )
    )


def _list_openai_models(chat_url: str, key: str) -> list[dict]:
    catalog_key = models_url(chat_url)
    try:
        rows = _request_models_payload(catalog_key, key)
    except JournalBudgetError:
        return []
    return _ids_from_rows(rows, id_keys=("id",))


def _list_ollama_tags(chat_url: str) -> list[dict]:
    parsed = urlparse(chat_url or "")
    if not parsed.scheme or not parsed.netloc:
        return []
    tags_url = f"{parsed.scheme}://{parsed.netloc}/api/tags"
    headers = {"Content-Type": "application/json"}
    try:
        response = httpx.get(tags_url, headers=headers, timeout=8.0)
    except httpx.HTTPError:
        return []
    if response.status_code >= 400:
        return []
    try:
        payload = response.json()
    except ValueError:
        return []
    rows = payload.get("models") if isinstance(payload, dict) else payload
    if not isinstance(rows, list):
        return []
    return _ids_from_rows(rows, id_keys=("name", "model", "id"))


def _ids_from_rows(rows: list, *, id_keys: tuple[str, ...]) -> list[dict]:
    seen: set[str] = set()
    items: list[dict] = []
    for row in rows:
        if not isinstance(row, dict):
            continue
        model_id = ""
        for key in id_keys:
            model_id = str(row.get(key) or "").strip()
            if model_id:
                break
        if not model_id or model_id in seen:
            continue
        seen.add(model_id)
        name = str(row.get("name") or model_id).strip() or model_id
        items.append({"id": model_id, "name": name})
    items.sort(key=lambda item: item["id"])
    return items