Cookbook

Practical recipes for common Axio patterns.

Agent with memory persistence

Save and restore conversation history:

import asyncio
from axio import Agent, MemoryContextStore
from axio.testing import StubTransport, make_text_response


transport = StubTransport([make_text_response("Hello!")])


async def main() -> None:
    context = MemoryContextStore()

    agent = Agent(
        system="You are a helpful assistant.",
        tools=[],
        transport=transport,
    )

    reply = await agent.run("Hi", context)
    print(reply)


asyncio.run(main())

Streaming in FastAPI

Build a web API with streaming events:

from axio import Agent, MemoryContextStore, TextDelta
from axio.testing import StubTransport, make_text_response


transport = StubTransport([make_text_response("Hello!")])


async def stream_events(message: str, agent: Agent, context: MemoryContextStore):
    """Pattern for streaming - yields events."""
    async for event in agent.run_stream(message, context):
        yield event

RAG with custom tools

Combine retrieval and generation:

from axio import Tool


async def retrieve_context(query: str) -> str:
    """Retrieve relevant context from a knowledge base."""
    return f"Results for: {query}"


async def generate_response(context: str, question: str) -> str:
    """Generate a response using retrieved context."""
    return f"Generated: {context[:50]}"


# Create tools
retrieve_tool = Tool(name="retrieve", handler=retrieve_context)
generate_tool = Tool(name="generate", handler=generate_response)

Multi-agent workflow

Coordinate multiple agents:

from axio import Agent, MemoryContextStore
from axio.testing import StubTransport, make_text_response


async def main():
    shared_context = MemoryContextStore()

    transport = StubTransport(
        [
            make_text_response("Research result"),
            make_text_response("Final summary"),
        ]
    )

    research_agent = Agent(
        system="Research the topic.",
        tools=[],
        transport=transport,
        context=shared_context,
    )

    writer_agent = Agent(
        system="Write a summary.",
        tools=[],
        transport=transport,
        context=shared_context,
    )

    research_result = await research_agent.run("What is async?")
    final_result = await writer_agent.run(f"Summary: {research_result}")
    print(final_result)

Retrying another transport

A transport that wraps another one and retries it. Retrying is only safe before the first event reaches the caller. Once a delta has been yielded the turn is half-delivered, and starting over repeats it.

import asyncio
from collections.abc import AsyncIterator
from typing import Any
from axio import (
    Agent,
    MemoryContextStore,
    StopReason,
    StreamEvent,
    TextDelta,
    Tool,
    Usage,
)
from axio.events import IterationEnd
from axio.exceptions import StreamError
from axio.messages import Message


class RetryTransport:
    """Retries the wrapped transport on StreamError, with exponential backoff."""

    def __init__(
        self, inner: Any, max_retries: int = 3, base_delay: float = 0.01
    ) -> None:
        self.inner = inner
        self.max_retries = max_retries
        self.base_delay = base_delay

    async def stream(
        self,
        messages: list[Message],
        tools: list[Tool[Any]],
        system: str,
    ) -> AsyncIterator[StreamEvent]:
        for attempt in range(1, self.max_retries + 1):
            delivered = False
            try:
                async for event in self.inner.stream(messages, tools, system):
                    delivered = True
                    yield event
                return
            except StreamError:
                # Nothing to retry once the caller has seen part of the turn.
                if delivered or attempt == self.max_retries:
                    raise
                await asyncio.sleep(self.base_delay * 2 ** (attempt - 1))


class Flaky:
    """Fails twice, then answers."""

    def __init__(self) -> None:
        self.calls = 0

    def stream(
        self,
        messages: list[Message],
        tools: list[Tool[Any]],
        system: str,
    ) -> AsyncIterator[StreamEvent]:
        self.calls += 1
        return self._answer()

    async def _answer(self) -> AsyncIterator[StreamEvent]:
        if self.calls < 3:
            raise StreamError("503 Service Unavailable")
        yield TextDelta(index=0, delta="ok")
        yield IterationEnd(
            iteration=0, stop_reason=StopReason.end_turn, usage=Usage(1, 1)
        )


async def main() -> None:
    flaky = Flaky()
    agent = Agent(system="", transport=RetryTransport(flaky))
    assert await agent.run("hi", MemoryContextStore()) == "ok"
    assert flaky.calls == 3


asyncio.run(main())

The shipped transports retry inside themselves rather than through a wrapper. They prefer the Retry-After header over their own backoff when the response carries one. See Writing Transports for what a transport is expected to do with errors.

Rate limiting tool

import asyncio
from axio import Tool, CONTEXT

RATE_LIMIT = 10
TIME_WINDOW = 60


async def rate_limited_action(data: str) -> str:
    """Tool with rate limiting."""
    calls: list[float] = CONTEXT.get()
    now = asyncio.get_event_loop().time()
    # Prune old calls outside the window
    calls[:] = [t for t in calls if now - t < TIME_WINDOW]
    if len(calls) >= RATE_LIMIT:
        raise RuntimeError(f"Rate limit: {RATE_LIMIT}/{TIME_WINDOW}s")
    calls.append(now)
    return "done"


call_log: list[float] = []
tool = Tool(name="rate_limited_action", handler=rate_limited_action, context=call_log)

API key guard

Check for required environment variables:

import os
from typing import Any
from axio import PermissionGuard, GuardError


class ApiKeyGuard(PermissionGuard):
    """Ensure required environment variables are set."""

    required_keys = ("OPENAI_API_KEY",)

    async def check(self, handler: Any) -> Any:
        missing = [k for k in self.required_keys if not os.environ.get(k)]
        if missing:
            raise GuardError(f"Missing: {', '.join(missing)}")
        return handler

Tool with guards

Apply guards to specific tools:

from typing import Any
from axio import Tool, PermissionGuard


async def sensitive_operation(data: str) -> str:
    """Process sensitive data."""
    return f"Processed: {data}"


class AllowGuard(PermissionGuard):
    async def check(self, tool: Any, **kwargs: Any) -> dict[str, Any]:
        return kwargs


sensitive_tool = Tool(
    name="sensitive_operation",
    handler=sensitive_operation,
    guards=(AllowGuard(),),
)

assert len(sensitive_tool.guards) == 1