Composing Agents (A2A)

Every agent and workflow in AIMU implements the single Runner interface (run() + messages), so they compose freely. This notebook shows two layers of composition:

A2A vs MCP: MCP (aimu.tools.MCPClient) shares tools across a process boundary; A2A shares whole agents. They are complementary; see docs/explanation/a2a-vs-mcp.md.

The A2A sections require the optional extra: pip install 'aimu[a2a]'.

Setup

The model-backed cells (the orchestrator and the routing agent) use a small local Ollama model so the notebook runs quickly. Swap in any tool-capable model, or aimu.client() to auto-resolve a local default.

import aimu
from aimu.models import OllamaModel

MODEL = OllamaModel.QWEN_3_8B  # small + tool-capable; change to taste

A - Composing runners in-process

The Runner interface is tiny: a run(task) method and a messages property. Here are two deterministic runners (no model needed) so the composition mechanics are easy to see. In real use these would be an Agent, Chain, Router, etc.

from aimu.agents import Runner


class FAQRunner(Runner):
    """A trivial knowledge-base lookup, used to demo composition without a model backend."""

    _FAQ = {
        "refund": "Refunds are processed within 5 business days.",
        "hours": "Support is available 9am-5pm, Monday to Friday.",
    }

    def __init__(self, name="faq"):
        self.name = name
        self.system_message = "Answer questions from the support knowledge base."

    def run(self, task, generate_kwargs=None, stream=False, images=None):
        for key, answer in self._FAQ.items():
            if key in task.lower():
                return answer
        return "No knowledge-base entry matched."

    @property
    def messages(self):
        return {self.name: []}


class GreetingRunner(Runner):
    def __init__(self, name="greeter"):
        self.name = name
        self.system_message = "Produce a friendly greeting for the user."

    def run(self, task, generate_kwargs=None, stream=False, images=None):
        return "Hello, and welcome to ACME support!"

    @property
    def messages(self):
        return {self.name: []}


faq = FAQRunner()
print(faq.run("What are your support hours?"))

runner.as_tool(): turn any runner into a tool

as_tool() wraps a runner’s run() as a @aimu.tool-style callable. The tool name comes from the runner’s name; the description from its system_message (or a generic fallback for workflows). This is the seam that lets an autonomous agent call another agent or workflow. No model needed to inspect it.

tool_fn = faq.as_tool()

print("tool name: ", tool_fn.__name__)
print("description:", tool_fn.__tool_spec__["function"]["description"])
print("schema:     ", tool_fn.__tool_spec__["function"]["parameters"])
print("call:       ", tool_fn("Tell me about your refund policy"))

OrchestratorAgent.assemble() accepts any Runner as a worker

assemble() wraps each worker via as_tool(), so workers can be any Runner: an Agent, a Chain, a Router, a remote A2A agent, or the deterministic runners above. The orchestrator itself is a model-driven Agent that dispatches to workers by name.

from aimu.agents import OrchestratorAgent

orchestrator = OrchestratorAgent.assemble(
    aimu.client(MODEL),
    "You coordinate support. Use the greeter tool to greet, and the faq tool for policy questions.",
    workers=[GreetingRunner(), FAQRunner()],
)

print("workers registered as tools:", [t.__name__ for t in orchestrator._orchestrator.model_client.tools])
print()
print(orchestrator.run("Greet me, then tell me your refund policy."))

A workflow as a tool inside an autonomous agent

Chain/Router/Parallel already accept other runners as members. The new piece is the reverse direction: handing a workflow to an autonomous Agent as a tool, via as_tool().

from aimu.agents import Agent, Chain

# A two-step chain (deterministic runners) exposed as a single tool.
lookup_then_greet = Chain(agents=[FAQRunner(), GreetingRunner()], name="lookup_then_greet")

agent = Agent(
    aimu.client(MODEL),
    "You are a router. For any support enquiry, call the lookup_then_greet tool, then report its result.",
    tools=[lookup_then_greet.as_tool()],
    max_iterations=3,
)
print(agent.run("A customer is asking about refunds."))

B - Cross-process interop with A2A

A2A exposes a whole agent over HTTP. AIMU gives you both directions:

  • Serve: serve_a2a(runner) (blocking) or build_a2a_app(runner) (returns a Starlette ASGI app).
  • Consume: RemoteAgent.connect(url) resolves the remote agent card and returns a local Runner.

serve_a2a blocks, so in a notebook we run the ASGI app under uvicorn in a background thread.

