diff --git a/docs/backlog/01_execution_backlog.md b/docs/backlog/01_execution_backlog.md index 0f0fb0a..ddd126c 100644 --- a/docs/backlog/01_execution_backlog.md +++ b/docs/backlog/01_execution_backlog.md @@ -36,8 +36,8 @@ ## Package D: Integrations -- [ ] model router providers -- [ ] MCP transport adapter +- [x] model router providers +- [x] MCP transport adapter - [ ] worker WebSocket gateway - [ ] heartbeat monitor - [ ] capability-aware dispatch diff --git a/migrations/versions/0001_initial_schema.py b/migrations/versions/0001_initial_schema.py index 6f91ab4..290fc19 100644 --- a/migrations/versions/0001_initial_schema.py +++ b/migrations/versions/0001_initial_schema.py @@ -119,8 +119,48 @@ def upgrade() -> None: op.create_index("ix_task_events_task_id", "task_events", ["task_id"]) op.create_index("ix_task_events_occurred_at", "task_events", ["occurred_at"]) + op.create_table( + "model_invocations", + sa.Column("invocation_id", sa.String(length=64), primary_key=True), + sa.Column("task_id", sa.String(length=64), nullable=False), + sa.Column("node_id", sa.String(length=64), nullable=True), + sa.Column("slot", sa.String(length=64), nullable=False), + sa.Column("provider", sa.String(length=64), nullable=False), + sa.Column("model", sa.String(length=255), nullable=False), + sa.Column("status", sa.String(length=64), nullable=False), + sa.Column("request_json", sa.JSON(), nullable=False), + sa.Column("response_json", sa.JSON(), nullable=False), + sa.Column("usage_json", sa.JSON(), nullable=False), + sa.Column("fallback_from_invocation_id", sa.String(length=64), nullable=True), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + ) + op.create_index("ix_model_invocations_task_id", "model_invocations", ["task_id"]) + op.create_index("ix_model_invocations_created_at", "model_invocations", ["created_at"]) + + op.create_table( + "tool_invocations", + sa.Column("invocation_id", sa.String(length=64), primary_key=True), + sa.Column("task_id", sa.String(length=64), nullable=False), + sa.Column("node_id", sa.String(length=64), nullable=True), + sa.Column("source_type", sa.String(length=64), nullable=False), + sa.Column("source_id", sa.String(length=128), nullable=False), + sa.Column("tool_name", sa.String(length=255), nullable=False), + sa.Column("status", sa.String(length=64), nullable=False), + sa.Column("request_json", sa.JSON(), nullable=False), + sa.Column("response_json", sa.JSON(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + ) + op.create_index("ix_tool_invocations_task_id", "tool_invocations", ["task_id"]) + op.create_index("ix_tool_invocations_created_at", "tool_invocations", ["created_at"]) + def downgrade() -> None: + op.drop_index("ix_tool_invocations_created_at", table_name="tool_invocations") + op.drop_index("ix_tool_invocations_task_id", table_name="tool_invocations") + op.drop_table("tool_invocations") + op.drop_index("ix_model_invocations_created_at", table_name="model_invocations") + op.drop_index("ix_model_invocations_task_id", table_name="model_invocations") + op.drop_table("model_invocations") op.drop_index("ix_task_events_occurred_at", table_name="task_events") op.drop_index("ix_task_events_task_id", table_name="task_events") op.drop_table("task_events") diff --git a/src/ai_orchestrator/application/ports.py b/src/ai_orchestrator/application/ports.py index b690dd2..6daa10e 100644 --- a/src/ai_orchestrator/application/ports.py +++ b/src/ai_orchestrator/application/ports.py @@ -35,6 +35,34 @@ class ToolInvocationResult: error: dict[str, Any] | None = None +@dataclass(slots=True) +class ModelInvocationRecord: + invocation_id: str + task_id: str + node_id: str | None + slot: str + provider: str + model: str + status: str + request: dict[str, Any] + response: dict[str, Any] + usage: dict[str, Any] = field(default_factory=dict) + fallback_from_invocation_id: str | None = None + + +@dataclass(slots=True) +class ToolInvocationRecord: + invocation_id: str + task_id: str + node_id: str | None + source_type: str + source_id: str + tool_name: str + status: str + request: dict[str, Any] + response: dict[str, Any] + + class TaskRepository(Protocol): def create(self, task: Task) -> Task: ... def save(self, task: Task) -> Task: ... @@ -63,6 +91,13 @@ class EventStore(Protocol): def list_by_task(self, task_id: str) -> list[DomainEvent]: ... +class InvocationStore(Protocol): + def save_model_invocation(self, record: ModelInvocationRecord) -> ModelInvocationRecord: ... + def save_tool_invocation(self, record: ToolInvocationRecord) -> ToolInvocationRecord: ... + def list_model_invocations(self, task_id: str) -> list[ModelInvocationRecord]: ... + def list_tool_invocations(self, task_id: str) -> list[ToolInvocationRecord]: ... + + class PolicyEvaluator(Protocol): def evaluate(self, action: ActionDescriptor, project_id: str) -> PolicyDecision: ... diff --git a/src/ai_orchestrator/application/services/router.py b/src/ai_orchestrator/application/services/router.py new file mode 100644 index 0000000..4579581 --- /dev/null +++ b/src/ai_orchestrator/application/services/router.py @@ -0,0 +1,236 @@ +from __future__ import annotations + +from dataclasses import dataclass +from uuid import uuid4 + +from ai_orchestrator.application.ports import ( + InvocationStore, + ModelInvocationRecord, + ModelInvocationResult, +) +from ai_orchestrator.config import ProjectConfig, ProviderType + + +@dataclass(slots=True) +class ProviderResponse: + status: str + message: dict[str, object] + tool_calls: list[dict[str, object]] + usage: dict[str, object] + error: dict[str, object] | None = None + + +class ModelProvider: + provider_name: str + + def generate( + self, + *, + model: str, + base_url: str | None, + messages: list[dict[str, object]], + task_context: dict[str, object], + ) -> ProviderResponse: + raise NotImplementedError + + +@dataclass(slots=True) +class ConfigurableModelRouter: + project_configs: dict[str, ProjectConfig] + providers: dict[str, ModelProvider] + invocation_store: InvocationStore + + def run( + self, + *, + project_id: str, + slot: str, + messages: list[dict[str, object]], + task_context: dict[str, object], + ) -> ModelInvocationResult: + project = self.project_configs[project_id] + if project.models is None: + raise ValueError(f"Project has no model config: {project_id}") + + if slot == "weak": + weak_result, weak_invocation_id = self._invoke_slot( + project=project, + slot_name="weak", + slot_config=project.models.weak, + messages=messages, + task_context=task_context, + ) + if weak_result.status == "success": + return weak_result + strong_config = project.models.strong + strong_disabled = strong_config.provider == ProviderType.DISABLED + strong_external_blocked = ( + strong_config.provider == ProviderType.EXTERNAL + and not project.execution.allow_external_models + ) + if strong_disabled or strong_external_blocked: + return ModelInvocationResult( + status="degraded", + provider=weak_result.provider, + model=weak_result.model, + message=weak_result.message, + tool_calls=weak_result.tool_calls, + usage=weak_result.usage, + error=weak_result.error + or {"code": "QUALITY_DEGRADED", "message": "Strong fallback unavailable"}, + ) + strong_result, _ = self._invoke_slot( + project=project, + slot_name="strong", + slot_config=strong_config, + messages=messages, + task_context=task_context, + fallback_from_invocation_id=weak_invocation_id, + ) + return strong_result + + slot_config = getattr(project.models, slot) + result, _ = self._invoke_slot( + project=project, + slot_name=slot, + slot_config=slot_config, + messages=messages, + task_context=task_context, + ) + return result + + def _invoke_slot( + self, + *, + project: ProjectConfig, + slot_name: str, + slot_config, + messages: list[dict[str, object]], + task_context: dict[str, object], + fallback_from_invocation_id: str | None = None, + ) -> tuple[ModelInvocationResult, str]: + if slot_config.provider == ProviderType.DISABLED: + result = ModelInvocationResult( + status="disabled", + provider=slot_config.provider.value, + model=slot_config.model or "disabled", + message={}, + tool_calls=[], + usage={}, + error={"code": "PROVIDER_DISABLED", "message": f"Slot {slot_name} is disabled"}, + ) + invocation_id = self._save_invocation( + project_id=project.project.id, + task_context=task_context, + slot_name=slot_name, + provider=result.provider, + model=result.model, + status=result.status, + request={"messages": messages}, + response={"message": result.message, "error": result.error}, + usage=result.usage, + fallback_from_invocation_id=fallback_from_invocation_id, + ) + return result, invocation_id + + provider = self.providers[slot_config.provider.value] + response = provider.generate( + model=slot_config.model or "", + base_url=slot_config.base_url, + messages=messages, + task_context=task_context, + ) + status = response.status + if status == "success" and not self._is_valid_response(response.message): + status = "invalid" + result = ModelInvocationResult( + status=status, + provider=slot_config.provider.value, + model=slot_config.model or "unknown", + message=response.message, + tool_calls=response.tool_calls, + usage=response.usage, + error=response.error, + ) + invocation_id = self._save_invocation( + project_id=project.project.id, + task_context=task_context, + slot_name=slot_name, + provider=result.provider, + model=result.model, + status=result.status, + request={"messages": messages}, + response={"message": result.message, "error": result.error}, + usage=result.usage, + fallback_from_invocation_id=fallback_from_invocation_id, + ) + if slot_name == "weak" and result.status != "success": + retry_response = provider.generate( + model=slot_config.model or "", + base_url=slot_config.base_url, + messages=messages, + task_context=task_context, + ) + retry_status = retry_response.status + if retry_status == "success" and not self._is_valid_response(retry_response.message): + retry_status = "invalid" + retry_result = ModelInvocationResult( + status=retry_status, + provider=slot_config.provider.value, + model=slot_config.model or "unknown", + message=retry_response.message, + tool_calls=retry_response.tool_calls, + usage=retry_response.usage, + error=retry_response.error, + ) + retry_invocation_id = self._save_invocation( + project_id=project.project.id, + task_context=task_context, + slot_name=slot_name, + provider=retry_result.provider, + model=retry_result.model, + status=retry_result.status, + request={"messages": messages, "retry": True}, + response={"message": retry_result.message, "error": retry_result.error}, + usage=retry_result.usage, + fallback_from_invocation_id=invocation_id, + ) + return retry_result, retry_invocation_id + return result, invocation_id + + @staticmethod + def _is_valid_response(message: dict[str, object]) -> bool: + content = message.get("content") + return isinstance(content, str) and bool(content.strip()) + + def _save_invocation( + self, + *, + project_id: str, + task_context: dict[str, object], + slot_name: str, + provider: str, + model: str, + status: str, + request: dict[str, object], + response: dict[str, object], + usage: dict[str, object], + fallback_from_invocation_id: str | None, + ) -> str: + invocation_id = f"minv_{uuid4().hex}" + self.invocation_store.save_model_invocation( + ModelInvocationRecord( + invocation_id=invocation_id, + task_id=str(task_context.get("task_id", project_id)), + node_id=str(task_context["node_id"]) if task_context.get("node_id") else None, + slot=slot_name, + provider=provider, + model=model, + status=status, + request=request, + response=response, + usage=usage, + fallback_from_invocation_id=fallback_from_invocation_id, + ) + ) + return invocation_id diff --git a/src/ai_orchestrator/config.py b/src/ai_orchestrator/config.py index 5314f23..d8d6704 100644 --- a/src/ai_orchestrator/config.py +++ b/src/ai_orchestrator/config.py @@ -59,6 +59,13 @@ class PolicyConfig(BaseModel): resources: PolicyResourceConfig = Field(default_factory=PolicyResourceConfig) +class McpServerConfig(BaseModel): + transport: str = "http" + base_url: str + api_key_env: str | None = None + timeout_ms: int = 30000 + + class ProjectMetadata(BaseModel): id: str name: str @@ -67,6 +74,7 @@ class ProjectMetadata(BaseModel): class ProjectConfig(BaseModel): project: ProjectMetadata models: ModelsConfig | None = None + mcp_servers: dict[str, McpServerConfig] = Field(default_factory=dict) execution: ExecutionConfig = Field(default_factory=ExecutionConfig) policy: PolicyConfig = Field(default_factory=PolicyConfig) diff --git a/src/ai_orchestrator/delivery/http/app.py b/src/ai_orchestrator/delivery/http/app.py index 3526106..0c90fbf 100644 --- a/src/ai_orchestrator/delivery/http/app.py +++ b/src/ai_orchestrator/delivery/http/app.py @@ -13,6 +13,7 @@ from ai_orchestrator.application.services.execution import ( ) from ai_orchestrator.application.services.lifecycle import TaskLifecycleService from ai_orchestrator.application.services.orchestrator import CreateTaskRequest, OrchestratorService +from ai_orchestrator.application.services.router import ConfigurableModelRouter from ai_orchestrator.application.services.workers import RegisterWorkerRequest, WorkerService from ai_orchestrator.config import AppSettings from ai_orchestrator.delivery.http.schemas import ( @@ -28,6 +29,7 @@ from ai_orchestrator.delivery.http.schemas import ( WorkerRegisterRequest, ) from ai_orchestrator.infrastructure.config_loader import load_project_config +from ai_orchestrator.infrastructure.model_router import StaticMockModelProvider from ai_orchestrator.infrastructure.policy import StaticProjectPolicyEvaluator from ai_orchestrator.infrastructure.storage.factory import create_storage_bundle @@ -41,9 +43,18 @@ def create_app(settings: AppSettings | None = None) -> FastAPI: confirmation_repository = storage.confirmation_repository worker_repository = storage.worker_repository event_store = storage.event_store + invocation_store = storage.invocation_store policy_evaluator = StaticProjectPolicyEvaluator( projects={project_config.project.id: project_config} ) + model_router = ConfigurableModelRouter( + project_configs={project_config.project.id: project_config}, + providers={ + "local": StaticMockModelProvider(provider_name="local"), + "external": StaticMockModelProvider(provider_name="external"), + }, + invocation_store=invocation_store, + ) runner_registry = RunnerRegistry() runner_registry.register(PlannerNodeRunner()) runner_registry.register(FinalizerNodeRunner(finalizer=DefaultFinalizer())) @@ -73,6 +84,8 @@ def create_app(settings: AppSettings | None = None) -> FastAPI: app.state.worker_repository = worker_repository app.state.event_store = event_store app.state.storage = storage + app.state.invocation_store = invocation_store + app.state.model_router = model_router app.state.orchestrator = orchestrator app.state.execution_engine = execution_engine app.state.lifecycle_service = lifecycle_service diff --git a/src/ai_orchestrator/infrastructure/mcp_client.py b/src/ai_orchestrator/infrastructure/mcp_client.py new file mode 100644 index 0000000..a25c033 --- /dev/null +++ b/src/ai_orchestrator/infrastructure/mcp_client.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +import os +from dataclasses import dataclass +from uuid import uuid4 + +import httpx + +from ai_orchestrator.application.ports import ( + InvocationStore, + ToolGateway, + ToolInvocationRecord, + ToolInvocationResult, +) +from ai_orchestrator.config import ProjectConfig + + +@dataclass(slots=True) +class McpHttpClient: + transport: httpx.BaseTransport | None = None + + def call( + self, + *, + base_url: str, + method: str, + params: dict[str, object], + timeout_ms: int, + api_key: str | None = None, + ) -> dict[str, object]: + headers: dict[str, str] = {} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + with httpx.Client(timeout=timeout_ms / 1000, transport=self.transport) as client: + response = client.post( + f"{base_url.rstrip('/')}/rpc", + json={ + "jsonrpc": "2.0", + "id": uuid4().hex, + "method": method, + "params": params, + }, + headers=headers, + ) + response.raise_for_status() + payload = response.json() + if "error" in payload: + return { + "status": "error", + "error": payload["error"], + "content": {}, + "artifacts": [], + "logs": [], + } + return payload.get("result", {}) + + +@dataclass(slots=True) +class McpToolGateway(ToolGateway): + project_configs: dict[str, ProjectConfig] + client: McpHttpClient + invocation_store: InvocationStore + + def call( + self, + *, + project_id: str, + server_id: str, + tool_name: str, + args: dict[str, object], + task_context: dict[str, object], + ) -> ToolInvocationResult: + project = self.project_configs[project_id] + server = project.mcp_servers[server_id] + api_key = os.getenv(server.api_key_env) if server.api_key_env else None + payload = self.client.call( + base_url=server.base_url, + method="tools/call", + params={"server_id": server_id, "tool": tool_name, "args": args}, + timeout_ms=server.timeout_ms, + api_key=api_key, + ) + result = ToolInvocationResult( + status=str(payload.get("status", "success")), + content=dict(payload.get("content", {})), + artifacts=list(payload.get("artifacts", [])), + logs=[str(item) for item in payload.get("logs", [])], + error=payload.get("error"), + ) + self.invocation_store.save_tool_invocation( + ToolInvocationRecord( + invocation_id=f"tinv_{uuid4().hex}", + task_id=str(task_context.get("task_id", project_id)), + node_id=str(task_context["node_id"]) if task_context.get("node_id") else None, + source_type="mcp", + source_id=server_id, + tool_name=tool_name, + status=result.status, + request={"args": args}, + response={ + "content": result.content, + "artifacts": result.artifacts, + "logs": result.logs, + "error": result.error, + }, + ) + ) + return result diff --git a/src/ai_orchestrator/infrastructure/model_router.py b/src/ai_orchestrator/infrastructure/model_router.py new file mode 100644 index 0000000..55f41ad --- /dev/null +++ b/src/ai_orchestrator/infrastructure/model_router.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +import os +from dataclasses import dataclass, field + +import httpx + +from ai_orchestrator.application.services.router import ModelProvider, ProviderResponse + + +@dataclass(slots=True) +class StaticMockModelProvider(ModelProvider): + provider_name: str + responses: list[ProviderResponse] = field(default_factory=list) + + def generate( + self, + *, + model: str, + base_url: str | None, + messages: list[dict[str, object]], + task_context: dict[str, object], + ) -> ProviderResponse: + del model, base_url, messages, task_context + if self.responses: + return self.responses.pop(0) + return ProviderResponse( + status="success", + message={"role": "assistant", "content": "mock response"}, + tool_calls=[], + usage={"input_tokens": 0, "output_tokens": 0, "cost": 0}, + ) + + +@dataclass(slots=True) +class OpenAICompatibleModelProvider(ModelProvider): + provider_name: str + api_key_env: str | None = None + timeout_s: float = 30.0 + transport: httpx.BaseTransport | None = None + + def generate( + self, + *, + model: str, + base_url: str | None, + messages: list[dict[str, object]], + task_context: dict[str, object], + ) -> ProviderResponse: + del task_context + if not base_url: + return ProviderResponse( + status="error", + message={}, + tool_calls=[], + usage={}, + error={"code": "MISSING_BASE_URL", "message": "Provider base_url is required"}, + ) + headers: dict[str, str] = {} + if self.api_key_env and os.getenv(self.api_key_env): + headers["Authorization"] = f"Bearer {os.getenv(self.api_key_env)}" + with httpx.Client(timeout=self.timeout_s, transport=self.transport) as client: + response = client.post( + f"{base_url.rstrip('/')}/chat/completions", + json={"model": model, "messages": messages}, + headers=headers, + ) + response.raise_for_status() + payload = response.json() + choice = payload.get("choices", [{}])[0] + message = choice.get("message", {}) + usage = payload.get("usage", {}) + return ProviderResponse( + status="success", + message=message, + tool_calls=message.get("tool_calls", []) or [], + usage=usage, + ) diff --git a/src/ai_orchestrator/infrastructure/storage/factory.py b/src/ai_orchestrator/infrastructure/storage/factory.py index b44bb6e..28e5e52 100644 --- a/src/ai_orchestrator/infrastructure/storage/factory.py +++ b/src/ai_orchestrator/infrastructure/storage/factory.py @@ -7,6 +7,7 @@ from ai_orchestrator.infrastructure.storage.memory import ( InMemoryConfirmationRepository, InMemoryEventStore, InMemoryGraphRepository, + InMemoryInvocationStore, InMemoryTaskRepository, InMemoryWorkerRepository, ) @@ -20,6 +21,7 @@ class StorageBundle: confirmation_repository: object worker_repository: object event_store: object + invocation_store: object engine: object | None = None session_factory: object | None = None @@ -32,16 +34,18 @@ def create_storage_bundle(settings: AppSettings) -> StorageBundle: confirmation_repository=InMemoryConfirmationRepository(), worker_repository=InMemoryWorkerRepository(), event_store=InMemoryEventStore(), + invocation_store=InMemoryInvocationStore(), ) if settings.storage_backend in {"sqlite", "postgres"}: bundle = create_sqlalchemy_storage(settings.database_url) return StorageBundle( - engine=bundle.engine, - session_factory=bundle.session_factory, task_repository=bundle.task_repository, graph_repository=bundle.graph_repository, confirmation_repository=bundle.confirmation_repository, worker_repository=bundle.worker_repository, event_store=bundle.event_store, + invocation_store=bundle.invocation_store, + engine=bundle.engine, + session_factory=bundle.session_factory, ) raise ValueError(f"Unsupported storage backend: {settings.storage_backend}") diff --git a/src/ai_orchestrator/infrastructure/storage/memory.py b/src/ai_orchestrator/infrastructure/storage/memory.py index 75de0e8..e09d744 100644 --- a/src/ai_orchestrator/infrastructure/storage/memory.py +++ b/src/ai_orchestrator/infrastructure/storage/memory.py @@ -6,7 +6,10 @@ from ai_orchestrator.application.ports import ( ConfirmationRepository, EventStore, GraphRepository, + InvocationStore, + ModelInvocationRecord, TaskRepository, + ToolInvocationRecord, WorkerRepository, ) from ai_orchestrator.domain.events import DomainEvent @@ -83,3 +86,22 @@ class InMemoryEventStore(EventStore): def list_by_task(self, task_id: str) -> list[DomainEvent]: return [event for event in self.items if event.task_id == task_id] + +@dataclass(slots=True) +class InMemoryInvocationStore(InvocationStore): + model_items: list[ModelInvocationRecord] = field(default_factory=list) + tool_items: list[ToolInvocationRecord] = field(default_factory=list) + + def save_model_invocation(self, record: ModelInvocationRecord) -> ModelInvocationRecord: + self.model_items.append(record) + return record + + def save_tool_invocation(self, record: ToolInvocationRecord) -> ToolInvocationRecord: + self.tool_items.append(record) + return record + + def list_model_invocations(self, task_id: str) -> list[ModelInvocationRecord]: + return [record for record in self.model_items if record.task_id == task_id] + + def list_tool_invocations(self, task_id: str) -> list[ToolInvocationRecord]: + return [record for record in self.tool_items if record.task_id == task_id] diff --git a/src/ai_orchestrator/infrastructure/storage/sqlalchemy.py b/src/ai_orchestrator/infrastructure/storage/sqlalchemy.py index 183cabf..991e554 100644 --- a/src/ai_orchestrator/infrastructure/storage/sqlalchemy.py +++ b/src/ai_orchestrator/infrastructure/storage/sqlalchemy.py @@ -11,7 +11,10 @@ from ai_orchestrator.application.ports import ( ConfirmationRepository, EventStore, GraphRepository, + InvocationStore, + ModelInvocationRecord, TaskRepository, + ToolInvocationRecord, WorkerRepository, ) from ai_orchestrator.domain.enums import ( @@ -128,6 +131,38 @@ class TaskEventRecord(Base): occurred_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), index=True) +class ModelInvocationDbRecord(Base): + __tablename__ = "model_invocations" + + invocation_id: Mapped[str] = mapped_column(String(64), primary_key=True) + task_id: Mapped[str] = mapped_column(String(64), index=True) + node_id: Mapped[str | None] = mapped_column(String(64), nullable=True) + slot: Mapped[str] = mapped_column(String(64)) + provider: Mapped[str] = mapped_column(String(64)) + model: Mapped[str] = mapped_column(String(255)) + status: Mapped[str] = mapped_column(String(64)) + request_json: Mapped[dict] = mapped_column(JSON) + response_json: Mapped[dict] = mapped_column(JSON) + usage_json: Mapped[dict] = mapped_column(JSON) + fallback_from_invocation_id: Mapped[str | None] = mapped_column(String(64), nullable=True) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), index=True) + + +class ToolInvocationDbRecord(Base): + __tablename__ = "tool_invocations" + + invocation_id: Mapped[str] = mapped_column(String(64), primary_key=True) + task_id: Mapped[str] = mapped_column(String(64), index=True) + node_id: Mapped[str | None] = mapped_column(String(64), nullable=True) + source_type: Mapped[str] = mapped_column(String(64)) + source_id: Mapped[str] = mapped_column(String(128)) + tool_name: Mapped[str] = mapped_column(String(255)) + status: Mapped[str] = mapped_column(String(64)) + request_json: Mapped[dict] = mapped_column(JSON) + response_json: Mapped[dict] = mapped_column(JSON) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), index=True) + + def build_engine(database_url: str): if database_url.startswith("sqlite:///"): db_path = database_url.removeprefix("sqlite:///") @@ -319,6 +354,68 @@ def _record_to_event(record: TaskEventRecord) -> DomainEvent: ) +def _model_invocation_to_record(record: ModelInvocationRecord) -> ModelInvocationDbRecord: + return ModelInvocationDbRecord( + invocation_id=record.invocation_id, + task_id=record.task_id, + node_id=record.node_id, + slot=record.slot, + provider=record.provider, + model=record.model, + status=record.status, + request_json=record.request, + response_json=record.response, + usage_json=record.usage, + fallback_from_invocation_id=record.fallback_from_invocation_id, + created_at=datetime.now(UTC), + ) + + +def _record_to_model_invocation(record: ModelInvocationDbRecord) -> ModelInvocationRecord: + return ModelInvocationRecord( + invocation_id=record.invocation_id, + task_id=record.task_id, + node_id=record.node_id, + slot=record.slot, + provider=record.provider, + model=record.model, + status=record.status, + request=record.request_json, + response=record.response_json, + usage=record.usage_json or {}, + fallback_from_invocation_id=record.fallback_from_invocation_id, + ) + + +def _tool_invocation_to_record(record: ToolInvocationRecord) -> ToolInvocationDbRecord: + return ToolInvocationDbRecord( + invocation_id=record.invocation_id, + task_id=record.task_id, + node_id=record.node_id, + source_type=record.source_type, + source_id=record.source_id, + tool_name=record.tool_name, + status=record.status, + request_json=record.request, + response_json=record.response, + created_at=datetime.now(UTC), + ) + + +def _record_to_tool_invocation(record: ToolInvocationDbRecord) -> ToolInvocationRecord: + return ToolInvocationRecord( + invocation_id=record.invocation_id, + task_id=record.task_id, + node_id=record.node_id, + source_type=record.source_type, + source_id=record.source_id, + tool_name=record.tool_name, + status=record.status, + request=record.request_json, + response=record.response_json, + ) + + @dataclass(slots=True) class SqlAlchemyTaskRepository(TaskRepository): session_factory: sessionmaker[Session] @@ -492,6 +589,41 @@ class SqlAlchemyEventStore(EventStore): return [_record_to_event(record) for record in records] +@dataclass(slots=True) +class SqlAlchemyInvocationStore(InvocationStore): + session_factory: sessionmaker[Session] + + def save_model_invocation(self, record: ModelInvocationRecord) -> ModelInvocationRecord: + with self.session_factory() as session: + session.add(_model_invocation_to_record(record)) + session.commit() + return record + + def save_tool_invocation(self, record: ToolInvocationRecord) -> ToolInvocationRecord: + with self.session_factory() as session: + session.add(_tool_invocation_to_record(record)) + session.commit() + return record + + def list_model_invocations(self, task_id: str) -> list[ModelInvocationRecord]: + with self.session_factory() as session: + records = session.scalars( + select(ModelInvocationDbRecord) + .where(ModelInvocationDbRecord.task_id == task_id) + .order_by(ModelInvocationDbRecord.created_at.asc()) + ).all() + return [_record_to_model_invocation(record) for record in records] + + def list_tool_invocations(self, task_id: str) -> list[ToolInvocationRecord]: + with self.session_factory() as session: + records = session.scalars( + select(ToolInvocationDbRecord) + .where(ToolInvocationDbRecord.task_id == task_id) + .order_by(ToolInvocationDbRecord.created_at.asc()) + ).all() + return [_record_to_tool_invocation(record) for record in records] + + @dataclass(slots=True) class SqlAlchemyStorageBundle: engine: object @@ -501,6 +633,7 @@ class SqlAlchemyStorageBundle: confirmation_repository: SqlAlchemyConfirmationRepository worker_repository: SqlAlchemyWorkerRepository event_store: SqlAlchemyEventStore + invocation_store: SqlAlchemyInvocationStore def create_sqlalchemy_storage(database_url: str) -> SqlAlchemyStorageBundle: @@ -514,4 +647,5 @@ def create_sqlalchemy_storage(database_url: str) -> SqlAlchemyStorageBundle: confirmation_repository=SqlAlchemyConfirmationRepository(session_factory), worker_repository=SqlAlchemyWorkerRepository(session_factory), event_store=SqlAlchemyEventStore(session_factory), + invocation_store=SqlAlchemyInvocationStore(session_factory), ) diff --git a/tests/integration/test_sqlalchemy_storage.py b/tests/integration/test_sqlalchemy_storage.py index 618eaeb..2956535 100644 --- a/tests/integration/test_sqlalchemy_storage.py +++ b/tests/integration/test_sqlalchemy_storage.py @@ -1,5 +1,6 @@ from pathlib import Path +from ai_orchestrator.application.ports import ModelInvocationRecord, ToolInvocationRecord from ai_orchestrator.config import AppSettings from ai_orchestrator.domain.enums import NodeType, PolicyDecisionType from ai_orchestrator.domain.events import DomainEvent @@ -77,3 +78,43 @@ def test_sqlite_storage_persists_confirmation_and_worker(tmp_path: Path) -> None assert storage.confirmation_repository.get(confirmation.confirmation_id) is not None assert storage.worker_repository.get(worker.session_id) is not None assert len(storage.worker_repository.list_active()) == 1 + + +def test_sqlite_storage_persists_invocation_records(tmp_path: Path) -> None: + db_path = tmp_path / "orchestrator.sqlite3" + settings = AppSettings( + storage_backend="sqlite", + database_url=f"sqlite:///{db_path.as_posix()}", + ) + storage = create_storage_bundle(settings) + + storage.invocation_store.save_model_invocation( + ModelInvocationRecord( + invocation_id="minv_1", + task_id="task_1", + node_id="node_1", + slot="weak", + provider="local", + model="mock", + status="success", + request={"messages": []}, + response={"message": {"content": "ok"}}, + usage={"input_tokens": 1}, + ) + ) + storage.invocation_store.save_tool_invocation( + ToolInvocationRecord( + invocation_id="tinv_1", + task_id="task_1", + node_id="node_2", + source_type="mcp", + source_id="one_c", + tool_name="one_c.run_sql", + status="success", + request={"args": {"query": "select 1"}}, + response={"content": {"rows": []}}, + ) + ) + + assert len(storage.invocation_store.list_model_invocations("task_1")) == 1 + assert len(storage.invocation_store.list_tool_invocations("task_1")) == 1 diff --git a/tests/unit/test_mcp_client.py b/tests/unit/test_mcp_client.py new file mode 100644 index 0000000..6bf9e62 --- /dev/null +++ b/tests/unit/test_mcp_client.py @@ -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 diff --git a/tests/unit/test_model_router.py b/tests/unit/test_model_router.py new file mode 100644 index 0000000..ee4ca35 --- /dev/null +++ b/tests/unit/test_model_router.py @@ -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