How do I trace a research agent built with Pydantic AI?

Rius traces a Pydantic AI research agent with two lines: rius.init() and Agent.instrument_all(True). Pydantic AI then emits a span per agent run, one generation per model request and one span per tool call. You add a few observe wrappers for the research stages, a retrieval span for each search, and a session and user scope for each request.

How it works

Pydantic AI is OpenTelemetry-native, so no auto-instrumentation extra is involved. rius.init() installs the global tracer provider, and Agent.instrument_all(True) tells the framework to emit its spans through it. Rius reads the framework’s GenAI attributes and computes cost from the model and token counts.

What the framework gives you:

• Each agent run is a span named invoke_agent followed by the agent name. It is the root unless you wrap it.
• Each model request is a generation named chat followed by the model, with tokens and cost.
• Each tool call is a span named execute_tool followed by the tool name. You do not wrap tools yourself.

What a research agent needs on top is structure the framework does not know about:

• A run root. One observe function around the whole research job gives you a single parent for every stage. A bare observe makes a span of the default CHAIN kind, named after the function.
• A step per stage. Wrap each stage, such as gathering sources and writing the report, in its own observe function. Each agent run then sits under the stage that started it.
• A retrieval span for each search. start_as_current_span with SpanKind.RETRIEVER records the data source and how many documents you asked for. set_retrieved_documents records document ids and scores, not text. Put retrieved text in set_output if you want it captured as content.

See Tracing your code for the span kinds and naming rules.

Set it up

1. Install the SDK and Pydantic AI. No extra is needed.

pip install glassflow-rius pydantic-ai

2. Initialize Rius, turn on Pydantic AI instrumentation, and add your own spans. The tool, the stage functions and the scopes below are examples. Replace search, format_docs and the ids with your own.

import rius
from rius import SpanKind
from pydantic_ai import Agent
 
rius.init(api_key="ri_...", service_name="research-agent")
Agent.instrument_all(True)
 
research_agent = Agent(
    "openai:gpt-5-mini",
    instructions="Research the question. Search the knowledge base before you answer.",
)
 
@research_agent.tool_plain
def search_kb(query: str) -> str:
    """Search the knowledge base."""
    with rius.start_as_current_span(
        kind=SpanKind.RETRIEVER, data_source_id="product-kb", top_k=5
    ) as obs:
        obs.set_input(query)
        docs = search(query)
        obs.set_retrieved_documents([{"id": d.id, "score": d.score} for d in docs])
        return format_docs(docs)
 
@rius.observe
def gather(question: str) -> str:
    result = research_agent.run_sync(f"Collect sources for: {question}")
    return result.output
 
@rius.observe
def write_up(notes: str) -> str:
    result = research_agent.run_sync(f"Write a report from these notes: {notes}")
    return result.output
 
@rius.observe
def research(question: str) -> str:
    return write_up(gather(question))
 
with rius.session(conversation_id), rius.user(user_id):
    report = research(question)

3. Set the session and user around each request, as in the last two lines. rius.session() groups the traces of one conversation or research job under your own id. rius.user() ties them to the person you served. The framework’s spans inherit both scopes like any other span. Prefer an opaque user id to an email.

4. Run one request that searches, then open the trace in the console. You should see your research root, the gather and write_up steps, invoke_agent spans under them, chat generations with model, tokens and cost, execute_tool spans, and the retrieval span for each search.

The full framework setup is on the Pydantic AI integration page.

What you get

• One trace per research request, with your own stage names at the top and the framework’s agent, model and tool spans beneath.
• Tokens and cost on every model request, as list-price estimates in US dollars that Rius computes. You never send a cost value.
• Time to first token and reasoning tokens on Pydantic AI runs.
• Retrieval spans that show which data source each search hit and which documents came back with their scores.
• Sessions that group a multi-turn research conversation, and a Users view that groups everything one person asked.

Two behaviours to know. Successful Pydantic AI spans arrive with an UNSET status, though the trace status still reads Ok. What the spans carry (prompts, completions, tool arguments) is set by Pydantic AI’s InstrumentationSettings, such as include_content, on top of the SDK’s own privacy controls.

When something else fits

• If your research agent is written in TypeScript, there is no Pydantic AI integration. Wrap your own calls with manual generations and spans instead.
• If you do not want the Rius SDK, any OpenTelemetry setup can send the framework’s spans to the ingest endpoint. You then set sessions and user ids yourself.
• For other frameworks, see the integrations list.

Related