How do I trace a RAG chatbot built with Pydantic AI?

Rius traces a Pydantic AI RAG chatbot by reading the framework’s own OpenTelemetry spans for the agent run, each model request and each tool call. You add a few spans of your own for retrieval and the other stages, and scope each conversation with a session and each person with a user. Model calls arrive with tokens and cost, which Rius computes.

How it works

Pydantic AI is OpenTelemetry-native. rius.init() installs the global tracer provider and Agent.instrument_all(True) tells the framework to emit its spans through it. No extra and no wrappers are needed for the framework’s part. See the Pydantic AI integration page for details. In the trace:

• An agent run is the root span, named invoke_agent followed by the agent’s name.
• A model request is a generation named chat followed by the model, with tokens and cost.
• A tool call is a span named execute_tool followed by the tool’s name.

A RAG chatbot needs more than that. Pydantic AI does not know about your vector search, your reranking or the stages around the agent run, so those are spans you add with observe and span helpers from Tracing your code:

• A retrieval span of kind RETRIEVER. It records the data source, how many documents you asked for and the ids and scores of what came back.
• A rerank span of kind RERANKER, if you reorder documents.
• A run root, a plain observe function around the whole turn. It is a CHAIN span named after your function and holds retrieval, reranking and the agent run as steps of one turn.

If the agent retrieves by calling a tool, the framework already records the execute_tool span. Open the retrieval span inside that tool function to add the data source and the documents. Don’t wrap that tool in observe again.

Don’t add manual generations for the model calls. The framework already produces them with tokens and cost.

Set it up

1. Install the SDK and Pydantic AI. No extra is needed. The integrations overview covers the other frameworks and clients.

pip install glassflow-rius pydantic-ai

2. Initialize Rius and turn on Pydantic AI’s instrumentation once at startup, before you create agents:

import rius
from pydantic_ai import Agent
 
rius.init(api_key="ri_...", service_name="my-agent")
Agent.instrument_all(True)

3. Wrap each turn in a run root, open a RETRIEVER span around the search, and trace the rerank step. The names search, product-kb and the prompt format are placeholders for your own:

import rius
from rius import SpanKind
from pydantic_ai import Agent
 
chatbot = Agent(
    "openai:gpt-5-mini",
    instructions="Answer from the provided context.",
)
 
@rius.observe(name="rerank", kind=SpanKind.RERANKER, capture_output=False)
def rerank(docs):
    ...
 
@rius.observe
def answer_turn(question: str):
    with rius.start_as_current_span(
        kind=SpanKind.RETRIEVER, data_source_id="product-kb", 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])
    ranked = rerank(docs)
    return chatbot.run_sync(f"Context: {ranked}\n\nQuestion: {question}").output

set_retrieved_documents() records ids and scores, not document text, so it survives content stripping. To capture the retrieved text as content, pass it to set_output() as well.

4. Scope each turn with the conversation and the user, so the framework’s spans and yours all carry them:

with rius.session(conversation_id), rius.user(user_id):
    answer_turn(message)

Use your conversation id for the session so all turns of one chat group together. For the user, prefer an opaque id to an email. Don’t set a process-wide session in a chatbot server that handles many users, because every conversation would merge into one.

5. Run one question that retrieves and answers, then open the trace in the console.

What you get

• One trace per turn. The answer_turn root holds the retrieval, the rerank, and the invoke_agent run with its chat generations and any execute_tool spans.
• Tokens, time to first token and reasoning tokens on the model calls, and a cost per call. Cost is a list-price estimate in US dollars that Rius computes, so you never send it.
• Retrieval spans named retrieval product-kb, so you can filter by kind in the Spans tab and see which data source and how many documents each turn asked for.
• Turns grouped into a session per conversation, and the user’s traces, failures and cost grouped by user.id.
• Content in the framework’s spans (prompts, completions, tool arguments) is controlled by Pydantic AI’s InstrumentationSettings, for example include_content, on top of the SDK’s privacy controls.

When something else fits

• If retrieval and the agent run are a single tool call, the framework’s own spans may be enough, so skip the run root and the extra retrieval span.
• If you use Pydantic AI without the Rius SDK, point any OpenTelemetry setup at the ingest endpoint instead. You then set up sessions and privacy controls yourself.

Related