ai-agent/backend/models.py

"""
The Claude models a spawned session can run on.

The composer lets a run be pointed at a specific model (a tag per model, e.g.
`opus 4.8` / `haiku 4.5`), and the sidecar passes it to `claude --model`. That
list has to stay current without a redeploy every time Anthropic ships a model,
so it is **fetched from the Anthropic API** (`GET /v1/models`, same key the token
counter already uses) and cached in-process for a few hours. When the key is
missing or the call fails we fall back to a small static list, so the picker
always has something to offer.

**One model per family** — only the newest opus / sonnet / haiku / fable. Older
snapshots are filtered out (see `_shape`): four chips, one per family, is the
whole choice a new session needs.

Each entry carries:
  ``id``           the exact ``--model`` argument (a versioned model id);
  ``family``       opus | sonnet | haiku | fable — how the picker groups/colours;
  ``version``      the human version within the family ("4.8"), for the chip label;
  ``alias``        the CLI shorthand (``opus``/``sonnet``/``haiku``) — the frontend
                   sends it instead of the id so a run follows the CLI's own
                   resolution of "latest";
  ``latest``       always true now that only the newest of a family is listed.
"""

import json
import os
import re
import threading
import time
import urllib.error
import urllib.request

API_URL = "https://api.anthropic.com/v1/models?limit=100"
API_VERSION = "2023-06-01"
CACHE_TTL_S = 6 * 3600

# The families the picker offers, in display order. Anything the API returns
# outside these (legacy claude-3-*, etc.) is dropped: the point is a short,
# pickable list, not an exhaustive catalogue.
FAMILIES = ("opus", "sonnet", "haiku", "fable")
# Families whose shorthand `claude --model <alias>` understands.
CLI_ALIASES = {"opus", "sonnet", "haiku"}
# What a fresh spawn runs on when the user hasn't picked anything.
DEFAULT_MODEL = os.environ.get("SPAWN_DEFAULT_MODEL", "opus")

# Context-window size (tokens) per family — the fallback when the API doesn't
# carry `max_input_tokens` for a model (or the API is unreachable). Read by the
# conversation page's context-fill gizmo via ModelInfo.contextWindow.
CONTEXT_WINDOWS = {
    "opus": 1_000_000,
    "sonnet": 1_000_000,
    "fable": 1_000_000,
    "haiku": 200_000,
    "qwen": 262_144,
}
DEFAULT_CONTEXT_WINDOW = 200_000

# ── pi-harness models ────────────────────────────────────────────────────────
# Sessions on the pi.dev harness (services/ai-agent/runner) run open models
# instead of the claude CLI. Static list: ids with a vendor prefix route to
# OpenRouter, bare ids to the EVOX2 LM Studio provider (see runner/cli.mjs).
# Override with PI_MODELS_JSON (a JSON array of the same shape) if needed.
PI_MODELS_DEFAULT = [
    {
        "id": "qwen/qwen3.6-35b-a3b",
        "displayName": "Qwen 3.6 35B-A3B (OpenRouter)",
        "family": "qwen", "version": "3.6", "label": "qwen 3.6",
        "alias": None, "latest": True,
        "harness": "pi", "provider": "openrouter",
        "contextWindow": 262_144,
    },
    {
        "id": "qwen3.6-35b-a3b",
        "displayName": "Qwen 3.6 35B-A3B (EVOX2 LM Studio, local)",
        "family": "qwen", "version": "3.6", "label": "qwen 3.6 local",
        "alias": None, "latest": True,
        "harness": "pi", "provider": "evox2",
        "contextWindow": 32_768,
    },
]


def _pi_models() -> list[dict]:
    raw = os.environ.get("PI_MODELS_JSON", "")
    if raw:
        try:
            models = json.loads(raw)
            if isinstance(models, list):
                return models
        except ValueError:
            pass
    return PI_MODELS_DEFAULT


