87 lines
3.0 KiB
Python
87 lines
3.0 KiB
Python
"""Tests for the curated model-capability table (#81)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from backend.model_capabilities import (
|
|
CAPABILITY_KEYS,
|
|
get_capabilities_for_models,
|
|
get_model_capabilities,
|
|
model_supports_vision,
|
|
)
|
|
|
|
|
|
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
|