How do I trace a document processing pipeline built with LlamaIndex?

Rius traces a LlamaIndex document processing pipeline with its Python integration, which records retrievers, embeddings and model calls as spans. You add a run root and one span per stage with observe, and set the session and user around each run. Cost is computed by Rius from the model and token counts on the model-call spans.

How it works

The LlamaIndex integration turns each query into a nested trace of typed spans:

• Retrieval steps become RETRIEVER spans. The retriever span’s output shows what the model was given to work with.
• Embedding calls become EMBEDDING spans.
• Synthesis model calls become LLM spans with model, token usage and messages.

The integration covers what happens inside LlamaIndex. It does not know what a stage of your pipeline is, so you add that structure with the tracing helpers:

• A run root. Wrap the function that processes one batch or one document set with observe. Every stage and every LlamaIndex span nests under it, so one run is one trace.
• A span per stage. Wrap each of your own stages, such as loading, parsing, chunking, indexing and answering, with observe. A bare observe makes a CHAIN span named after the function.
• Retrieval you run yourself. If a stage searches a store outside LlamaIndex, open it with rius.start_as_current_span and SpanKind.RETRIEVER. Pass data_source_id and top_k, and record hits with set_retrieved_documents(). That records document ids and scores, which stay visible even when content is stripped.
• Your own tool functions. A function that calls an external service and that no integration already records gets observe with kind=SpanKind.TOOL. A bare observe would make a CHAIN span instead.
• Session and user. rius.session stamps session.id on every span in its scope, and rius.user stamps user.id. Both also reach the spans the integration creates.

If the openai extra is installed too, an LLM call is also traced as a second LLM span under the LlamaIndex span. It is still counted once in tokens and cost. Calls from MultiModalLLM models, embedding models and retrievers are not nested, so their OpenAI span starts a trace of its own.

Set it up

1. Install the integration. The extra installs the instrumentation, not llama-index, so name both in a fresh environment:

pip install "glassflow-rius[llama-index]" llama-index

2. Call init() once at startup. With a default init(), every installed integration turns on automatically. Installation and configuration covers every setting, and Reliability covers how spans are exported.

3. Wrap the run and its stages, then scope the session and user around each run:

import rius
from rius import SpanKind
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
 
rius.init(api_key="ri_...", service_name="doc-pipeline")
 
@rius.observe
def load_documents(path: str):
    return SimpleDirectoryReader(path).load_data()
 
@rius.observe
def build_index(documents):
    return VectorStoreIndex.from_documents(documents)
 
@rius.observe
def answer(index, question: str) -> str:
    return str(index.as_query_engine().query(question))
 
@rius.observe(name="process_batch")
def process_batch(path: str, questions: list[str]):
    documents = load_documents(path)
    index = build_index(documents)
    return [answer(index, q) for q in questions]
 
with rius.session(batch_id), rius.user(requested_by):
    process_batch("./incoming", questions)

Here batch_id and requested_by are your own values.

4. Choose the session id. Pass an id you already have, such as a batch or job id, so a run that spans several traces stays together. For a batch job that handles exactly one session per process, you can set RIUS_SESSION_ID once instead. Do not do that in a server that handles many users, because every user’s traces would merge into one session.

5. Choose the user id. Pass an opaque identifier rather than an email. If you pass no user, spans are anonymous.

6. Run one batch and open the trace in the console. process_batch is the root, your stages sit under it, and the LlamaIndex retriever, embedding and LLM spans nest inside the stages that call them, with the LLM leaf carrying model, tokens and cost.

What you get

• One trace per pipeline run, with a root you named and a span for each stage, so you can see which stage was slow or failed.
• Retrieval, embedding and synthesis spans from LlamaIndex inside those stages, with model, tokens and cost on the model calls.
• Cost per call as a list-price estimate in US dollars, computed by Rius. You send token usage, never cost.
• Runs grouped by session.id, and the people who requested them tied together by user.id.
• Retrieved document ids and scores that stay visible even with content stripped. Retrieved document content and messages are content attributes, so Rius’s privacy controls cover them.

When something else fits

The LlamaIndex integration is Python only. For LlamaIndex.TS no query-engine or retriever structure is captured, but the OpenAI integration still traces provider calls made through the OpenAI client, and manual generations cover the pipeline around them. If your stages do not use LlamaIndex at all, observe and the span helpers on the tracing page cover them directly.

Related