def context_window(model_id: str) -> int:
    """Context-window size for any model id, via its family (static fallback)."""
    m = (model_id or "").lower()
    for family, size in CONTEXT_WINDOWS.items():
        if family in m:
            return size
    return DEFAULT_CONTEXT_WINDOW


# Used only when the API is unreachable — kept deliberately short.
FALLBACK = [
    {"id": "claude-opus-4-8", "displayName": "Claude Opus 4.8"},
    {"id": "claude-sonnet-5", "displayName": "Claude Sonnet 5"},
    {"id": "claude-fable-5", "displayName": "Claude Fable 5"},
    {"id": "claude-haiku-4-5-20251001", "displayName": "Claude Haiku 4.5"},
]

_lock = threading.Lock()
_cache: tuple[float, list[dict]] | None = None


def _version(model_id: str, family: str) -> str:
    """The version inside a family: `claude-opus-4-8` → "4.8", `claude-sonnet-5` → "5".

    Model ids carry the version as dash-separated digits after the family, with an
    optional trailing snapshot date (`claude-haiku-4-5-20251001`) we drop."""
    rest = model_id.split(f"{family}-", 1)[-1] if f"{family}-" in model_id else ""
    parts = [p for p in rest.split("-") if p.isdigit()]
    if parts and len(parts[-1]) == 8:  # trailing YYYYMMDD snapshot
        parts = parts[:-1]
    return ".".join(parts)


def _shape(raw: list[dict]) -> list[dict]:
    """Turn the API's model records into the picker's entries: one per family.

    The API returns newest-first, so a family's first model is its latest — and
    that is the only one we keep. Older snapshots (opus 4.7, sonnet 4.5, …) are
    dropped: they'd make the tag row a wall of near-identical chips, and there's
    no reason to start a fresh session on a superseded model. Reading a *past*
    conversation is unaffected — its tags come from the transcript, not this list,
    so a run on an older model still renders with its own version."""
    out: list[dict] = []
    seen_family: set[str] = set()
    for m in raw:
        mid = m.get("id") or ""
        family = next((f for f in FAMILIES if re.search(rf"\b{f}\b", mid)), None)
        if not family or family in seen_family:
            continue
        seen_family.add(family)
        out.append({
            "id": mid,
            "displayName": m.get("displayName") or m.get("display_name") or mid,
            "family": family,
            "version": _version(mid, family),
            "alias": family if family in CLI_ALIASES else None,
            "latest": True,
            "harness": "claude",
            # The API's own window size when present (max_input_tokens, added
            # to /v1/models in Mar 2026), else the per-family static value.
            "contextWindow": m.get("max_input_tokens")
                             or CONTEXT_WINDOWS.get(family, DEFAULT_CONTEXT_WINDOW),
        })
    out.sort(key=lambda e: FAMILIES.index(e["family"]))
    return out


def _fetch() -> list[dict]:
    key = os.environ.get("ANTHROPIC_API_KEY", "")
    if not key:
        return _shape(FALLBACK)
    req = urllib.request.Request(
        API_URL, headers={"x-api-key": key, "anthropic-version": API_VERSION})
    try:
        with urllib.request.urlopen(req, timeout=10) as r:
            data = json.loads(r.read().decode())
        models = _shape(data.get("data") or [])
        return models or _shape(FALLBACK)
    except (urllib.error.URLError, OSError, ValueError):
        return _shape(FALLBACK)


def list_models(force: bool = False) -> list[dict]:
    """The pickable models, cached for {CACHE_TTL_S}s (one API call per TTL).

    Claude models (live from the Anthropic API) first, then the pi-harness
    models — the composer renders one chip per entry, and picking a pi entry
    routes the session through the pi runner instead of the claude CLI."""
    global _cache
    with _lock:
        if not force and _cache and time.time() - _cache[0] < CACHE_TTL_S:
            return _cache[1]
    models = _fetch() + _pi_models()  # outside the lock: slow call, no blocking
    with _lock:
        _cache = (time.time(), models)
    return models