How do I trace a document processing pipeline built with Pydantic AI?

Rius traces a Pydantic AI document pipeline by reading the spans the framework emits for each agent run, model request and tool call, and you add a few spans of your own for the pipeline around them. Call rius.init() and Agent.instrument_all(True), then wrap the pipeline in one observe function so every stage lands in one trace. Add a retrieval span where you look documents up, and scope each run with a session and a user. The result is one trace per document, with tokens and cost on every model call.

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) makes the framework emit its spans through it. The Pydantic AI integration gives you:

• A span named invoke_agent <agent name> for each agent run, which is the trace root only when nothing wraps it.
• A generation named chat <model> for each model request, with tokens and cost. Rius computes the cost, so you never send it.
• A span named execute_tool <tool name> for each tool call.

A pipeline usually has more than one agent run: for example one agent extracts fields, another classifies, another summarizes. Each run is its own root unless something encloses them, so the framework alone gives you separate pieces. What the pipeline needs on top comes from tracing your code:

• A run root. An observe function around the whole pipeline. A bare observe makes a span named after the function, and the agent runs inside it nest beneath it.
• A step per stage. An observe function for each non-model stage, such as parsing or chunking, so slow or failing stages show up by name.
• A retrieval span. start_as_current_span with SpanKind.RETRIEVER for lookups. It is named retrieval <data source id> and records top_k and the retrieved document ids and scores. It stores no document text. Put retrieved text in set_output() if you want it captured.
• A rerank step. If you reorder retrieved documents, use kind=SpanKind.RERANKER.

You do not wrap the tools registered on your agents. The framework already records them as execute_tool spans.

Set it up

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

   pip install glassflow-rius pydantic-ai

2. Initialize Rius once at startup and turn on Pydantic AI instrumentation:

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

3. Wrap each stage and the whole run. The names below (process_document, parse, fetch_templates, extract_agent, search) are your own:

   from rius import SpanKind
    
   extract_agent = Agent(
       "openai:gpt-5-mini",
       instructions="Extract the key fields from the document.",
   )
    
   @rius.observe
   def parse(raw: bytes) -> str:
       ...
    
   def fetch_templates(text: str):
       with rius.start_as_current_span(
           kind=SpanKind.RETRIEVER, data_source_id="templates", top_k=5
       ) as obs:
           obs.set_input(text[:200])
           docs = search(text)
           obs.set_retrieved_documents(
               [{"id": d.id, "score": d.score} for d in docs]
           )
           return docs
    
   @rius.observe
   def extract(text: str, docs):
       return extract_agent.run_sync(text).output
    
   @rius.observe
   def process_document(raw: bytes):
       text = parse(raw)
       docs = fetch_templates(text)
       return extract(text, docs)

4. Scope each run with a session and a user. Use your own job or document id for the session, and an opaque id for the user:

   with rius.user(request.user_id):
       with rius.session(job_id):
           process_document(raw)

   One session groups the traces of related work, such as every document in a batch job. With no session, each trace is its own session. A span outside a user scope is anonymous. observe takes no user argument because a user belongs to each request. Do not set a process-wide session in a server that handles many users, since their traces would merge.

5. Run one document, then open the trace. You should see the process_document root with parse, the retrieval span, extract, and under it invoke_agent, chat generations with model, tokens and cost, and any execute_tool spans.

What you get

• One trace per document, with every stage as a named span and the model calls nested under the stage that made them.
• Tokens and cost on each chat generation. Costs are list-price estimates in US dollars.
• Time to first token and reasoning tokens for Pydantic AI runs, on glassflow-rius 1.0.1 and 1.0.2 or later.
• A retrieval span that shows which index was searched, how many documents were requested and which ids came back with what score.
• Traces grouped by session and filterable by user, so you can follow a batch job or see what one person’s documents cost.

Two behaviors to know. Successful framework spans arrive with status UNSET, and the trace status still reads Ok. What prompts and completions the spans carry is set by Pydantic AI’s InstrumentationSettings (for example include_content), on top of the SDK’s own privacy controls.

When something else fits

• If you want to send Pydantic AI spans without the Rius SDK, point any OpenTelemetry setup at the ingest endpoint instead. You then set up sessions and exporter wiring yourself.
• If a stage calls a model directly rather than through a Pydantic AI agent, trace that call with a manual generation, as the tracing page shows.

Related