ToolSearch 工具分组

import asyncio
import json

from a13n_harness import AgentSpec, HarnessBuilder
from a13n_harness.capabilities import ToolSearchCapability, ToolSearchGroup
from pydantic_ai.messages import ModelRequest, ToolReturnPart
from pydantic_ai.models.function import DeltaToolCall, FunctionModel
from pydantic_ai.toolsets import FunctionToolset


def customer_name(customer_id: int) -> str:
    """Look up the display name for a customer ID."""
    return {42: "Ada"}.get(customer_id, "Unknown")


async def respond(messages, info):
    names = {tool.name for tool in info.function_tools}
    assert "customer_name" not in names and "crm__customer_name" not in names
    returns = [
        part
        for message in messages
        if isinstance(message, ModelRequest)
        for part in message.parts
        if isinstance(part, ToolReturnPart)
    ]
    if not returns:
        name, args = "discover_tools", {"query": "customer name", "group": "crm"}
    elif returns[-1].tool_name == "discover_tools":
        assert returns[-1].content["detail"] == "summary"
        match = returns[-1].content["tools"][0]
        name, args = "describe_tool", {"group": match["group"], "tool": match["tool"]}
    elif returns[-1].tool_name == "describe_tool":
        definition = returns[-1].content["tool"]
        assert "customer_id" in definition["parameters_json_schema"]["properties"]
        name, args = "call_tool", {
            "group": "crm", "tool": "customer_name", "arguments": {"customer_id": 42}
        }
    else:
        assert returns[-1].content == "Ada"
        yield "Customer Ada"
        return
    yield {0: DeltaToolCall(name=name, json_args=json.dumps(args))}


async def main():
    source = FunctionToolset(
        [customer_name], instructions="Use numeric customer IDs, not display names."
    )
    proxy = ToolSearchCapability(
        groups={
            "crm": ToolSearchGroup(source=source, description="Customer records")
        }
    )
    agent = HarnessBuilder().build(
        AgentSpec(),
        output_type=str,
        model=FunctionModel(stream_function=respond),
        capabilities=(proxy,),
    )
    result = await agent.run("Find customer 42")
    assert result.output_or_raise() == "Customer Ada"
    print(result.output_or_raise())


if __name__ == "__main__":
    asyncio.run(main())

uv run --no-sync python search_demo.py

Customer Ada

from a13n_harness.capabilities import ToolSearchCapability, ProxyToolSearchConfig, ToolSearchGroup

proxy = ToolSearchCapability(
    groups={"crm": ToolSearchGroup(source=source, description="Customer records")},
    config=ProxyToolSearchConfig(
        search_name="find_operations",
        describe_name="describe_operation",
        call_name="invoke_operation",
        list_name="list_operations",
        groups_name="list_domains",
        max_results=10,
        max_search_bytes=32_768,
    )
)

from a13n_harness.capabilities import ToolSearchCapability, ToolSearchGroup
from a13n_harness.mcp import (
    ContextualMCP,
    MCPContextHeaderBinding,
    MCPContextHeaders,
    MCPContextHeadersConfig,
)

knowledge = ToolSearchGroup(
    source=ContextualMCP(
        "https://mcp.example.com/mcp",
        id="knowledge-server",
        native=False,
        local=True,
        headers_factory=MCPContextHeaders(
            MCPContextHeadersConfig(
                headers={
                    "X-Run-ID": MCPContextHeaderBinding("context.run_id"),
                    "X-Tenant-ID": MCPContextHeaderBinding("identity.tenant_id"),
                }
            )
        ),
    ),
    description="Search and read the tenant's knowledge base",
)
capabilities = (ToolSearchCapability(groups={"knowledge": knowledge}),)

from a13n_harness.capabilities import CodeActCapability, ToolSearchCapability, ToolSearchGroup
from a13n_harness.toolsets import CodeActPolicyToolset, CodeActToolPolicy
from pydantic_ai.toolsets import FunctionToolset

source = CodeActPolicyToolset(
    wrapped=FunctionToolset([customer_name]),
    policy=CodeActToolPolicy(tools={"customer_name": True}),
    reject_unknown_tools=True,
)
proxy = ToolSearchCapability(
    groups={"crm": ToolSearchGroup(source=source, description="Customer records")},
)
capabilities = (proxy, CodeActCapability())

found = await discover_tools(query="customer name", group="crm", detail="full")
match = found["tools"][0]
name = await call_tool(
    group=match["group"],
    tool=match["tool"],
    arguments={"customer_id": 42},
)
name

async def main(inputs):
    return await call_tool(
        group="crm",
        tool="customer_name",
        arguments={"customer_id": inputs["customer_id"]},
    )

from a13n_harness.capabilities import ToolSearchPlan, ToolSearchSelection
from pydantic_ai.capabilities import Capability

# Inside main(), replace the source/proxy/builder setup:
customer_tools = Capability(id="customer-tools", tools=[customer_name])
agent = HarnessBuilder().build(
    AgentSpec(),
    output_type=str,
    model=FunctionModel(stream_function=respond),
    capabilities=(customer_tools,),
    tool_search=ToolSearchPlan(
        groups={
            "crm": ToolSearchSelection(
                "Customer records", capabilities=(customer_tools,)
            )
        }
    ),
)

from a13n_harness import AbstractHarnessPlugin
from a13n_harness.capabilities import ToolSearchPlan, ToolSearchSelection
from pydantic_ai.capabilities import Capability


class CustomerPlugin(AbstractHarnessPlugin):
    @property
    def plugin_id(self):
        return "customer-plugin"

    def get_capabilities(self):
        return (Capability(id="customer-tools", tools=[customer_name]),)


# Inside main(), replace the source/proxy/builder setup:
agent = HarnessBuilder().build(
    AgentSpec(),
    output_type=str,
    model=FunctionModel(stream_function=respond),
    plugins=(CustomerPlugin(),),
    tool_search=ToolSearchPlan(
        groups={
            "crm": ToolSearchSelection(
                "Customer records", plugins=("customer-plugin",)
            )
        }
    ),
)