136 lines
5.4 KiB
Python
136 lines
5.4 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from app.ai_provider.config.models import ModelConfig
|
|
from app.ai_provider.core.registry import ModelRegistry
|
|
from app.ai_provider.routing.multimodal import MultiModalDecision
|
|
from app.ai_provider.routing.selector import ModelSelector
|
|
from app.ai_provider.types.errors import AIProviderError
|
|
|
|
|
|
class TestModelSelector:
|
|
@pytest.fixture
|
|
def registry(self):
|
|
models = [
|
|
ModelConfig(
|
|
name="fast-model",
|
|
provider="deepseek",
|
|
model="deepseek-v4-flash",
|
|
format="openai",
|
|
tags=["default"],
|
|
weight=100,
|
|
),
|
|
ModelConfig(
|
|
name="pro-model",
|
|
provider="deepseek",
|
|
model="deepseek-v4-pro",
|
|
format="openai",
|
|
tags=["fallback"],
|
|
weight=50,
|
|
),
|
|
ModelConfig(
|
|
name="vision-model",
|
|
provider="anthropic",
|
|
model="claude-sonnet-4-6",
|
|
format="anthropic",
|
|
supports_vision=True,
|
|
supports_thinking=True,
|
|
supports_tools=True,
|
|
),
|
|
ModelConfig(
|
|
name="thinking-model",
|
|
provider="deepseek",
|
|
model="deepseek-v4-thinking",
|
|
format="openai",
|
|
supports_reasoning_effort=True,
|
|
supports_thinking=True,
|
|
),
|
|
]
|
|
return ModelRegistry(models)
|
|
|
|
def test_select_default_model(self, registry):
|
|
selector = ModelSelector(registry)
|
|
model = selector.select({})
|
|
assert model.name == "fast-model"
|
|
|
|
def test_select_with_preferred_model(self, registry):
|
|
selector = ModelSelector(registry)
|
|
model = selector.select({"preferred": "pro-model"})
|
|
assert model.name == "pro-model"
|
|
|
|
def test_select_preferred_with_capability(self, registry):
|
|
selector = ModelSelector(registry)
|
|
model = selector.select({"preferred": "vision-model", "capabilities": {"vision"}})
|
|
assert model.name == "vision-model"
|
|
|
|
def test_select_preferred_missing_required_capability(self, registry):
|
|
"""Preferred model doesn't have required capability — falls back."""
|
|
selector = ModelSelector(registry)
|
|
model = selector.select({"preferred": "fast-model", "capabilities": {"vision"}})
|
|
assert model.name != "fast-model"
|
|
assert "vision" in model.resolved_capabilities()
|
|
assert model.resolved_capabilities()["vision"] is True
|
|
|
|
def test_select_for_vision(self, registry):
|
|
selector = ModelSelector(registry)
|
|
model = selector.select({"capabilities": {"vision"}})
|
|
assert model.name == "vision-model"
|
|
|
|
def test_select_for_thinking(self, registry):
|
|
selector = ModelSelector(registry)
|
|
model = selector.select({"capabilities": {"thinking"}})
|
|
assert model.resolved_capabilities()["thinking"] is True
|
|
|
|
def test_select_for_tool_use(self, registry):
|
|
selector = ModelSelector(registry)
|
|
model = selector.select({"capabilities": {"tool_use"}})
|
|
assert model.resolved_capabilities()["tool_use"] is True
|
|
|
|
def test_select_for_multiple_capabilities(self, registry):
|
|
selector = ModelSelector(registry)
|
|
model = selector.select({"capabilities": {"vision", "thinking"}})
|
|
assert model.resolved_capabilities()["vision"] is True
|
|
assert model.resolved_capabilities()["thinking"] is True
|
|
|
|
def test_select_from_multi_modal_decision(self, registry):
|
|
selector = ModelSelector(registry)
|
|
model = selector.select_from_decisions(
|
|
{MultiModalDecision.VISION, MultiModalDecision.THINKING}
|
|
)
|
|
assert model.resolved_capabilities()["vision"] is True
|
|
assert model.resolved_capabilities()["thinking"] is True
|
|
|
|
def test_select_from_plain_text_decision(self, registry):
|
|
selector = ModelSelector(registry)
|
|
model = selector.select_from_decisions({MultiModalDecision.PLAIN_TEXT})
|
|
assert model.name == "fast-model"
|
|
|
|
def test_select_with_invalid_preferred(self, registry):
|
|
selector = ModelSelector(registry)
|
|
with pytest.raises(AIProviderError):
|
|
selector.select({"preferred": "nonexistent-model"})
|
|
|
|
def test_select_with_capability_no_match(self, registry):
|
|
"""No model supports an unknown capability."""
|
|
selector = ModelSelector(registry)
|
|
with pytest.raises(AIProviderError):
|
|
selector.select({"capabilities": {"telepathy"}})
|
|
|
|
def test_select_allows_preferred_with_superset_capabilities(self, registry):
|
|
"""Preferred model has MORE capabilities than required — should still match."""
|
|
selector = ModelSelector(registry)
|
|
model = selector.select({"preferred": "vision-model", "capabilities": {"tool_use"}})
|
|
assert model.name == "vision-model"
|
|
|
|
def test_select_from_decisions_with_empty_set(self, registry):
|
|
selector = ModelSelector(registry)
|
|
model = selector.select_from_decisions(set())
|
|
assert model.name == "fast-model"
|
|
|
|
def test_select_for_avoid_list(self, registry):
|
|
"""Avoid certain models — prefer a non-avoided model."""
|
|
selector = ModelSelector(registry)
|
|
model = selector.select({"capabilities": {"chat"}, "avoid": {"fast-model"}})
|
|
assert model.name != "fast-model"
|