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,63 @@
+import httpx
+
+from ai_orchestrator.config import (
+    ExecutionConfig,
+    McpServerConfig,
+    ModelsConfig,
+    ModelSlotConfig,
+    PolicyConfig,
+    ProjectConfig,
+    ProjectMetadata,
+)
+from ai_orchestrator.infrastructure.mcp_client import McpHttpClient, McpToolGateway
+from ai_orchestrator.infrastructure.storage.memory import InMemoryInvocationStore
+
+
+def test_mcp_tool_gateway_normalizes_result_and_persists_invocation() -> None:
+    def handler(request: httpx.Request) -> httpx.Response:
+        assert request.url.path == "/rpc"
+        payload = {
+            "jsonrpc": "2.0",
+            "id": "1",
+            "result": {
+                "status": "success",
+                "content": {"rows": [{"id": 1}]},
+                "artifacts": [],
+                "logs": ["done"],
+                "error": None,
+            },
+        }
+        return httpx.Response(200, json=payload)
+
+    transport = httpx.MockTransport(handler)
+    project = ProjectConfig(
+        project=ProjectMetadata(id="default", name="Default"),
+        models=ModelsConfig(
+            weak=ModelSlotConfig(provider="local", model="weak"),
+            strong=ModelSlotConfig(provider="local", model="strong"),
+            vision=ModelSlotConfig(provider="disabled"),
+            embedding=ModelSlotConfig(provider="local", model="embed"),
+        ),
+        mcp_servers={"one_c": McpServerConfig(base_url="http://mcp.test")},
+        execution=ExecutionConfig(),
+        policy=PolicyConfig(),
+    )
+    invocation_store = InMemoryInvocationStore()
+    gateway = McpToolGateway(
+        project_configs={"default": project},
+        client=McpHttpClient(transport=transport),
+        invocation_store=invocation_store,
+    )
+
+    result = gateway.call(
+        project_id="default",
+        server_id="one_c",
+        tool_name="one_c.run_sql",
+        args={"query": "select 1"},
+        task_context={"task_id": "task_1", "node_id": "node_2"},
+    )
+
+    assert result.status == "success"
+    assert result.content["rows"][0]["id"] == 1
+    assert result.logs == ["done"]
+    assert len(invocation_store.list_tool_invocations("task_1")) == 1
@@ -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