see that
This commit is contained in:
@@ -82,6 +82,9 @@ def test_service_compiles_and_runs_raw_plan() -> None:
|
|||||||
|
|
||||||
assert run.status == RunStatus.COMPLETED
|
assert run.status == RunStatus.COMPLETED
|
||||||
assert run.output == {"result": "final:hello"}
|
assert run.output == {"result": "final:hello"}
|
||||||
|
event_kinds = [event.kind for event in service.list_events()]
|
||||||
|
assert "workflow_run_started" in event_kinds
|
||||||
|
assert "workflow_run_completed" in event_kinds
|
||||||
|
|
||||||
|
|
||||||
def test_service_refreshes_catalog_from_adapter() -> None:
|
def test_service_refreshes_catalog_from_adapter() -> None:
|
||||||
@@ -163,3 +166,61 @@ def test_service_refreshes_catalog_from_adapter() -> None:
|
|||||||
},
|
},
|
||||||
}
|
}
|
||||||
]
|
]
|
||||||
|
event_kinds = [event.kind for event in service.list_events()]
|
||||||
|
assert "catalog_refresh_started" in event_kinds
|
||||||
|
assert "catalog_refresh_completed" in event_kinds
|
||||||
|
|
||||||
|
|
||||||
|
def test_service_records_tool_call_events() -> None:
|
||||||
|
service = WfMcpService(store=FileStore(local_temp_root() / "event_store"))
|
||||||
|
service.register_connection(
|
||||||
|
ConnectionConfig(id="demo.personal", server="demo", account="personal")
|
||||||
|
)
|
||||||
|
service.register_adapter("demo", FakeAdapter())
|
||||||
|
|
||||||
|
asyncio.run(service.refresh_connection_catalog("demo.personal"))
|
||||||
|
|
||||||
|
plan = RawWorkflowPlan(
|
||||||
|
name="tool_only_plan",
|
||||||
|
input_schema={
|
||||||
|
"type": "object",
|
||||||
|
"properties": {"text": {"type": "string"}},
|
||||||
|
"required": ["text"],
|
||||||
|
},
|
||||||
|
state_schema={
|
||||||
|
"fields": {
|
||||||
|
"echoed": {"type": "string"},
|
||||||
|
}
|
||||||
|
},
|
||||||
|
output_schema={
|
||||||
|
"type": "object",
|
||||||
|
"properties": {"echoed": {"type": "string"}},
|
||||||
|
"required": ["echoed"],
|
||||||
|
},
|
||||||
|
start="echo",
|
||||||
|
nodes=[
|
||||||
|
{
|
||||||
|
"id": "echo",
|
||||||
|
"type": "node",
|
||||||
|
"node": "demo.personal.echo_tool",
|
||||||
|
"in_map": {"input.text": "text"},
|
||||||
|
"out_map": {"echoed": "state.echoed"},
|
||||||
|
}
|
||||||
|
],
|
||||||
|
edges=[
|
||||||
|
{"from": "echo", "outcome": "ok", "to": END},
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
run = asyncio.run(service.run_workflow_from_plan(plan, {"text": "hello"}))
|
||||||
|
|
||||||
|
assert run.status == RunStatus.COMPLETED
|
||||||
|
tool_events = [
|
||||||
|
event for event in service.list_events() if "tool_call" in event.kind
|
||||||
|
]
|
||||||
|
assert [event.kind for event in tool_events] == [
|
||||||
|
"tool_call_started",
|
||||||
|
"tool_call_completed",
|
||||||
|
]
|
||||||
|
assert tool_events[0].capability_id == "demo.personal.echo_tool"
|
||||||
|
assert tool_events[1].payload["outcome"] == "ok"
|
||||||
|
|||||||
@@ -12,6 +12,7 @@ from .discovery import (
|
|||||||
discover_connection_capabilities,
|
discover_connection_capabilities,
|
||||||
specs_from_discovered_tools,
|
specs_from_discovered_tools,
|
||||||
)
|
)
|
||||||
|
from .events import McpEvent, make_event
|
||||||
from .models import (
|
from .models import (
|
||||||
AuthRecord,
|
AuthRecord,
|
||||||
CatalogNodeEntry,
|
CatalogNodeEntry,
|
||||||
@@ -41,12 +42,14 @@ __all__ = [
|
|||||||
"DiscoveredResource",
|
"DiscoveredResource",
|
||||||
"DiscoveredTool",
|
"DiscoveredTool",
|
||||||
"FileStore",
|
"FileStore",
|
||||||
|
"McpEvent",
|
||||||
"McpSdkAdapter",
|
"McpSdkAdapter",
|
||||||
"RawWorkflowPlan",
|
"RawWorkflowPlan",
|
||||||
"Store",
|
"Store",
|
||||||
"ToolCallResult",
|
"ToolCallResult",
|
||||||
"WfMcpService",
|
"WfMcpService",
|
||||||
"discover_connection_capabilities",
|
"discover_connection_capabilities",
|
||||||
|
"make_event",
|
||||||
"parse_connection_id",
|
"parse_connection_id",
|
||||||
"qualify_node_name",
|
"qualify_node_name",
|
||||||
"specs_from_discovered_tools",
|
"specs_from_discovered_tools",
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from collections.abc import Callable
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
@@ -11,6 +12,7 @@ from .adapters import (
|
|||||||
DiscoveredResource,
|
DiscoveredResource,
|
||||||
DiscoveredTool,
|
DiscoveredTool,
|
||||||
)
|
)
|
||||||
|
from .events import McpEvent
|
||||||
from .models import AuthRecord, ConnectionConfig
|
from .models import AuthRecord, ConnectionConfig
|
||||||
from .wrappers import wrap_discovered_tool
|
from .wrappers import wrap_discovered_tool
|
||||||
|
|
||||||
@@ -47,6 +49,7 @@ def specs_from_discovered_tools(
|
|||||||
auth: AuthRecord | None,
|
auth: AuthRecord | None,
|
||||||
adapter: BackendAdapter,
|
adapter: BackendAdapter,
|
||||||
tools: list[DiscoveredTool],
|
tools: list[DiscoveredTool],
|
||||||
|
emit_event: Callable[[McpEvent], None] | None = None,
|
||||||
) -> list[NodeSpec[Any, Any]]:
|
) -> list[NodeSpec[Any, Any]]:
|
||||||
return [
|
return [
|
||||||
wrap_discovered_tool(
|
wrap_discovered_tool(
|
||||||
@@ -54,6 +57,7 @@ def specs_from_discovered_tools(
|
|||||||
auth=auth,
|
auth=auth,
|
||||||
adapter=adapter,
|
adapter=adapter,
|
||||||
tool=tool,
|
tool=tool,
|
||||||
|
emit_event=emit_event,
|
||||||
)
|
)
|
||||||
for tool in tools
|
for tool in tools
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -0,0 +1,33 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import time
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(slots=True)
|
||||||
|
class McpEvent:
|
||||||
|
kind: str
|
||||||
|
timestamp_epoch_ms: int
|
||||||
|
connection_id: str | None = None
|
||||||
|
capability_id: str | None = None
|
||||||
|
workflow_name: str | None = None
|
||||||
|
payload: dict[str, Any] = field(default_factory=dict)
|
||||||
|
|
||||||
|
|
||||||
|
def make_event(
|
||||||
|
kind: str,
|
||||||
|
*,
|
||||||
|
connection_id: str | None = None,
|
||||||
|
capability_id: str | None = None,
|
||||||
|
workflow_name: str | None = None,
|
||||||
|
payload: dict[str, Any] | None = None,
|
||||||
|
) -> McpEvent:
|
||||||
|
return McpEvent(
|
||||||
|
kind=kind,
|
||||||
|
timestamp_epoch_ms=int(time.time() * 1000),
|
||||||
|
connection_id=connection_id,
|
||||||
|
capability_id=capability_id,
|
||||||
|
workflow_name=workflow_name,
|
||||||
|
payload=payload or {},
|
||||||
|
)
|
||||||
@@ -44,29 +44,31 @@ def _tool_to_discovered(tool: McpTool) -> DiscoveredTool:
|
|||||||
input_schema=tool.inputSchema,
|
input_schema=tool.inputSchema,
|
||||||
output_schema=output_schema,
|
output_schema=output_schema,
|
||||||
outcomes=("ok", "error"),
|
outcomes=("ok", "error"),
|
||||||
metadata=tool.model_dump(by_alias=True),
|
metadata=tool.model_dump(by_alias=True, mode="json"),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def _resource_to_discovered(resource: McpResource) -> DiscoveredResource:
|
def _resource_to_discovered(resource: McpResource) -> DiscoveredResource:
|
||||||
|
local_name = resource.name or str(resource.uri)
|
||||||
return DiscoveredResource(
|
return DiscoveredResource(
|
||||||
uri=str(resource.uri),
|
uri=str(resource.uri),
|
||||||
name=str(resource.uri),
|
name=local_name,
|
||||||
description=resource.description,
|
description=resource.description,
|
||||||
mime_type=resource.mimeType,
|
mime_type=resource.mimeType,
|
||||||
metadata=resource.model_dump(by_alias=True),
|
metadata=resource.model_dump(by_alias=True, mode="json"),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def _prompt_to_discovered(prompt: McpPrompt) -> DiscoveredPrompt:
|
def _prompt_to_discovered(prompt: McpPrompt) -> DiscoveredPrompt:
|
||||||
arguments = [
|
arguments = [
|
||||||
argument.model_dump(by_alias=True) for argument in prompt.arguments or []
|
argument.model_dump(by_alias=True, mode="json")
|
||||||
|
for argument in prompt.arguments or []
|
||||||
]
|
]
|
||||||
return DiscoveredPrompt(
|
return DiscoveredPrompt(
|
||||||
name=prompt.name,
|
name=prompt.name,
|
||||||
description=prompt.description,
|
description=prompt.description,
|
||||||
arguments=arguments,
|
arguments=arguments,
|
||||||
metadata=prompt.model_dump(by_alias=True),
|
metadata=prompt.model_dump(by_alias=True, mode="json"),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
+64
-1
@@ -11,6 +11,7 @@ from .adapters import BackendAdapter
|
|||||||
from .catalog import CombinedCatalog, snapshot_from_specs
|
from .catalog import CombinedCatalog, snapshot_from_specs
|
||||||
from .connections import ConnectionRegistry, parse_connection_id, qualify_node_name
|
from .connections import ConnectionRegistry, parse_connection_id, qualify_node_name
|
||||||
from .discovery import discover_connection_capabilities, specs_from_discovered_tools
|
from .discovery import discover_connection_capabilities, specs_from_discovered_tools
|
||||||
|
from .events import McpEvent, make_event
|
||||||
from .models import AuthRecord, CatalogSnapshot, ConnectionConfig, RawWorkflowPlan
|
from .models import AuthRecord, CatalogSnapshot, ConnectionConfig, RawWorkflowPlan
|
||||||
from .store import Store
|
from .store import Store
|
||||||
|
|
||||||
@@ -36,16 +37,31 @@ class WfMcpService:
|
|||||||
specs_by_connection: dict[str, dict[str, NodeSpec[Any, Any]]] = field(
|
specs_by_connection: dict[str, dict[str, NodeSpec[Any, Any]]] = field(
|
||||||
default_factory=dict
|
default_factory=dict
|
||||||
)
|
)
|
||||||
|
events: list[McpEvent] = field(default_factory=list)
|
||||||
|
|
||||||
def register_connection(self, connection: ConnectionConfig) -> None:
|
def register_connection(self, connection: ConnectionConfig) -> None:
|
||||||
parse_connection_id(connection.id)
|
parse_connection_id(connection.id)
|
||||||
self.connections.register(connection)
|
self.connections.register(connection)
|
||||||
|
self._record_event(
|
||||||
|
make_event(
|
||||||
|
"connection_registered",
|
||||||
|
connection_id=connection.id,
|
||||||
|
payload={"server": connection.server, "account": connection.account},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
def register_adapter(self, server: str, adapter: BackendAdapter) -> None:
|
def register_adapter(self, server: str, adapter: BackendAdapter) -> None:
|
||||||
self.adapters[server] = adapter
|
self.adapters[server] = adapter
|
||||||
|
|
||||||
def save_auth(self, record: AuthRecord) -> None:
|
def save_auth(self, record: AuthRecord) -> None:
|
||||||
self.store.save_auth(record)
|
self.store.save_auth(record)
|
||||||
|
self._record_event(
|
||||||
|
make_event(
|
||||||
|
"auth_saved",
|
||||||
|
connection_id=record.connection_id,
|
||||||
|
payload={"scheme": record.scheme},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
def load_auth(self, connection_id: str) -> AuthRecord | None:
|
def load_auth(self, connection_id: str) -> AuthRecord | None:
|
||||||
return self.store.load_auth(connection_id)
|
return self.store.load_auth(connection_id)
|
||||||
@@ -71,6 +87,13 @@ class WfMcpService:
|
|||||||
max_age_seconds=max_age_seconds or self.default_catalog_max_age_seconds,
|
max_age_seconds=max_age_seconds or self.default_catalog_max_age_seconds,
|
||||||
)
|
)
|
||||||
self.store.save_catalog(snapshot)
|
self.store.save_catalog(snapshot)
|
||||||
|
self._record_event(
|
||||||
|
make_event(
|
||||||
|
"specs_registered",
|
||||||
|
connection_id=connection_id,
|
||||||
|
payload={"node_count": len(qualified_specs)},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
def get_catalog(self) -> CombinedCatalog:
|
def get_catalog(self) -> CombinedCatalog:
|
||||||
snapshots: dict[str, CatalogSnapshot] = {}
|
snapshots: dict[str, CatalogSnapshot] = {}
|
||||||
@@ -92,6 +115,13 @@ class WfMcpService:
|
|||||||
raise KeyError(f"no adapter registered for server {connection.server!r}")
|
raise KeyError(f"no adapter registered for server {connection.server!r}")
|
||||||
|
|
||||||
auth = self.load_auth(connection_id)
|
auth = self.load_auth(connection_id)
|
||||||
|
self._record_event(
|
||||||
|
make_event(
|
||||||
|
"catalog_refresh_started",
|
||||||
|
connection_id=connection_id,
|
||||||
|
payload={"server": connection.server},
|
||||||
|
)
|
||||||
|
)
|
||||||
capabilities = await discover_connection_capabilities(
|
capabilities = await discover_connection_capabilities(
|
||||||
connection=connection,
|
connection=connection,
|
||||||
auth=auth,
|
auth=auth,
|
||||||
@@ -102,6 +132,7 @@ class WfMcpService:
|
|||||||
auth=auth,
|
auth=auth,
|
||||||
adapter=adapter,
|
adapter=adapter,
|
||||||
tools=capabilities.tools,
|
tools=capabilities.tools,
|
||||||
|
emit_event=self._record_event,
|
||||||
)
|
)
|
||||||
self.register_specs(
|
self.register_specs(
|
||||||
connection_id,
|
connection_id,
|
||||||
@@ -118,6 +149,17 @@ class WfMcpService:
|
|||||||
max_age_seconds=max_age_seconds or self.default_catalog_max_age_seconds,
|
max_age_seconds=max_age_seconds or self.default_catalog_max_age_seconds,
|
||||||
)
|
)
|
||||||
self.store.save_catalog(snapshot)
|
self.store.save_catalog(snapshot)
|
||||||
|
self._record_event(
|
||||||
|
make_event(
|
||||||
|
"catalog_refresh_completed",
|
||||||
|
connection_id=connection_id,
|
||||||
|
payload={
|
||||||
|
"node_count": len(snapshot.nodes),
|
||||||
|
"resource_count": len(snapshot.resources),
|
||||||
|
"prompt_count": len(snapshot.prompts),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
def compile_plan(self, plan: RawWorkflowPlan) -> Workflow:
|
def compile_plan(self, plan: RawWorkflowPlan) -> Workflow:
|
||||||
node_defs: dict[str, Any] = {}
|
node_defs: dict[str, Any] = {}
|
||||||
@@ -145,6 +187,13 @@ class WfMcpService:
|
|||||||
plan: RawWorkflowPlan,
|
plan: RawWorkflowPlan,
|
||||||
workflow_input: dict[str, Any],
|
workflow_input: dict[str, Any],
|
||||||
):
|
):
|
||||||
|
self._record_event(
|
||||||
|
make_event(
|
||||||
|
"workflow_run_started",
|
||||||
|
workflow_name=plan.name,
|
||||||
|
payload={"input_keys": sorted(workflow_input.keys())},
|
||||||
|
)
|
||||||
|
)
|
||||||
workflow = self.compile_plan(plan)
|
workflow = self.compile_plan(plan)
|
||||||
specs = [
|
specs = [
|
||||||
self._get_qualified_spec(node.node)
|
self._get_qualified_spec(node.node)
|
||||||
@@ -152,7 +201,18 @@ class WfMcpService:
|
|||||||
if isinstance(node, NodeUse)
|
if isinstance(node, NodeUse)
|
||||||
]
|
]
|
||||||
registry = build_async_registry(*specs)
|
registry = build_async_registry(*specs)
|
||||||
return await execute_workflow_async(workflow, workflow_input, registry)
|
run = await execute_workflow_async(workflow, workflow_input, registry)
|
||||||
|
self._record_event(
|
||||||
|
make_event(
|
||||||
|
"workflow_run_completed",
|
||||||
|
workflow_name=plan.name,
|
||||||
|
payload={"status": run.status.value},
|
||||||
|
)
|
||||||
|
)
|
||||||
|
return run
|
||||||
|
|
||||||
|
def list_events(self) -> list[McpEvent]:
|
||||||
|
return list(self.events)
|
||||||
|
|
||||||
def _get_qualified_spec(self, qualified_name: str) -> NodeSpec[Any, Any]:
|
def _get_qualified_spec(self, qualified_name: str) -> NodeSpec[Any, Any]:
|
||||||
connection_id, _ = qualified_name.rsplit(".", 1)
|
connection_id, _ = qualified_name.rsplit(".", 1)
|
||||||
@@ -160,3 +220,6 @@ class WfMcpService:
|
|||||||
if specs is None or qualified_name not in specs:
|
if specs is None or qualified_name not in specs:
|
||||||
raise KeyError(f"unknown qualified node {qualified_name!r}")
|
raise KeyError(f"unknown qualified node {qualified_name!r}")
|
||||||
return specs[qualified_name]
|
return specs[qualified_name]
|
||||||
|
|
||||||
|
def _record_event(self, event: McpEvent) -> None:
|
||||||
|
self.events.append(event)
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from collections.abc import Callable
|
||||||
from typing import Any, cast
|
from typing import Any, cast
|
||||||
|
|
||||||
from pydantic import BaseModel, ConfigDict, Field, create_model
|
from pydantic import BaseModel, ConfigDict, Field, create_model
|
||||||
@@ -8,6 +9,7 @@ from wf_authoring import NodeReturn, NodeSpec
|
|||||||
from wf_core import RuntimeContext
|
from wf_core import RuntimeContext
|
||||||
|
|
||||||
from .adapters import BackendAdapter, DiscoveredTool
|
from .adapters import BackendAdapter, DiscoveredTool
|
||||||
|
from .events import McpEvent, make_event
|
||||||
from .models import AuthRecord, ConnectionConfig
|
from .models import AuthRecord, ConnectionConfig
|
||||||
|
|
||||||
|
|
||||||
@@ -35,6 +37,7 @@ def wrap_discovered_tool(
|
|||||||
auth: AuthRecord | None,
|
auth: AuthRecord | None,
|
||||||
adapter: BackendAdapter,
|
adapter: BackendAdapter,
|
||||||
tool: DiscoveredTool,
|
tool: DiscoveredTool,
|
||||||
|
emit_event: Callable[[McpEvent], None] | None = None,
|
||||||
) -> NodeSpec[BaseModel, BaseModel]:
|
) -> NodeSpec[BaseModel, BaseModel]:
|
||||||
input_model = _model_from_schema(
|
input_model = _model_from_schema(
|
||||||
f"{connection.id}_{tool.name}_Input",
|
f"{connection.id}_{tool.name}_Input",
|
||||||
@@ -49,12 +52,33 @@ def wrap_discovered_tool(
|
|||||||
payload: BaseModel,
|
payload: BaseModel,
|
||||||
ctx: RuntimeContext,
|
ctx: RuntimeContext,
|
||||||
) -> NodeReturn[BaseModel]:
|
) -> NodeReturn[BaseModel]:
|
||||||
|
if emit_event is not None:
|
||||||
|
emit_event(
|
||||||
|
make_event(
|
||||||
|
"tool_call_started",
|
||||||
|
connection_id=connection.id,
|
||||||
|
capability_id=f"{connection.id}.{tool.name}",
|
||||||
|
payload={"input": payload.model_dump()},
|
||||||
|
)
|
||||||
|
)
|
||||||
result = await adapter.call_tool(
|
result = await adapter.call_tool(
|
||||||
connection=connection,
|
connection=connection,
|
||||||
auth=auth,
|
auth=auth,
|
||||||
tool_name=tool.name,
|
tool_name=tool.name,
|
||||||
payload=payload.model_dump(),
|
payload=payload.model_dump(),
|
||||||
)
|
)
|
||||||
|
if emit_event is not None:
|
||||||
|
emit_event(
|
||||||
|
make_event(
|
||||||
|
"tool_call_completed",
|
||||||
|
connection_id=connection.id,
|
||||||
|
capability_id=f"{connection.id}.{tool.name}",
|
||||||
|
payload={
|
||||||
|
"outcome": result.outcome,
|
||||||
|
"meta": result.meta,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
)
|
||||||
return NodeReturn(
|
return NodeReturn(
|
||||||
outcome=result.outcome,
|
outcome=result.outcome,
|
||||||
output=output_model.model_validate(result.output),
|
output=output_model.model_validate(result.output),
|
||||||
|
|||||||
Reference in New Issue
Block a user