Add model router and MCP integration adapters

2026-07-03 21:25:48 +03:00
parent 4db700bd4d
commit b681a90eae
14 changed files with 863 additions and 4 deletions
@@ -0,0 +1,77 @@
+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