Source code for model_capabilities

"""model_capabilities.py — Provider capability registry for LLM model dispatch.

Replaces the fragile ``"gemini" in model.lower()`` substring checks in
:mod:`openrouter_client` with an explicit pattern-matched registry.

Usage::

    from model_capabilities import get_capabilities

    caps = get_capabilities(model)
    if caps.provider == "google":
        ...  # inject reasoning field, apply Gemini-specific sanitization
    if caps.requires_reasoning_field:
        ...
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from fnmatch import fnmatch

logger = logging.getLogger(__name__)


[docs] @dataclass(frozen=True) class ModelCapabilities: """Immutable capability record for a matched model pattern. A frozen dataclass describing what an LLM endpoint supports, so callers can branch on provider quirks instead of repeating brittle substring checks. The flags drive provider-specific request shaping in the OpenRouter client -- for example whether to inject a separate ``reasoning`` / ``thinking`` field, whether a ``system`` role is accepted, and whether multimodal parts or tool calls are allowed. Instances are constructed in this module's ``_REGISTRY`` (one per glob pattern) and the ``_DEFAULT`` fallback, and are returned by ``get_capabilities``. Being frozen, a single instance is shared by every model that matches a given pattern, so it must never be mutated by callers. """ provider: str = "openrouter" """Canonical provider name: 'google', 'anthropic', 'openai', 'deepseek', 'openrouter'.""" requires_reasoning_field: bool = False """True when the API requires a separate ``reasoning`` / ``thinking`` JSON field.""" supports_multimodal: bool = True """True when the model accepts image/audio/video content parts.""" supports_system_role: bool = True """True when the model accepts a ``system`` role message.""" supports_tool_calls: bool = True """True when the model supports function/tool calling."""
# Ordered list of (glob_pattern, capabilities). First match wins. # Patterns are matched case-insensitively against the full model identifier. _REGISTRY: list[tuple[str, ModelCapabilities]] = [ # Google / Gemini ("gemini-*", ModelCapabilities(provider="google", requires_reasoning_field=True)), ( "google/gemini-*", ModelCapabilities(provider="google", requires_reasoning_field=True), ), ("gemini/*", ModelCapabilities(provider="google", requires_reasoning_field=True)), # Anthropic / Claude ("anthropic/*", ModelCapabilities(provider="anthropic")), ("claude-*", ModelCapabilities(provider="anthropic")), # OpenAI — o-series lacks system role support ("o1-*", ModelCapabilities(provider="openai", supports_system_role=False)), ("o3-*", ModelCapabilities(provider="openai", supports_system_role=False)), ("o4-*", ModelCapabilities(provider="openai", supports_system_role=False)), ("openai/o1-*", ModelCapabilities(provider="openai", supports_system_role=False)), ("openai/o3-*", ModelCapabilities(provider="openai", supports_system_role=False)), ("openai/*", ModelCapabilities(provider="openai")), ("gpt-*", ModelCapabilities(provider="openai")), # DeepSeek ("deepseek/*", ModelCapabilities(provider="deepseek")), ("deepseek-*", ModelCapabilities(provider="deepseek")), # Meta / Llama (via OpenRouter) ("meta-llama/*", ModelCapabilities(provider="meta")), ("llama-*", ModelCapabilities(provider="meta")), ] _DEFAULT = ModelCapabilities()
[docs] def get_capabilities(model: str) -> ModelCapabilities: """Return the :class:`ModelCapabilities` for *model*. Matches the model identifier against ``_REGISTRY`` using :func:`fnmatch.fnmatch` (case-insensitive). Falls back to legacy substring heuristics for unregistered models that still contain well-known provider keywords, then returns safe defaults. Args: model: The model identifier string (e.g. ``"gemini-2.5-pro"``). Returns: :class:`ModelCapabilities` instance; never raises. """ if not model: return _DEFAULT m = model.lower().strip() for pattern, caps in _REGISTRY: if fnmatch(m, pattern): return caps # Legacy substring fallback: preserves existing behaviour for custom # aliases or unregistered models that contain provider keywords. if "gemini" in m: logger.debug( "model_capabilities: unregistered model %r matched via 'gemini' substring fallback", model, ) return ModelCapabilities(provider="google", requires_reasoning_field=True) if "claude" in m: logger.debug( "model_capabilities: unregistered model %r matched via 'claude' substring fallback", model, ) return ModelCapabilities(provider="anthropic") return _DEFAULT
# --------------------------------------------------------------------------- # Thin compatibility shims — used by any external callers that previously # called the private helpers directly. Deprecated; use get_capabilities(). # --------------------------------------------------------------------------- def _model_targets_gemini(model: str) -> bool: # noqa: N802 """Report whether a model routes to Google's Gemini API. Deprecated backwards-compatibility shim that simply checks whether ``get_capabilities(model).provider`` is ``"google"``; prefer that call directly in new code. It exists so the Gemini-specific branches in the OpenRouter client keep working unchanged. Called by ``openrouter_client/transport.py`` and ``openrouter_client/sanitization.py`` to gate Gemini-only request shaping (reasoning-field injection and schema sanitization), and exercised by the model-capabilities and multimodal test suites. Args: model: The model identifier string (e.g. ``"gemini-2.5-pro"``). Returns: bool: True when the model's resolved provider is Google/Gemini. """ return get_capabilities(model).provider == "google" def _model_targets_claude(model: str) -> bool: # noqa: N802 """Report whether a model routes to Anthropic's Claude API. Deprecated backwards-compatibility shim that simply checks whether ``get_capabilities(model).provider`` is ``"anthropic"``; prefer that call directly in new code. It exists so the Claude-specific branches in the OpenRouter client keep working unchanged. Called by ``openrouter_client/transport.py`` and ``openrouter_client/sanitization.py`` to gate Claude-only handling (image clamping and temperature capping, among others), and exercised by the model-capabilities and image-clamp test suites. Args: model: The model identifier string (e.g. ``"anthropic/claude-3.5-sonnet"``). Returns: bool: True when the model's resolved provider is Anthropic/Claude. """ return get_capabilities(model).provider == "anthropic"