ObsiGate/backend/model_capabilities.py

"""Model-capability metadata for the AI assistant.

Two layers, in order of trust:

1. **Provider-declared** (:mod:`backend.provider_capabilities`) — when the
   provider publishes per-model capabilities in its models endpoint (Mistral
   ``capabilities``, OpenRouter ``architecture``), that declaration *wins* for
   every flag it mentions. The snapshot is cached by the model-list endpoint.
2. **Curated table** (this module) — a static, hand-maintained map of known
   model-name patterns to capability flags with per-provider defaults. It fills
   the flags the provider stays silent about, and is the only source for
   providers and models that declare nothing (offline, no API key, DeepSeek,
   NVIDIA, QwenCloud, Xiaomi…).

The UI uses the result to show, when a model is selected, which features it
supports (Chat, Embeddings, Rerank, Images, Video, Audio Speech, Audio
Transcriptions, Vision).

The curated layer is intentionally conservative: an unknown model falls back to
the provider default (usually ``chat`` only), so we never claim a capability the
model may not have. A provider declaration, on the other hand, is authoritative
in both directions — it can also *revoke* a flag the curated table guessed
wrongly (e.g. ``mistral-embed`` declares no chat).
"""

from __future__ import annotations

from typing import Any

from backend.provider_capabilities import get_declared_capabilities

# Ordered list of capability keys exposed to the UI. Keep in sync with the
# frontend ``AI_CAPABILITY_KEYS`` and the i18n ``ai.cap_*`` labels.
CAPABILITY_KEYS: tuple[str, ...] = (
    "chat",
    "embeddings",
    "rerank",
    "images",
    "video",
    "audio_speech",
    "audio_transcription",
    "vision",
)


def _caps(**kwargs: bool) -> dict[str, bool]:
    """Build a full capability dict (missing flags default to False)."""
    base = {key: False for key in CAPABILITY_KEYS}
    base.update(kwargs)
    return base


# Provider defaults, used when no model-name pattern matches.
_PROVIDER_DEFAULTS: dict[str, dict[str, bool]] = {
    "deepseek": _caps(chat=True),
    "openrouter": _caps(chat=True),
    "gemini": _caps(chat=True, vision=True, embeddings=True, audio_speech=True),
    "ollama": _caps(chat=True, embeddings=True),
    "nvidia": _caps(chat=True),
    "qwencloud": _caps(chat=True),
    "xiaomi": _caps(chat=True),
    # Mistral: chat only. ``embeddings`` used to be assumed for every Mistral
    # model, which wrongly labelled mistral-large / codestral as embedders
    # (BUG-044); the ``embed`` rule below covers the real embedding models.
    "mistral": _caps(chat=True),
}

# Ordered (substrings, capabilities) rules — the first matching rule wins.
# More specific modalities are listed before the broad vision/chat rule.
_MODEL_RULES: list[tuple[tuple[str, ...], dict[str, bool]]] = [
    # Rerankers.
    (("rerank", "cross-encoder"), _caps(rerank=True)),
    # Embedding models (usually not chat-capable).
    (
        ("text-embedding", "embed", "bge-", "e5-", "nomic-embed", "gte-"),
        _caps(embeddings=True),
    ),
    # Speech-to-text / transcription.
    (
        ("whisper", "transcrib", "-asr", "asr-", "speech-to-text"),
        _caps(audio_transcription=True),
    ),
    # Text-to-speech.
    (
        ("-tts", "text-to-speech", "voiceclone", "voicedesign"),
        _caps(audio_speech=True),
    ),
    # Image generation.
    (
        ("dall-e", "stable-diffusion", "flux", "imagen", "image-gen"),
        _caps(images=True),
    ),
    # Video generation.
    (("veo-", "sora", "video-gen", "-video"), _caps(video=True)),
    # Document OCR models — image input, but not a chat endpoint.
    (("mistral-ocr",), _caps(vision=True)),
    # Vision-capable chat models (multimodal input).
    (
        (
            "gpt-4o",
            "gpt-4.1",
            "gpt-5",
            "o3",
            "o4",
            "claude-3",
            "claude-4",
            "gemini",
            "qwen-vl",
            "-vl-",
            "vl-",
            "vision",
            "llava",
            "pixtral",
            "internvl",
            "minicpm-v",
            "llama-3.2-vision",
            "mimo-vl",
            # Mistral vision families (BUG-044): text-only mistral-large,
            # codestral and voxtral are deliberately absent.
            "ministral",
            "magistral",
            "mistral-small",
            "mistral-medium",
            "mistral-vibe-cli",
            "labs-leanstral",
        ),
        _caps(chat=True, vision=True),
    ),
]


def _curated_capabilities(provider: str, model: str) -> dict[str, bool]:
    """Curated table lookup (model rules first, then the provider default)."""
    if model:
        for needles, caps in _MODEL_RULES:
            if any(needle in model for needle in needles):
                return dict(caps)
    return dict(_PROVIDER_DEFAULTS.get(provider, _caps(chat=True)))


def get_model_capabilities(provider: str, model: str) -> dict[str, bool]:
    """Return the capability flags for a ``provider``/``model`` pair.

    A provider declaration (see :mod:`backend.provider_capabilities`) overrides
    the curated table for every flag it mentions; the curated table supplies the
    rest.

    Args:
        provider: Provider identifier (e.g. ``"deepseek"``). Case-insensitive.
        model: Model identifier (e.g. ``"deepseek-chat"``). May be empty, in
            which case the provider default is returned.

    Returns:
        A dict with every key of :data:`CAPABILITY_KEYS` and boolean values.
    """
    provider = (provider or "").strip().lower()
    name = (model or "").strip().lower()
    caps = _curated_capabilities(provider, name)
    declared = get_declared_capabilities(provider, name)
    if declared:
        caps.update({key: value for key, value in declared.items() if key in CAPABILITY_KEYS})
    return caps


def get_capabilities_for_models(
    provider: str, models: list[str]
) -> dict[str, dict[str, bool]]:
    """Return a ``{model: capabilities}`` map for a list of models."""
    return {model: get_model_capabilities(provider, model) for model in models}


def model_supports_vision(provider: str, model: str) -> bool:
    """True when the given model can accept image input."""
    return bool(get_model_capabilities(provider, model).get("vision"))


def capabilities_payload(provider: str, model: str) -> dict[str, Any]:
    """Serialize capabilities for API responses."""
    return {
        "provider": provider,
        "model": model,
        "capabilities": get_model_capabilities(provider, model),
    }