import socket
import threading
import time
from contextlib import contextmanager

import uvicorn


@contextmanager
def serve_in_background(app):
    """Run an ASGI app on a free port in a daemon thread; yield its base URL."""
    sock = socket.socket()
    sock.bind(("127.0.0.1", 0))
    port = sock.getsockname()[1]
    sock.close()
    config = uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning")
    server = uvicorn.Server(config)
    thread = threading.Thread(target=server.run, daemon=True)
    thread.start()
    try:
        while not server.started:
            time.sleep(0.02)
        yield f"http://127.0.0.1:{port}/"
    finally:
        server.should_exit = True
        thread.join(timeout=5)

Expose a runner as an A2A server

build_a2a_app wraps the runner in an A2A executor and advertises it via an agent card at /.well-known/agent-card.json. We serve the deterministic FAQRunner here so calls are instant, but any Runner (including a model-backed Agent or Chain) works identically.

from aimu.agents.a2a import build_a2a_app

faq_server = FAQRunner(name="support-faq")
app = build_a2a_app(faq_server, url="http://placeholder/", description="ACME support FAQ agent.")

server_ctx = serve_in_background(app)
base_url = server_ctx.__enter__()  # keep the server alive across cells
print("serving at", base_url)

Inspect the agent card

Any A2A client (AIMU’s or another framework’s) discovers the agent by fetching its card.

import httpx

card = httpx.get(base_url + ".well-known/agent-card.json").json()
print("name:       ", card["name"])
print("description:", card["description"])
print("skills:     ", [s["name"] for s in card["skills"]])
print("streaming:  ", card["capabilities"].get("streaming"))

Consume the remote agent as a local Runner

RemoteAgent.connect(url) resolves the card and returns a Runner. From here it behaves like any local runner: run(), messages, and as_tool() all work, and it slots into workflows.

from aimu.agents import RemoteAgent

remote = RemoteAgent.connect(base_url)
print("connected to:", remote.name)
print("run():       ", remote.run("What is your refund policy?"))
print("messages:    ", remote.messages)

The payoff: a remote agent composes exactly like a local one

Because RemoteAgent is a Runner, the same composition tools from Part A apply with zero A2A-specific wiring: as_tool() into a local agent, or a Chain step.

# 1. As a tool for a local autonomous agent (description comes from the remote card).
remote_tool = remote.as_tool()
print("remote as tool:", remote_tool.__name__, "->", remote_tool("refund"))

# 2. As a step in a local workflow.
pipeline = Chain(agents=[remote])
print("remote in Chain:", pipeline.run("What are your hours?"))

Tear down the server

remote.close()
server_ctx.__exit__(None, None, None)
print("server stopped")

C - Async (aimu.aio)

The async surface mirrors the sync one. aio.a2a.RemoteAgent uses the A2A SDK’s async client natively (no thread portal) and supports real incremental streaming via message/stream. Jupyter allows top-level await.

from aimu.aio.agent import AsyncRunner
from aimu.aio.a2a import build_a2a_app as aio_build_a2a_app


class AsyncFAQRunner(AsyncRunner):
    def __init__(self, name="async-faq"):
        self.name = name
        self.system_message = "Async support FAQ."

    async def run(self, task, generate_kwargs=None, stream=False, images=None):
        return "Refunds take 5 business days." if "refund" in task.lower() else "Ask me about refunds."

    @property
    def messages(self):
        return {self.name: []}


aio_app = aio_build_a2a_app(AsyncFAQRunner(), url="http://placeholder/")
aio_ctx = serve_in_background(aio_app)
aio_url = aio_ctx.__enter__()
print("async server at", aio_url)
from aimu.aio.a2a import RemoteAgent as AsyncRemoteAgent

aio_remote = await AsyncRemoteAgent.connect(aio_url)
print("await run():", await aio_remote.run("refund please"))

# Streaming: message/stream -> StreamChunks
stream = await aio_remote.run("refund please", stream=True)
async for chunk in stream:
    print(chunk.phase.value, repr(chunk.content))

await aio_remote.aclose()
_ = aio_ctx.__exit__(None, None, None)

Serving from the command line

To expose an agent as a standalone service (the agent-level analog of python -m aimu.tools.mcp, which serves tools):

python -m aimu.agents.a2a --model anthropic:claude-sonnet-4-6 \
    --system "You are a helpful research assistant." --port 9000

Then from anywhere: RemoteAgent.connect("http://localhost:9000").