ObsiGate/tests/test_model_capabilities.py

"""Tests for the curated model-capability table (#81)."""

from __future__ import annotations

import pytest

from backend.model_capabilities import (
    CAPABILITY_KEYS,
    get_capabilities_for_models,
    get_model_capabilities,
    model_supports_vision,
)


@pytest.fixture(autouse=True)
def _no_declared_caps():
    """Curated-table tests must not see the process-wide declaration cache."""
    from backend.provider_capabilities import clear_declared_capabilities

    clear_declared_capabilities()
    yield
    clear_declared_capabilities()


class TestCapabilityShape:
    def test_every_result_has_all_keys(self):
        for provider, model in [
            ("deepseek", "deepseek-chat"),
            ("gemini", "gemini-2.0-flash"),
            ("qwencloud", "qwen-vl-max"),
            ("openrouter", "unknown-model"),
        ]:
            caps = get_model_capabilities(provider, model)
            assert set(caps.keys()) == set(CAPABILITY_KEYS)
            assert all(isinstance(v, bool) for v in caps.values())


class TestVisionDetection:
    def test_vision_models_detected(self):
        assert model_supports_vision("qwencloud", "qwen-vl-max")
        assert model_supports_vision("openrouter", "openai/gpt-4o")
        assert model_supports_vision("gemini", "gemini-2.0-flash")
        assert model_supports_vision("mistral", "pixtral-large-latest")

    def test_text_models_not_vision(self):
        assert not model_supports_vision("deepseek", "deepseek-chat")
        assert not model_supports_vision("xiaomi", "mimo-v2.5-pro")

    def test_vision_implies_chat(self):
        caps = get_model_capabilities("qwencloud", "qwen-vl-plus")
        assert caps["vision"] is True
        assert caps["chat"] is True


class TestSpecialModalities:
    def test_embeddings(self):
        caps = get_model_capabilities("openai", "text-embedding-3-small")
        assert caps["embeddings"] is True
        assert caps["chat"] is False

    def test_rerank(self):
        caps = get_model_capabilities("nvidia", "nvidia/llama-3.2-nv-rerankqa-1b-v2")
        assert caps["rerank"] is True

    def test_audio_transcription(self):
        caps = get_model_capabilities("xiaomi", "mimo-v2.5-asr")
        assert caps["audio_transcription"] is True

    def test_audio_speech(self):
        caps = get_model_capabilities("xiaomi", "mimo-v2.5-tts")
        assert caps["audio_speech"] is True

    def test_image_generation(self):
        caps = get_model_capabilities("openai", "dall-e-3")
        assert caps["images"] is True


class TestProviderDefaults:
    def test_unknown_model_uses_provider_default(self):
        caps = get_model_capabilities("deepseek", "some-new-model")
        assert caps["chat"] is True
        assert caps["vision"] is False

    def test_unknown_provider_defaults_to_chat(self):
        caps = get_model_capabilities("nope", "mystery")
        assert caps["chat"] is True

    def test_empty_model_uses_provider_default(self):
        assert get_model_capabilities("gemini", "")["vision"] is True


class TestBatch:
    def test_capabilities_for_models_map(self):
        result = get_capabilities_for_models(
            "qwencloud", ["qwen-max", "qwen-vl-max"]
        )
        assert result["qwen-max"]["vision"] is False
        assert result["qwen-vl-max"]["vision"] is True


class TestMistralFamilies:
    """BUG-044 — the curated fallback must know the current Mistral families.

    This fallback is what answers when the provider declaration is unavailable
    (no API key, offline, cold cache), so it has to agree with what
    ``api.mistral.ai/v1/models`` declares today. Before the fix, no Mistral
    model at all was reported as vision-capable — only the dead ``pixtral``
    pattern matched.
    """

    @pytest.mark.parametrize(
        "model",
        [
            "mistral-small-latest",
            "mistral-small-2603",
            "mistral-medium-latest",
            "mistral-medium-3.5",
            "ministral-3b-latest",
            "ministral-8b-2512",
            "ministral-14b-latest",
            "magistral-small-latest",
            "magistral-medium-latest",
            "mistral-vibe-cli-latest",
            "pixtral-12b-2409",
        ],
    )
    def test_vision_chat_families(self, model):
        caps = get_model_capabilities("mistral", model)
        assert caps["vision"] is True, model
        assert caps["chat"] is True, model

    @pytest.mark.parametrize(
        "model",
        [
            "mistral-large-latest",
            "codestral-latest",
            "open-mixtral-8x7b",
            "voxtral-small-latest",
            "mistral-code-latest",
        ],
    )
    def test_text_only_models_are_not_vision(self, model):
        caps = get_model_capabilities("mistral", model)
        assert caps["vision"] is False, model
        assert caps["chat"] is True, model

    def test_ocr_models_are_vision_without_chat(self):
        caps = get_model_capabilities("mistral", "mistral-ocr-latest")
        assert caps["vision"] is True
        assert caps["chat"] is False

    def test_embeddings_only_for_the_embed_models(self):
        assert get_model_capabilities("mistral", "mistral-embed")["embeddings"] is True
        # Regression: every Mistral model used to inherit embeddings=True from
        # the provider default, wrongly branding chat models as embedders.
        assert get_model_capabilities("mistral", "mistral-large-latest")["embeddings"] is False
        assert get_model_capabilities("mistral", "some-unknown-mistral")["embeddings"] is False

    def test_current_api_vision_models_are_all_covered(self):
        """Every model the API declares vision=true is vision-capable offline.

        List captured from ``GET https://api.mistral.ai/v1/models`` (2026-09-15);
        the live declaration layer covers renames, this guards the offline path.
        """
        api_vision_models = [
            "magistral-medium-latest", "magistral-small-latest",
            "ministral-14b-2512", "ministral-14b-latest",
            "ministral-3b-2512", "ministral-3b-latest",
            "ministral-8b-2512", "ministral-8b-latest",
            "mistral-medium", "mistral-medium-2604", "mistral-medium-3",
            "mistral-medium-3-5", "mistral-medium-3.5", "mistral-medium-latest",
            "mistral-ocr-2512", "mistral-ocr-3", "mistral-ocr-3-0",
            "mistral-ocr-4", "mistral-ocr-4-0", "mistral-ocr-4-1",
            "mistral-ocr-latest",
            "mistral-small-2603", "mistral-small-latest",
            "mistral-vibe-cli-fast", "mistral-vibe-cli-latest",
            "mistral-vibe-cli-with-tools",
        ]
        missing = [m for m in api_vision_models if not model_supports_vision("mistral", m)]
        assert missing == [], f"curated fallback misses vision for: {missing}"