adapters
This commit is contained in:
+90
-1
@@ -2,12 +2,21 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from wf_authoring import NodeReturn, node
|
||||
from wf_core import END, RuntimeContext, RunStatus
|
||||
from wf_mcp import AuthRecord, ConnectionConfig, FileStore, RawWorkflowPlan, WfMcpService
|
||||
from wf_mcp import (
|
||||
AuthRecord,
|
||||
ConnectionConfig,
|
||||
DiscoveredTool,
|
||||
FileStore,
|
||||
RawWorkflowPlan,
|
||||
ToolCallResult,
|
||||
WfMcpService,
|
||||
)
|
||||
|
||||
|
||||
class EchoInput(BaseModel):
|
||||
@@ -47,6 +56,44 @@ def _local_temp_root() -> Path:
|
||||
return root
|
||||
|
||||
|
||||
class FakeAdapter:
|
||||
async def list_tools(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
) -> list[DiscoveredTool]:
|
||||
return [
|
||||
DiscoveredTool(
|
||||
name="echo_tool",
|
||||
description="Echo text back",
|
||||
input_schema={
|
||||
"type": "object",
|
||||
"properties": {"text": {"type": "string"}},
|
||||
"required": ["text"],
|
||||
},
|
||||
output_schema={
|
||||
"type": "object",
|
||||
"properties": {"echoed": {"type": "string"}},
|
||||
"required": ["echoed"],
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
async def call_tool(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
tool_name: str,
|
||||
payload: dict[str, Any],
|
||||
) -> ToolCallResult:
|
||||
if tool_name != "echo_tool":
|
||||
raise KeyError(tool_name)
|
||||
return ToolCallResult(
|
||||
outcome="ok",
|
||||
output={"echoed": str(payload["text"])},
|
||||
)
|
||||
|
||||
|
||||
def test_file_store_round_trips_auth() -> None:
|
||||
store = FileStore(_local_temp_root() / "auth_store")
|
||||
record = AuthRecord(
|
||||
@@ -129,3 +176,45 @@ def test_service_compiles_and_runs_raw_plan() -> None:
|
||||
|
||||
assert run.status == RunStatus.COMPLETED
|
||||
assert run.output == {"result": "final:hello"}
|
||||
|
||||
|
||||
def test_service_refreshes_catalog_from_adapter() -> None:
|
||||
service = WfMcpService(store=FileStore(_local_temp_root() / "adapter_store"))
|
||||
service.register_connection(
|
||||
ConnectionConfig(id="demo.personal", server="demo", account="personal")
|
||||
)
|
||||
service.save_auth(
|
||||
AuthRecord(
|
||||
connection_id="demo.personal",
|
||||
scheme="token",
|
||||
payload={"token": "abc"},
|
||||
)
|
||||
)
|
||||
service.register_adapter("demo", FakeAdapter())
|
||||
|
||||
asyncio.run(service.refresh_connection_catalog("demo.personal"))
|
||||
|
||||
payload = service.get_catalog().as_payload()
|
||||
assert payload["nodes"] == [
|
||||
{
|
||||
"qualified_name": "demo.personal.echo_tool",
|
||||
"connection_id": "demo.personal",
|
||||
"local_name": "echo_tool",
|
||||
"description": "Echo text back",
|
||||
"outcomes": ["ok"],
|
||||
"input_schema": {
|
||||
"additionalProperties": True,
|
||||
"properties": {"text": {"title": "Text"}},
|
||||
"required": ["text"],
|
||||
"title": "demo.personal_echo_tool_Input",
|
||||
"type": "object",
|
||||
},
|
||||
"output_schema": {
|
||||
"additionalProperties": True,
|
||||
"properties": {"echoed": {"title": "Echoed"}},
|
||||
"required": ["echoed"],
|
||||
"title": "demo.personal_echo_tool_Output",
|
||||
"type": "object",
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
from .adapters import BackendAdapter, DiscoveredTool, ToolCallResult
|
||||
from .catalog import CombinedCatalog
|
||||
from .connections import ConnectionRegistry, parse_connection_id, qualify_node_name
|
||||
from .models import (
|
||||
@@ -9,18 +10,23 @@ from .models import (
|
||||
)
|
||||
from .service import WfMcpService
|
||||
from .store import FileStore, Store
|
||||
from .wrappers import wrap_discovered_tool
|
||||
|
||||
__all__ = [
|
||||
"AuthRecord",
|
||||
"BackendAdapter",
|
||||
"CatalogNodeEntry",
|
||||
"CatalogSnapshot",
|
||||
"CombinedCatalog",
|
||||
"ConnectionConfig",
|
||||
"ConnectionRegistry",
|
||||
"DiscoveredTool",
|
||||
"FileStore",
|
||||
"RawWorkflowPlan",
|
||||
"Store",
|
||||
"ToolCallResult",
|
||||
"WfMcpService",
|
||||
"parse_connection_id",
|
||||
"qualify_node_name",
|
||||
"wrap_discovered_tool",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Protocol
|
||||
|
||||
from .models import AuthRecord, ConnectionConfig
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class DiscoveredTool:
|
||||
name: str
|
||||
description: str | None
|
||||
input_schema: dict[str, Any]
|
||||
output_schema: dict[str, Any]
|
||||
outcomes: tuple[str, ...] = ("ok",)
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ToolCallResult:
|
||||
outcome: str
|
||||
output: dict[str, Any] = field(default_factory=dict)
|
||||
meta: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
class BackendAdapter(Protocol):
|
||||
async def list_tools(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
) -> list[DiscoveredTool]: ...
|
||||
|
||||
async def call_tool(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
tool_name: str,
|
||||
payload: dict[str, Any],
|
||||
) -> ToolCallResult: ...
|
||||
@@ -7,10 +7,12 @@ from typing import Any
|
||||
from wf_authoring import NodeSpec, build_async_registry
|
||||
from wf_core import NodeUse, Workflow, execute_workflow_async
|
||||
|
||||
from .adapters import BackendAdapter
|
||||
from .catalog import CombinedCatalog, snapshot_from_specs
|
||||
from .connections import ConnectionRegistry, parse_connection_id, qualify_node_name
|
||||
from .models import AuthRecord, CatalogSnapshot, ConnectionConfig, RawWorkflowPlan
|
||||
from .store import Store
|
||||
from .wrappers import wrap_discovered_tool
|
||||
|
||||
|
||||
def _qualify_spec(connection_id: str, spec: NodeSpec[Any, Any]) -> NodeSpec[Any, Any]:
|
||||
@@ -30,6 +32,7 @@ class WfMcpService:
|
||||
store: Store
|
||||
default_catalog_max_age_seconds: int = 300
|
||||
connections: ConnectionRegistry = field(default_factory=ConnectionRegistry)
|
||||
adapters: dict[str, BackendAdapter] = field(default_factory=dict)
|
||||
specs_by_connection: dict[str, dict[str, NodeSpec[Any, Any]]] = field(
|
||||
default_factory=dict
|
||||
)
|
||||
@@ -38,6 +41,9 @@ class WfMcpService:
|
||||
parse_connection_id(connection.id)
|
||||
self.connections.register(connection)
|
||||
|
||||
def register_adapter(self, server: str, adapter: BackendAdapter) -> None:
|
||||
self.adapters[server] = adapter
|
||||
|
||||
def save_auth(self, record: AuthRecord) -> None:
|
||||
self.store.save_auth(record)
|
||||
|
||||
@@ -72,6 +78,34 @@ class WfMcpService:
|
||||
snapshots[connection.id] = snapshot
|
||||
return CombinedCatalog(snapshots=snapshots)
|
||||
|
||||
async def refresh_connection_catalog(
|
||||
self,
|
||||
connection_id: str,
|
||||
*,
|
||||
max_age_seconds: int | None = None,
|
||||
) -> None:
|
||||
connection = self.connections.get(connection_id)
|
||||
adapter = self.adapters.get(connection.server)
|
||||
if adapter is None:
|
||||
raise KeyError(f"no adapter registered for server {connection.server!r}")
|
||||
|
||||
auth = self.load_auth(connection_id)
|
||||
tools = await adapter.list_tools(connection, auth)
|
||||
specs = [
|
||||
wrap_discovered_tool(
|
||||
connection=connection,
|
||||
auth=auth,
|
||||
adapter=adapter,
|
||||
tool=tool,
|
||||
)
|
||||
for tool in tools
|
||||
]
|
||||
self.register_specs(
|
||||
connection_id,
|
||||
*specs,
|
||||
max_age_seconds=max_age_seconds,
|
||||
)
|
||||
|
||||
def compile_plan(self, plan: RawWorkflowPlan) -> Workflow:
|
||||
node_defs: dict[str, Any] = {}
|
||||
for step in plan.nodes:
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, cast
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, create_model
|
||||
|
||||
from wf_authoring import NodeReturn, NodeSpec
|
||||
from wf_core import RuntimeContext
|
||||
|
||||
from .adapters import BackendAdapter, DiscoveredTool
|
||||
from .models import AuthRecord, ConnectionConfig
|
||||
|
||||
|
||||
def _model_from_schema(name: str, schema: dict[str, Any]) -> type[BaseModel]:
|
||||
properties = cast(dict[str, Any], schema.get("properties", {}))
|
||||
required = set(cast(list[str], schema.get("required", [])))
|
||||
field_defs: dict[str, tuple[object, object]] = {}
|
||||
|
||||
for field_name in properties:
|
||||
default = ... if field_name in required else None
|
||||
field_defs[field_name] = (Any, Field(default=default))
|
||||
|
||||
raw_field_defs = cast(dict[str, Any], field_defs)
|
||||
model = create_model(
|
||||
name,
|
||||
__config__=ConfigDict(extra="allow"),
|
||||
**raw_field_defs,
|
||||
)
|
||||
return cast(type[BaseModel], model)
|
||||
|
||||
|
||||
def wrap_discovered_tool(
|
||||
*,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
adapter: BackendAdapter,
|
||||
tool: DiscoveredTool,
|
||||
) -> NodeSpec[BaseModel, BaseModel]:
|
||||
input_model = _model_from_schema(
|
||||
f"{connection.id}_{tool.name}_Input",
|
||||
tool.input_schema,
|
||||
)
|
||||
output_model = _model_from_schema(
|
||||
f"{connection.id}_{tool.name}_Output",
|
||||
tool.output_schema,
|
||||
)
|
||||
|
||||
async def invoke_tool(
|
||||
payload: BaseModel,
|
||||
ctx: RuntimeContext,
|
||||
) -> NodeReturn[BaseModel]:
|
||||
result = await adapter.call_tool(
|
||||
connection=connection,
|
||||
auth=auth,
|
||||
tool_name=tool.name,
|
||||
payload=payload.model_dump(),
|
||||
)
|
||||
return NodeReturn(
|
||||
outcome=result.outcome,
|
||||
output=output_model.model_validate(result.output),
|
||||
)
|
||||
|
||||
return NodeSpec(
|
||||
name=tool.name,
|
||||
input_model=input_model,
|
||||
output_model=output_model,
|
||||
outcomes=tool.outcomes,
|
||||
fn=invoke_tool,
|
||||
description=tool.description,
|
||||
is_async=True,
|
||||
)
|
||||
Reference in New Issue
Block a user