Model Picker — Python source

"""Model Picker — filter + rank the AI model catalog.

Language: Python (3.9+, standard library only)
Source:   CosmoDev polyglot showcase port of the Model Picker tool, ported from
          src/lib/modelPicker.ts (the canonical TypeScript implementation).
Live:     https://dev.cosmolabs.org/tools/model-picker
License:  display source — part of CosmoDev's polyglot tool pages.

Design goals:
  - Pure + deterministic; never raises.
  - Functionally equivalent to the TS reference: same inputs -> same outputs.
  - Self-contained: stdlib only (no pip packages).

Port notes: the TS lib imports ``listModels`` from src/lib/ai/models.ts and
lets ``models`` default to the bundled pricing snapshot
(src/data/ai-models.json). A dependency-free port cannot load that file, so
the filter step is inlined here and ``models`` is an explicit **keyword-only**
parameter on both entry points — Python cannot put a required parameter after
an optional one, and silently defaulting the data to ``[]`` would hide bugs.
Sort-key strings ("price" | "context" | "tokensPerDollar" | "released") keep
the TS spellings so the two APIs map 1:1.

Every sort places null/missing values last and breaks ties on id ascending,
so output order is fully deterministic for a given model list.
"""

from __future__ import annotations

import functools
from dataclasses import dataclass
from typing import Dict, List, Literal, Optional

__all__ = [
    "Tier",
    "SortKey",
    "AiModel",
    "ModelFilter",
    "Preset",
    "PRESETS",
    "list_models",
    "pick_models",
]

Tier = Literal["flagship", "balanced", "fast", "budget"]
SortKey = Literal["price", "context", "tokensPerDollar", "released"]


@dataclass
class AiModel:
    """One catalog entry. Fields mirror the TS ``AiModel`` interface; ``None``
    plays the role of TS ``null`` (unpriced / unknown).

    Required fields come first, then everything with a default, because
    dataclass fields without defaults cannot follow ones with defaults.
    """

    id: str
    name: str
    vendor: str
    family: str
    tier: Tier
    context_window: int
    max_output: int
    modalities: List[str]

    # USD per 1M tokens; None for open/unpriced models.
    input_per_m: Optional[float] = None
    output_per_m: Optional[float] = None
    cache_read_per_m: Optional[float] = None
    cache_write_per_m: Optional[float] = None
    knowledge_cutoff: Optional[str] = None
    # ISO date string; None when unknown.
    released: Optional[str] = None
    open_weights: bool = False
    reasoning: bool = False
    tool_call: bool = False


@dataclass
class ModelFilter:
    """Mirrors the TS ``ModelFilter`` interface: every field is optional, and
    ``None`` means "do not apply this constraint" (TS ``undefined``)."""

    vendor: Optional[str] = None
    tier: Optional[Tier] = None
    # Minimum usable context window in tokens.
    min_context: Optional[int] = None
    # Maximum input price (USD per 1M tokens); None-priced models are skipped.
    max_input_per_m: Optional[float] = None
    # Required input modality, e.g. 'image'.
    modality: Optional[str] = None
    # Case-insensitive substring match on id, name, and vendor.
    search: Optional[str] = None


@dataclass
class Preset:
    """A curated entry point: a filter + the sort that makes that filter useful."""

    filter: ModelFilter
    sort: SortKey


PRESETS: Dict[str, Preset] = {
    "long-context": Preset(ModelFilter(min_context=500_000), "context"),
    "cheap-bulk": Preset(ModelFilter(max_input_per_m=1), "price"),
    "flagship": Preset(ModelFilter(tier="flagship"), "tokensPerDollar"),
}


def list_models(models: List[AiModel], filter: Optional[ModelFilter] = None) -> List[AiModel]:
    """Filter the model list. Inline port of ``listModels()`` from
    src/lib/ai/models.ts: ``None`` prices never satisfy ``max_input_per_m``,
    and search lowercases both sides before the substring check."""
    f = filter or ModelFilter()
    q = f.search.lower() if f.search is not None else None
    out: List[AiModel] = []
    for m in models:
        if f.vendor is not None and m.vendor != f.vendor:
            continue
        if f.tier is not None and m.tier != f.tier:
            continue
        if f.min_context is not None and m.context_window < f.min_context:
            continue
        if f.max_input_per_m is not None and (
            m.input_per_m is None or m.input_per_m > f.max_input_per_m
        ):
            continue
        if f.modality is not None and f.modality not in m.modalities:
            continue
        if q is not None and q not in f"{m.id} {m.name} {m.vendor}".lower():
            continue
        out.append(m)
    return out


def _by_id(a: AiModel, b: AiModel) -> int:
    """id ascending — the shared stable tie-break for every sort."""
    if a.id < b.id:
        return -1
    if a.id > b.id:
        return 1
    return 0


def _cmp_price(a: AiModel, b: AiModel) -> int:
    if a.input_per_m is None and b.input_per_m is None:
        return _by_id(a, b)
    if a.input_per_m is None:
        return 1
    if b.input_per_m is None:
        return -1
    if a.input_per_m != b.input_per_m:
        return -1 if a.input_per_m < b.input_per_m else 1
    return _by_id(a, b)


def _cmp_context(a: AiModel, b: AiModel) -> int:
    if a.context_window != b.context_window:
        return -1 if a.context_window > b.context_window else 1
    return _by_id(a, b)


def _cmp_tokens_per_dollar(a: AiModel, b: AiModel) -> int:
    ta = None if a.output_per_m is None else 1_000_000 / a.output_per_m
    tb = None if b.output_per_m is None else 1_000_000 / b.output_per_m
    if ta is None and tb is None:
        return _by_id(a, b)
    if ta is None:
        return 1
    if tb is None:
        return -1
    if ta != tb:
        return 1 if ta < tb else -1  # more tokens per dollar first
    return _by_id(a, b)


def _cmp_released(a: AiModel, b: AiModel) -> int:
    if a.released is None and b.released is None:
        return _by_id(a, b)
    if a.released is None:
        return 1
    if b.released is None:
        return -1
    if a.released != b.released:
        return 1 if a.released < b.released else -1  # newest first
    return _by_id(a, b)


# Sort comparator per key: value order first, None last, id asc on ties.
_COMPARATORS = {
    "price": _cmp_price,
    "context": _cmp_context,
    "tokensPerDollar": _cmp_tokens_per_dollar,
    "released": _cmp_released,
}


def pick_models(
    filter: Optional[ModelFilter] = None,
    sort: SortKey = "price",
    *,
    models: List[AiModel],
) -> List[AiModel]:
    """Filter the model list, then sort it (see ``_COMPARATORS`` for the exact
    order). ``sorted`` is stable, so rows comparing equal — including fully
    identical duplicates — keep their input order, matching
    ``Array.prototype.sort`` in the TS original."""
    ranked = list_models(models, filter)
    ranked.sort(key=functools.cmp_to_key(_COMPARATORS[sort]))
    return ranked