How do I trace a data analysis agent built with LlamaIndex?

Rius traces a LlamaIndex data analysis agent in Python through its LlamaIndex integration, plus a few observe spans of your own for the run and its stages. The integration records query engines, retrievers, embeddings and model calls with tokens and cost. Your own spans add one root per analysis run and one span per stage, and you set the session and user so traces group by conversation and by person.

How it works

The LlamaIndex integration turns each query into a nested trace of typed spans. Retrieval steps become retriever spans, embedding calls become embedding spans, and synthesis model calls become model-call spans with model, token usage and messages. Rius computes cost from the model and token counts, so you never send a cost value.

A data analysis agent does more than query an index, so three things come from your own code with observe:

• A run root. One span with kind AGENT wraps the whole analysis, so every stage and every LlamaIndex span nests under one trace.
• A span per stage. A plain observe makes a span named after the function, which suits stages such as planning the question or summarizing results.
• Your own tools. A function that runs SQL or a dataframe query is not a LlamaIndex call. Trace it with the TOOL kind, so it shows under the Tool kind in the Spans tab.

If you retrieve from a source outside LlamaIndex, such as a schema catalog, open a retriever span yourself with the data source id and the number of documents you asked for. Document ids and scores stay visible even when content capture is off.

This integration is Python only. See Integrations for the full list.

Set it up

1. Install the LlamaIndex instrumentation. The extra installs the instrumentation, not llama-index itself:

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

2. Call init() once at startup. With a global init(), the installed integration turns on automatically, so build your indexes and query engines as before. Installation and keys are covered in full on the LlamaIndex integration page.

3. Wrap the run, its stages and your own tools, and scope the session and user around each request:

import rius
from rius import SpanKind
 
rius.init(api_key="ri_...", service_name="data-analysis-agent")
 
@rius.observe(kind=SpanKind.TOOL)
def run_sql(query: str) -> str:
    ...  # your own database call
 
@rius.observe
def summarize(rows: str) -> str:
    ...  # your LlamaIndex query engine call, traced by the integration
 
@rius.observe(kind=SpanKind.AGENT)
def analyze(question: str) -> str:
    with rius.start_as_current_span(
        kind=SpanKind.RETRIEVER, data_source_id="schema-catalog", top_k=5
    ) as obs:
        obs.set_input(question)
        tables = ...  # your own lookup
        obs.set_retrieved_documents([{"id": t.id, "score": t.score} for t in tables])
    rows = run_sql(...)
    return summarize(rows)
 
with rius.session(conversation_id):
    with rius.user(request.user_id):
        analyze(question)

4. Replace conversation_id and request.user_id with the ids your application already has. Use an opaque user id rather than an email, because the value is stored on every span and shown in the console.

5. Do not add a manual generation for a model call the integration already records. Leave model calls to the integration, and trace only your own functions.

What you get

• One trace per analysis run, with the AGENT root on top and the retrieval, SQL and summary stages beneath it.
• Retrieval spans from LlamaIndex whose output shows what the model was given to work with, which is where debugging a wrong answer usually starts.
• Model calls with model, tokens and a list-price cost estimate in US dollars, so you see which stage spends the most.
• Your SQL or dataframe tool calls as tool spans, with errors recorded on the span and grouped with other errors.
• Sessions that group the traces of one conversation, and a Users view of what each person’s analyses cost. With no session set, each trace is its own session, and with no user set it counts as anonymous.
• If the openai extra is also installed, the provider call appears as a second model-call span under the LlamaIndex span, and its tokens and cost are counted once.

When something else fits

• If your agent runs on LlamaIndex.TS, there is no LlamaIndex integration in TypeScript. Trace the pipeline with manual spans and generations, and the OpenAI integration still records provider calls.
• If a model call is not covered by any integration, record it as a manual generation with its own token usage, and turn off the integration that would also record it.

Related