初始化提交

backend/tests/ai_provider/routing/test_selector.py

@@ -0,0 +1,135 @@
+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"