322 lines
11 KiB
Python
322 lines
11 KiB
Python
import httpx
|
|
|
|
from ai_orchestrator.application.services.execution import (
|
|
DefaultFinalizer,
|
|
FinalizerNodeRunner,
|
|
GraphExecutionEngine,
|
|
NoOpReviewer,
|
|
PlannerNodeRunner,
|
|
RunnerRegistry,
|
|
)
|
|
from ai_orchestrator.application.services.orchestrator import CreateTaskRequest, OrchestratorService
|
|
from ai_orchestrator.application.services.router import ConfigurableModelRouter, ProviderResponse
|
|
from ai_orchestrator.application.services.workers import (
|
|
CapabilityAwareWorkerGateway,
|
|
InMemoryWorkerConnectionManager,
|
|
RegisterWorkerRequest,
|
|
WorkerCommandResult,
|
|
WorkerService,
|
|
)
|
|
from ai_orchestrator.config import (
|
|
ExecutionConfig,
|
|
McpServerConfig,
|
|
ModelsConfig,
|
|
ModelSlotConfig,
|
|
PolicyConfig,
|
|
ProjectConfig,
|
|
ProjectMetadata,
|
|
)
|
|
from ai_orchestrator.domain.enums import RiskLevel
|
|
from ai_orchestrator.domain.models import ActionDescriptor
|
|
from ai_orchestrator.infrastructure.mcp_client import McpHttpClient, McpToolGateway
|
|
from ai_orchestrator.infrastructure.model_router import StaticMockModelProvider
|
|
from ai_orchestrator.infrastructure.policy import StaticProjectPolicyEvaluator
|
|
from ai_orchestrator.infrastructure.storage.memory import (
|
|
InMemoryConfirmationRepository,
|
|
InMemoryEventStore,
|
|
InMemoryGraphRepository,
|
|
InMemoryInvocationStore,
|
|
InMemoryTaskRepository,
|
|
InMemoryWorkerRepository,
|
|
)
|
|
|
|
|
|
def _project_config(*, allow_external: bool = True, default_mode: str = "confirm") -> ProjectConfig:
|
|
return ProjectConfig(
|
|
project=ProjectMetadata(id="default", name="Default"),
|
|
models=ModelsConfig(
|
|
weak=ModelSlotConfig(provider="local", model="weak"),
|
|
strong=ModelSlotConfig(
|
|
provider="external" if allow_external else "disabled",
|
|
model="strong" if allow_external else None,
|
|
),
|
|
vision=ModelSlotConfig(provider="disabled"),
|
|
embedding=ModelSlotConfig(provider="local", model="embed"),
|
|
),
|
|
mcp_servers={"one_c": McpServerConfig(base_url="http://mcp.test")},
|
|
execution=ExecutionConfig(allow_external_models=allow_external),
|
|
policy=PolicyConfig(
|
|
default_mode=default_mode,
|
|
resources={
|
|
"filesystem": default_mode,
|
|
"shell": default_mode,
|
|
"sql": default_mode,
|
|
"mcp": default_mode,
|
|
"external_models": default_mode if allow_external else "disabled",
|
|
"desktop": default_mode,
|
|
"browser": default_mode,
|
|
"network": default_mode,
|
|
"cost": default_mode,
|
|
"system": default_mode,
|
|
},
|
|
),
|
|
)
|
|
|
|
|
|
def test_smoke_simple_chat_and_graph_and_events() -> None:
|
|
task_repository = InMemoryTaskRepository()
|
|
graph_repository = InMemoryGraphRepository()
|
|
confirmation_repository = InMemoryConfirmationRepository()
|
|
event_store = InMemoryEventStore()
|
|
orchestrator = OrchestratorService(
|
|
task_repository=task_repository,
|
|
graph_repository=graph_repository,
|
|
confirmation_repository=confirmation_repository,
|
|
event_store=event_store,
|
|
policy_evaluator=StaticProjectPolicyEvaluator(projects={"default": _project_config()}),
|
|
)
|
|
registry = RunnerRegistry()
|
|
registry.register(PlannerNodeRunner())
|
|
registry.register(FinalizerNodeRunner(finalizer=DefaultFinalizer()))
|
|
engine = GraphExecutionEngine(
|
|
event_store=event_store,
|
|
runner_registry=registry,
|
|
reviewer=NoOpReviewer(),
|
|
)
|
|
|
|
task = orchestrator.create_task(
|
|
CreateTaskRequest(project_id="default", goal="hello", inputs={})
|
|
)
|
|
graph = orchestrator.plan_task(task.task_id)
|
|
engine.execute_ready_nodes(task=task, graph=graph)
|
|
|
|
assert len(graph.nodes) == 2
|
|
assert any(
|
|
event.event_type == "node_started" for event in event_store.list_by_task(task.task_id)
|
|
)
|
|
|
|
|
|
def test_smoke_model_router_retry_and_fallback_and_external_disabled() -> None:
|
|
invocation_store = InMemoryInvocationStore()
|
|
project = _project_config(allow_external=True)
|
|
router = ConfigurableModelRouter(
|
|
project_configs={"default": project},
|
|
providers={
|
|
"local": StaticMockModelProvider(
|
|
provider_name="local",
|
|
responses=[
|
|
ProviderResponse(
|
|
status="success",
|
|
message={"content": ""},
|
|
tool_calls=[],
|
|
usage={},
|
|
),
|
|
ProviderResponse(
|
|
status="success",
|
|
message={"content": ""},
|
|
tool_calls=[],
|
|
usage={},
|
|
),
|
|
],
|
|
),
|
|
"external": StaticMockModelProvider(
|
|
provider_name="external",
|
|
responses=[
|
|
ProviderResponse(
|
|
status="success",
|
|
message={"content": "strong"},
|
|
tool_calls=[],
|
|
usage={},
|
|
)
|
|
],
|
|
),
|
|
},
|
|
invocation_store=invocation_store,
|
|
)
|
|
result = router.run(
|
|
project_id="default",
|
|
slot="weak",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
task_context={"task_id": "task_1"},
|
|
)
|
|
|
|
disabled_router = ConfigurableModelRouter(
|
|
project_configs={"default": _project_config(allow_external=False)},
|
|
providers={"local": StaticMockModelProvider(provider_name="local")},
|
|
invocation_store=InMemoryInvocationStore(),
|
|
)
|
|
disabled_result = disabled_router.run(
|
|
project_id="default",
|
|
slot="weak",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
task_context={"task_id": "task_2"},
|
|
)
|
|
|
|
assert result.status == "success"
|
|
assert result.message["content"] == "strong"
|
|
assert disabled_result.status in {"success", "degraded", "disabled"}
|
|
|
|
|
|
def test_smoke_policy_confirm_and_full_auto() -> None:
|
|
confirmation_repository = InMemoryConfirmationRepository()
|
|
event_store = InMemoryEventStore()
|
|
orchestrator_confirm = OrchestratorService(
|
|
task_repository=InMemoryTaskRepository(),
|
|
graph_repository=InMemoryGraphRepository(),
|
|
confirmation_repository=confirmation_repository,
|
|
event_store=event_store,
|
|
policy_evaluator=StaticProjectPolicyEvaluator(
|
|
projects={"default": _project_config(default_mode="confirm")}
|
|
),
|
|
)
|
|
task = orchestrator_confirm.create_task(
|
|
CreateTaskRequest(project_id="default", goal="confirm", inputs={})
|
|
)
|
|
graph = orchestrator_confirm.plan_task(task.task_id)
|
|
confirmation = orchestrator_confirm.evaluate_action(
|
|
task_id=task.task_id,
|
|
node_id=graph.nodes[-1].node_id,
|
|
action=ActionDescriptor(
|
|
resource="filesystem",
|
|
risk_level=RiskLevel.WRITE,
|
|
action_name="file.write",
|
|
preview_available=True,
|
|
),
|
|
)
|
|
|
|
orchestrator_auto = OrchestratorService(
|
|
task_repository=InMemoryTaskRepository(),
|
|
graph_repository=InMemoryGraphRepository(),
|
|
confirmation_repository=InMemoryConfirmationRepository(),
|
|
event_store=InMemoryEventStore(),
|
|
policy_evaluator=StaticProjectPolicyEvaluator(
|
|
projects={"default": _project_config(default_mode="full_auto")}
|
|
),
|
|
)
|
|
auto_task = orchestrator_auto.create_task(
|
|
CreateTaskRequest(project_id="default", goal="auto", inputs={})
|
|
)
|
|
auto_graph = orchestrator_auto.plan_task(auto_task.task_id)
|
|
auto_confirmation = orchestrator_auto.evaluate_action(
|
|
task_id=auto_task.task_id,
|
|
node_id=auto_graph.nodes[-1].node_id,
|
|
action=ActionDescriptor(
|
|
resource="filesystem",
|
|
risk_level=RiskLevel.WRITE,
|
|
action_name="file.write",
|
|
),
|
|
)
|
|
|
|
assert confirmation is not None
|
|
assert auto_confirmation is None
|
|
|
|
|
|
def test_smoke_mcp_and_worker_results_are_stored_and_finalizer_runs() -> None:
|
|
invocation_store = InMemoryInvocationStore()
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"jsonrpc": "2.0",
|
|
"id": "1",
|
|
"result": {
|
|
"status": "success",
|
|
"content": {"rows": [1]},
|
|
"artifacts": [],
|
|
"logs": [],
|
|
"error": None,
|
|
},
|
|
},
|
|
)
|
|
|
|
mcp = McpToolGateway(
|
|
project_configs={"default": _project_config()},
|
|
client=McpHttpClient(transport=httpx.MockTransport(handler)),
|
|
invocation_store=invocation_store,
|
|
)
|
|
worker_repository = InMemoryWorkerRepository()
|
|
event_store = InMemoryEventStore()
|
|
worker_service = WorkerService(worker_repository=worker_repository, event_store=event_store)
|
|
gateway = CapabilityAwareWorkerGateway(
|
|
worker_repository=worker_repository,
|
|
event_store=event_store,
|
|
invocation_store=invocation_store,
|
|
connection_manager=InMemoryWorkerConnectionManager(),
|
|
)
|
|
worker = worker_service.register(
|
|
RegisterWorkerRequest(
|
|
worker_id="worker_home_pc",
|
|
name="Home PC",
|
|
machine="DESKTOP-1",
|
|
os="windows",
|
|
version="0.1.0",
|
|
capabilities=["file.read"],
|
|
)
|
|
).worker
|
|
queued = gateway.dispatch(
|
|
project_id="default",
|
|
worker_session_id=worker.session_id,
|
|
command_name="file.read",
|
|
args={"path": "D:/test.txt"},
|
|
task_context={"task_id": "task_1"},
|
|
)
|
|
command = gateway.poll_commands(worker.session_id)[0]
|
|
gateway.complete_command(
|
|
worker.session_id,
|
|
WorkerCommandResult(
|
|
command_id=command.command_id,
|
|
task_id="task_1",
|
|
tool="file.read",
|
|
status="success",
|
|
started_at=__import__("datetime").datetime.now(__import__("datetime").UTC),
|
|
finished_at=__import__("datetime").datetime.now(__import__("datetime").UTC),
|
|
duration_ms=1,
|
|
result={"content": "ok"},
|
|
),
|
|
)
|
|
tool_result = mcp.call(
|
|
project_id="default",
|
|
server_id="one_c",
|
|
tool_name="one_c.run_sql",
|
|
args={"query": "select 1"},
|
|
task_context={"task_id": "task_1"},
|
|
)
|
|
|
|
registry = RunnerRegistry()
|
|
registry.register(PlannerNodeRunner())
|
|
registry.register(FinalizerNodeRunner(finalizer=DefaultFinalizer()))
|
|
engine = GraphExecutionEngine(
|
|
event_store=event_store,
|
|
runner_registry=registry,
|
|
reviewer=NoOpReviewer(),
|
|
)
|
|
orchestrator = OrchestratorService(
|
|
task_repository=InMemoryTaskRepository(),
|
|
graph_repository=InMemoryGraphRepository(),
|
|
confirmation_repository=InMemoryConfirmationRepository(),
|
|
event_store=event_store,
|
|
policy_evaluator=StaticProjectPolicyEvaluator(projects={"default": _project_config()}),
|
|
)
|
|
task = orchestrator.create_task(
|
|
CreateTaskRequest(project_id="default", goal="finalize", inputs={})
|
|
)
|
|
graph = orchestrator.plan_task(task.task_id)
|
|
engine.execute_ready_nodes(task=task, graph=graph)
|
|
engine.execute_ready_nodes(task=task, graph=graph)
|
|
|
|
assert queued.status == "queued"
|
|
assert tool_result.status == "success"
|
|
assert len(invocation_store.list_tool_invocations("task_1")) >= 2
|
|
assert graph.nodes[-1].output_data["status"] == "completed"
|