60 lines
1.7 KiB
Python
60 lines
1.7 KiB
Python
import pytest
|
|
|
|
from app.integrations.model_router import ModelRouter
|
|
|
|
|
|
def test_auto_route_requires_orchestrator_choice_with_multiple_candidates() -> None:
|
|
pool = [
|
|
{
|
|
"id": "image-a",
|
|
"name": "Image A",
|
|
"model_id": "image-a-model",
|
|
"category": "multimodal",
|
|
"provider": "provider-a",
|
|
"enabled": True,
|
|
"capabilities": ["image_generation"],
|
|
},
|
|
{
|
|
"id": "image-b",
|
|
"name": "Image B",
|
|
"model_id": "image-b-model",
|
|
"category": "multimodal",
|
|
"provider": "provider-b",
|
|
"enabled": True,
|
|
"capabilities": ["image_generation"],
|
|
},
|
|
]
|
|
tests = {
|
|
"model:image-a": {"ok": True},
|
|
"model:image-b": {"ok": True},
|
|
}
|
|
router = ModelRouter({"model_pool": pool, "image_routing_mode": "auto"}, tests)
|
|
|
|
with pytest.raises(RuntimeError, match="总调度返回模型实例 ID"):
|
|
router.choose("image")
|
|
assert router.choose("image", "image-b")["model_id"] == "image-b-model"
|
|
|
|
|
|
def test_manual_route_uses_the_user_selected_model() -> None:
|
|
pool = [
|
|
{
|
|
"id": "spatial",
|
|
"name": "Spatial",
|
|
"model_id": "spatial-model",
|
|
"category": "multimodal",
|
|
"provider": "custom",
|
|
"enabled": True,
|
|
"capabilities": ["spatial_understanding"],
|
|
}
|
|
]
|
|
router = ModelRouter(
|
|
{
|
|
"model_pool": pool,
|
|
"spatial_routing_mode": "manual",
|
|
"spatial_model_id": "spatial",
|
|
},
|
|
{"model:spatial": {"ok": True}},
|
|
)
|
|
|
|
assert router.choose("spatial")["id"] == "spatial"
|