78 lines
2.7 KiB
Python
78 lines
2.7 KiB
Python
from ai_orchestrator.application.services.router import (
|
|
ConfigurableModelRouter,
|
|
ProviderResponse,
|
|
)
|
|
from ai_orchestrator.config import (
|
|
ExecutionConfig,
|
|
McpServerConfig,
|
|
ModelsConfig,
|
|
ModelSlotConfig,
|
|
PolicyConfig,
|
|
ProjectConfig,
|
|
ProjectMetadata,
|
|
)
|
|
from ai_orchestrator.infrastructure.model_router import StaticMockModelProvider
|
|
from ai_orchestrator.infrastructure.storage.memory import InMemoryInvocationStore
|
|
|
|
|
|
def test_model_router_retries_weak_and_falls_back_to_strong() -> None:
|
|
config = ProjectConfig(
|
|
project=ProjectMetadata(id="default", name="Default"),
|
|
models=ModelsConfig(
|
|
weak=ModelSlotConfig(provider="local", model="weak"),
|
|
strong=ModelSlotConfig(provider="external", model="strong"),
|
|
vision=ModelSlotConfig(provider="disabled"),
|
|
embedding=ModelSlotConfig(provider="local", model="embed"),
|
|
),
|
|
mcp_servers={"test": McpServerConfig(base_url="http://example.test")},
|
|
execution=ExecutionConfig(allow_external_models=True),
|
|
policy=PolicyConfig(),
|
|
)
|
|
invocation_store = InMemoryInvocationStore()
|
|
router = ConfigurableModelRouter(
|
|
project_configs={"default": config},
|
|
providers={
|
|
"local": StaticMockModelProvider(
|
|
provider_name="local",
|
|
responses=[
|
|
ProviderResponse(
|
|
status="success",
|
|
message={"role": "assistant", "content": ""},
|
|
tool_calls=[],
|
|
usage={},
|
|
),
|
|
ProviderResponse(
|
|
status="success",
|
|
message={"role": "assistant", "content": ""},
|
|
tool_calls=[],
|
|
usage={},
|
|
),
|
|
],
|
|
),
|
|
"external": StaticMockModelProvider(
|
|
provider_name="external",
|
|
responses=[
|
|
ProviderResponse(
|
|
status="success",
|
|
message={"role": "assistant", "content": "strong answer"},
|
|
tool_calls=[],
|
|
usage={"input_tokens": 1, "output_tokens": 2},
|
|
)
|
|
],
|
|
),
|
|
},
|
|
invocation_store=invocation_store,
|
|
)
|
|
|
|
result = router.run(
|
|
project_id="default",
|
|
slot="weak",
|
|
messages=[{"role": "user", "content": "Hello"}],
|
|
task_context={"task_id": "task_1", "node_id": "node_1"},
|
|
)
|
|
|
|
assert result.status == "success"
|
|
assert result.model == "strong"
|
|
assert result.message["content"] == "strong answer"
|
|
assert len(invocation_store.list_model_invocations("task_1")) == 3
|