How do I trace a data analysis agent built with Pydantic AI?

Rius traces a Pydantic AI data analysis agent with two lines: call rius.init() and Agent.instrument_all(True). Pydantic AI then sends its own spans for each agent run, model request and tool call, with tokens and cost on every model request. You add observe wrappers for your own stages, a retrieval span for the context you load, and a session and user scope around each request.

How it works

Pydantic AI is OpenTelemetry-native. rius.init() installs the global tracer provider, and Agent.instrument_all(True) tells Pydantic AI to emit its spans through it. No auto-instrumentation extra is involved, and Rius computes cost on the server from the model and token counts. The Pydantic AI integration page covers the full setup.

The framework gives you these spans without any wrapping:

• An agent run becomes the span invoke_agent <agent name>.
• Each model request becomes a generation named chat <model>, with tokens and cost.
• Each tool call becomes execute_tool <tool name>.

That covers the agent itself. A data analysis agent usually does more around the run: it loads context before the model is called and prepares or summarizes results afterwards. Those steps are your own code, so you trace them with observe and span helpers:

• A run root. An observe function that wraps the whole request puts every stage and every Pydantic AI run under one span, named after the function.
• A step per stage. Wrap each stage function (loading data, planning, summarizing) with observe. Each becomes a span, so you see which stage took the time or failed.
• Retrieval. Open a span of kind RETRIEVER around the lookup of documentation, schemas or earlier findings. It records the data source, how many documents were requested, and the id and score of each document returned.

Do not wrap the functions you register as Pydantic AI tools. The framework already records them as execute_tool spans, so a second wrapper would add a duplicate step.

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 names analyst, run_query, find_context, analyze, search and execute below are your own code:

import rius
from rius import SpanKind
from pydantic_ai import Agent
 
rius.init(api_key="ri_...", service_name="analysis-agent")
Agent.instrument_all(True)
 
analyst = Agent(
    "openai:gpt-5-mini",
    instructions="Answer questions about the dataset; use the tool to run queries.",
)
 
@analyst.tool_plain
def run_query(sql: str) -> str:
    """Run a read-only query against the dataset."""
    return execute(sql)
 
@rius.observe
def find_context(question: str) -> list[str]:
    with rius.start_as_current_span(
        kind=SpanKind.RETRIEVER, data_source_id="metrics-docs", top_k=5
    ) as obs:
        obs.set_input(question)
        docs = search(question)
        obs.set_retrieved_documents([{"id": d.id, "score": d.score} for d in docs])
        return [d.text for d in docs]
 
@rius.observe
def analyze(question: str):
    context = find_context(question)
    return analyst.run_sync(f"{question}\n\nContext: {context}").output

3. Scope each request with a session and a user. Reuse your conversation id as the session id, and pass the identifier your application already has for the person the request served:

with rius.session(conversation_id), rius.user(user_id):
    analyze(question)

Every span in the scope carries session.id and user.id, including the spans Pydantic AI emits. Use one session per conversation, so a follow-up question lands in the same session as the first. Prefer an opaque user id to an email or a name, because the value is stored on every span and shown in the console. A span outside any user scope is anonymous, and there is no environment variable for the user.

4. Run one request that calls the tool, then open the trace. You should see your root, the stage and retrieval spans, the invoke_agent span, chat generations with model, tokens and cost, and execute_tool spans:

analyze
├─ find_context
│  └─ retrieval metrics-docs
└─ invoke_agent analyst
   ├─ chat gpt-5-mini
   ├─ execute_tool run_query
   └─ chat gpt-5-mini

What you get

• One trace per request, with your stages, the retrieval step and the framework’s agent run in a single tree.
• Tokens and cost on each chat generation. Costs are list-price estimates in US dollars.
• Retrieved document ids and scores on the retrieval span. These are metadata, so they stay when content capture is off. Put retrieved text in set_output() only if you want it captured as content.
• Failures that point at a stage. An exception in your own code is recorded on that span with ERROR status and the exception class, then re-raised.
• Conversations grouped by session, and the Users page showing each end user your agent served, with their traces, failures and cost.
• Time to first token and reasoning tokens on Pydantic AI runs, with glassflow-rius 1.0.1 and 1.0.2 or later.

Two things to know about Pydantic AI spans. Successful spans arrive with status UNSET rather than OK, while the trace status still reads Ok. What the spans carry (prompts, completions, tool arguments) is controlled by Pydantic AI’s InstrumentationSettings, on top of the SDK’s own privacy controls.

When something else fits

• If your analysis pipeline also calls a model through a client library or another framework outside Pydantic AI, see the integrations overview for the integration that records those calls.
• If you already use Pydantic AI without the Rius SDK, point any OpenTelemetry setup at the ingest endpoint instead. You then set up sessions and exporter wiring yourself.

Related