Runtime Tool Registry Example

This page embeds the full source for examples/runtime_tool_registry.py.

"""Example: reuse mapped tool handlers with explicit runtime state."""

import asyncio
import json

from swarmforge.swarm import (
    AgentTurnResult,
    InMemorySessionStore,
    SwarmDefinition,
    SwarmNode,
    SwarmSession,
    function_tool,
    process_swarm_stream,
)


class OpsTurnRunner:
    def __init__(self) -> None:
        self._call_count = 0

    async def run_turn(self, *, agent_node, contents, config):
        del agent_node, contents, config
        self._call_count += 1
        if self._call_count == 1:
            return AgentTurnResult(tool_calls=[{"id": "tool-1", "name": "lookup_order", "args": {"order_id": "123"}}])
        return AgentTurnResult(response_text="Order 123 has shipped.")


async def lookup_order(order_id: str, state=None, context=None):
    """Fetch order status for a customer order.

    Args:
        order_id: The order identifier.
    """

    return {
        "order_id": order_id,
        "tenant": state["tenant"],
        "account_id": context.state.get("account_id"),
        "status": "shipped",
    }


TOOLS = {"lookup_order": lookup_order}


swarm = SwarmDefinition(
    id="ops",
    name="Ops Swarm",
    nodes=[
        SwarmNode(
            id="ops-agent",
            node_key="ops",
            name="Ops",
            system_prompt="Use tools when an order lookup is needed.",
            enabled_tools=[function_tool(handler=lookup_order)],
            is_entry_node=True,
        )
    ],
)


async def main():
    session = SwarmSession.from_state(
        id="session-1",
        swarm=swarm,
        state={"account_id": "ACME-991"},
    )
    store = InMemorySessionStore()
    async for event in process_swarm_stream(
        session,
        "Where is order 123?",
        store=store,
        turn_runner=OpsTurnRunner(),
        tool_registry=TOOLS,
        tool_state={"tenant": "acme"},
    ):
        print(json.dumps(event, indent=2))


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

Run:

python examples/runtime_tool_registry.py