"""Tests for provider-declared capabilities + the cache (BUG-044). Covers: - payload-shape parsing (Mistral ``capabilities``, OpenRouter ``architecture``); - the in-process cache (TTL, clearing, malformed payloads); - the merge rule: a declaration wins for the flags it mentions, the curated table fills the rest (``backend.model_capabilities``). """ from __future__ import annotations import time import pytest from backend import provider_capabilities as pc from backend.model_capabilities import CAPABILITY_KEYS, get_model_capabilities MISTRAL_PAYLOAD = { "object": "list", "data": [ { "id": "mistral-small-latest", "capabilities": { "completion_chat": True, "vision": True, "audio_transcription": False, "audio_speech": False, }, }, {"id": "mistral-embed", "capabilities": {"completion_chat": False, "vision": False}}, {"id": "mistral-ocr-latest", "capabilities": {"completion_chat": False, "vision": True}}, { "id": "voxtral-mini-latest", "capabilities": {"completion_chat": False, "audio_transcription": True}, }, {"id": "mistral-new-thing", "capabilities": {"completion_chat": True, "vision": True}}, {"id": "mistral-moderation-2603", "capabilities": {"completion_chat": False}}, {"id": "no-capabilities-declared"}, ], } OPENROUTER_PAYLOAD = { "data": [ { "id": "openai/gpt-4o", "architecture": { "input_modalities": ["text", "image"], "output_modalities": ["text"], }, }, { "id": "text-only/model", "architecture": {"input_modalities": ["text"], "output_modalities": ["text"]}, }, { "id": "tts/model", "architecture": {"input_modalities": ["text"], "output_modalities": ["audio"]}, }, { "id": "whisper/model", "architecture": {"input_modalities": ["audio"], "output_modalities": ["text"]}, }, { "id": "video/model", "architecture": {"input_modalities": ["text"], "output_modalities": ["video"]}, }, ] } @pytest.fixture(autouse=True) def _clean_cache(): """The capability cache is process-wide: never leak between tests.""" pc.clear_declared_capabilities() yield pc.clear_declared_capabilities() class TestParsing: def test_mistral_capability_flags_are_mapped(self): declared = pc.parse_declared_capabilities(MISTRAL_PAYLOAD["data"][0]) assert declared == { "chat": True, "vision": True, "audio_transcription": False, "audio_speech": False, } assert set(declared) <= set(CAPABILITY_KEYS) def test_mistral_ocr_is_vision_without_chat(self): declared = pc.parse_declared_capabilities(MISTRAL_PAYLOAD["data"][2]) assert declared == {"chat": False, "vision": True} def test_realtime_transcription_counts_as_transcription(self): entry = { "id": "model", "capabilities": {"completion_chat": False, "audio_transcription_realtime": True}, } assert pc.parse_declared_capabilities(entry) == { "chat": False, "audio_transcription": True, } def test_openrouter_modalities_are_mapped(self): entries = {e["id"]: e for e in OPENROUTER_PAYLOAD["data"]} assert pc.parse_declared_capabilities(entries["openai/gpt-4o"]) == { "chat": True, "vision": True, "images": False, "audio_transcription": False, "audio_speech": False, "video": False, } assert pc.parse_declared_capabilities(entries["tts/model"])["audio_speech"] is True assert pc.parse_declared_capabilities(entries["whisper/model"])["audio_transcription"] is True assert pc.parse_declared_capabilities(entries["video/model"])["video"] is True text_only = pc.parse_declared_capabilities(entries["text-only/model"]) assert text_only is not None and text_only["vision"] is False def test_undeclared_entries_return_none(self): for entry in [None, "not-a-dict", {}, {"architecture": {}}, {"capabilities": {}}]: assert pc.parse_declared_capabilities(entry) is None class TestCache: def test_remember_returns_cached_model_count(self): # 7 entries, 1 without any declaration. assert pc.remember_declared_capabilities("mistral", MISTRAL_PAYLOAD) == 6 assert pc.cache_info()["mistral"]["models"] == 6 def test_unknown_model_returns_none(self): pc.remember_declared_capabilities("mistral", MISTRAL_PAYLOAD) assert pc.get_declared_capabilities("mistral", "not-in-payload") is None def test_unknown_provider_returns_none(self): pc.remember_declared_capabilities("mistral", MISTRAL_PAYLOAD) assert pc.get_declared_capabilities("openrouter", "mistral-small-latest") is None def test_payload_can_be_a_bare_list(self): assert pc.remember_declared_capabilities("mistral", MISTRAL_PAYLOAD["data"]) == 6 def test_shape_that_declares_nothing_is_not_cached(self): """Gemini-shaped payloads (``{"models": [...]}``) declare nothing.""" gemini_like = {"models": [{"name": "models/gemini-2.0-flash"}]} assert pc.remember_declared_capabilities("gemini", gemini_like) == 0 assert "gemini" not in pc.cache_info() def test_gemini_models_prefix_is_normalised(self): payload = [{"name": "models/gemini-x", "capabilities": {"completion_chat": True}}] pc.remember_declared_capabilities("gemini", payload) assert pc.get_declared_capabilities("gemini", "gemini-x") == {"chat": True} assert pc.get_declared_capabilities("gemini", "models/gemini-x") == {"chat": True} def test_snapshot_expires_after_ttl(self, monkeypatch): pc.remember_declared_capabilities("mistral", MISTRAL_PAYLOAD) monkeypatch.setattr(pc, "TTL_SECONDS", 0.01) time.sleep(0.05) assert pc.get_declared_capabilities("mistral", "mistral-small-latest") is None assert pc.cache_info()["mistral"]["expired"] is True def test_clear_one_provider_only(self): pc.remember_declared_capabilities("mistral", MISTRAL_PAYLOAD) pc.remember_declared_capabilities("openrouter", OPENROUTER_PAYLOAD) pc.clear_declared_capabilities("mistral") assert pc.get_declared_capabilities("mistral", "mistral-small-latest") is None assert pc.get_declared_capabilities("openrouter", "openai/gpt-4o") is not None class TestMergeWithCuratedTable: def test_declaration_fixes_a_model_the_curated_table_never_heard_of(self): model = "mistral-new-thing" assert get_model_capabilities("mistral", model)["vision"] is False # curated alone pc.remember_declared_capabilities("mistral", MISTRAL_PAYLOAD) caps = get_model_capabilities("mistral", model) assert caps["vision"] is True assert caps["chat"] is True assert set(caps) == set(CAPABILITY_KEYS) def test_declaration_revokes_a_wrong_curated_guess(self): """voxtral is audio-only: the curated default called it a chat model.""" assert get_model_capabilities("mistral", "voxtral-mini-latest")["chat"] is True assert get_model_capabilities("mistral", "voxtral-mini-latest")["audio_transcription"] is False pc.remember_declared_capabilities("mistral", MISTRAL_PAYLOAD) caps = get_model_capabilities("mistral", "voxtral-mini-latest") assert caps["chat"] is False assert caps["audio_transcription"] is True def test_curated_table_fills_flags_the_provider_stays_silent_about(self): """Mistral never declares ``embedding``: only the absence of chat.""" pc.remember_declared_capabilities("mistral", MISTRAL_PAYLOAD) caps = get_model_capabilities("mistral", "mistral-embed") assert caps["chat"] is False assert caps["embeddings"] is True def test_openrouter_vision_comes_from_the_declaration(self): pc.remember_declared_capabilities("openrouter", OPENROUTER_PAYLOAD) assert get_model_capabilities("openrouter", "openai/gpt-4o")["vision"] is True assert get_model_capabilities("openrouter", "text-only/model")["vision"] is False def test_clearing_the_cache_restores_the_curated_answer(self): pc.remember_declared_capabilities("mistral", MISTRAL_PAYLOAD) assert get_model_capabilities("mistral", "mistral-new-thing")["vision"] is True pc.clear_declared_capabilities() assert get_model_capabilities("mistral", "mistral-new-thing")["vision"] is False