153 lines
4.5 KiB
Python
153 lines
4.5 KiB
Python
from __future__ import annotations
|
|
|
|
import dataclasses
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from wf_authoring import NodeSpec
|
|
from wf_sources_mcp.catalog import (
|
|
CombinedCatalog,
|
|
DiscoveredPrompt,
|
|
DiscoveredResource,
|
|
snapshot_from_specs,
|
|
)
|
|
|
|
|
|
class EchoInput(BaseModel):
|
|
message: str
|
|
|
|
|
|
class EchoOutput(BaseModel):
|
|
text: str
|
|
|
|
|
|
async def _echo(payload: EchoInput) -> EchoOutput:
|
|
return EchoOutput(text=payload.message)
|
|
|
|
|
|
def _echo_spec() -> NodeSpec[EchoInput, EchoOutput]:
|
|
return NodeSpec(
|
|
name="echo",
|
|
input_model=EchoInput,
|
|
output_model=EchoOutput,
|
|
outcomes=("ok",),
|
|
fn=_echo,
|
|
description="Echo message",
|
|
is_async=True,
|
|
input_schema_contract={
|
|
"type": "object",
|
|
"properties": {"message": {"type": "string"}},
|
|
"required": ["message"],
|
|
},
|
|
output_schema_contract={
|
|
"type": "object",
|
|
"properties": {"text": {"type": "string"}},
|
|
},
|
|
)
|
|
|
|
|
|
def test_snapshot_from_specs_qualifies_nodes_resources_and_prompts() -> None:
|
|
snapshot = snapshot_from_specs(
|
|
"demo.default",
|
|
specs={"echo": _echo_spec()},
|
|
tool_display_names={"echo": "Echo Tool"},
|
|
resources=[
|
|
DiscoveredResource(
|
|
uri="demo://docs/guide",
|
|
name="guide",
|
|
title="Guide",
|
|
description="Read me",
|
|
mime_type="text/markdown",
|
|
metadata={"kind": "doc"},
|
|
)
|
|
],
|
|
prompts=[
|
|
DiscoveredPrompt(
|
|
name="summarize",
|
|
title="Summarize",
|
|
description="Summarize text",
|
|
arguments=[{"name": "topic"}],
|
|
metadata={"kind": "prompt"},
|
|
)
|
|
],
|
|
metadata={"source": "test"},
|
|
fetched_at_epoch_ms=123,
|
|
max_age_seconds=60,
|
|
)
|
|
|
|
assert snapshot.connection_id == "demo.default"
|
|
assert snapshot.nodes[0].qualified_name == "demo.default.echo"
|
|
assert snapshot.nodes[0].local_name == "echo"
|
|
assert snapshot.nodes[0].title == "Echo Tool"
|
|
assert snapshot.resources[0].qualified_name == "demo.default.guide"
|
|
assert snapshot.resources[0].uri == "demo://docs/guide"
|
|
assert snapshot.prompts[0].qualified_name == "demo.default.summarize"
|
|
assert snapshot.prompts[0].arguments == [{"name": "topic"}]
|
|
assert snapshot.metadata == {"source": "test"}
|
|
|
|
|
|
def test_snapshot_from_specs_preserves_already_qualified_node_name() -> None:
|
|
spec = dataclasses.replace(_echo_spec(), name="demo.default.echo")
|
|
|
|
snapshot = snapshot_from_specs(
|
|
"demo.default",
|
|
specs={"demo.default.echo": spec},
|
|
fetched_at_epoch_ms=123,
|
|
max_age_seconds=60,
|
|
)
|
|
|
|
assert snapshot.nodes[0].qualified_name == "demo.default.echo"
|
|
assert snapshot.nodes[0].local_name == "echo"
|
|
|
|
|
|
def test_combined_catalog_sorts_entries_and_serializes_payload() -> None:
|
|
first = snapshot_from_specs(
|
|
"zeta.default",
|
|
specs={"echo": _echo_spec()},
|
|
resources=[
|
|
DiscoveredResource(
|
|
uri="zeta://guide", name="guide", title=None, description=None
|
|
)
|
|
],
|
|
prompts=[DiscoveredPrompt(name="prompt", title=None, description=None)],
|
|
metadata={"order": "second"},
|
|
fetched_at_epoch_ms=2,
|
|
max_age_seconds=60,
|
|
)
|
|
second = snapshot_from_specs(
|
|
"alpha.default",
|
|
specs={"echo": _echo_spec()},
|
|
resources=[
|
|
DiscoveredResource(
|
|
uri="alpha://guide", name="guide", title=None, description=None
|
|
)
|
|
],
|
|
prompts=[DiscoveredPrompt(name="prompt", title=None, description=None)],
|
|
metadata={"order": "first"},
|
|
fetched_at_epoch_ms=1,
|
|
max_age_seconds=60,
|
|
)
|
|
|
|
catalog = CombinedCatalog(
|
|
snapshots={
|
|
first.connection_id: first,
|
|
second.connection_id: second,
|
|
}
|
|
)
|
|
payload = catalog.as_payload()
|
|
|
|
assert [entry.qualified_name for entry in catalog.entries()] == [
|
|
"alpha.default.echo",
|
|
"zeta.default.echo",
|
|
]
|
|
assert catalog.find_resource("alpha.default.guide") is not None
|
|
assert catalog.find_prompt("zeta.default.prompt") is not None
|
|
assert [node["qualified_name"] for node in payload["nodes"]] == [
|
|
"alpha.default.echo",
|
|
"zeta.default.echo",
|
|
]
|
|
assert [item["connection_id"] for item in payload["connections"]] == [
|
|
"alpha.default",
|
|
"zeta.default",
|
|
]
|