How do I trace a research agent built with LlamaIndex?

Rius traces a LlamaIndex research agent with one install extra and init(), then you add observe spans for the run root and each stage. Retrieval, embedding and model calls inside LlamaIndex are captured for you. The session and user you set tie each run to a conversation and a person. The integration is Python only.

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 model-call spans with the model, token usage and messages. Rius computes the cost on its side from the model and token counts, so you never send a cost.

If the openai extra is also installed, a chat or completion call shows up as a second model-call span under the LlamaIndex span. With glassflow-rius 1.1.1 or later, the call 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.

What the integration does not know is the shape of your research agent: where a run starts, what the stages are, and which tools it calls outside LlamaIndex. You add that with observe and span helpers from Tracing your code:

• A run root: observe with the AGENT kind, so every stage nests under one span for the whole run.
• A step per stage: a bare observe on each stage function, such as planning, gathering and writing. It makes a span named after the function.
• Retrieval you run yourself, such as a search outside a LlamaIndex retriever: a span of the RETRIEVER kind with the data source and top_k.
• Tools you call yourself: observe with the TOOL kind. A bare observe makes a CHAIN span, which the Tool kind in the Spans tab’s All kinds menu and the Tool Loop alert never see.

Set it up

1. Install the instrumentation and call init(). The extra installs the instrumentation, not llama-index itself. Indexes and query engines are built exactly as before. The LlamaIndex integration page covers setup in full.

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

import rius
 
rius.init(api_key="ri_...", service_name="my-agent")
# build indexes and query engines exactly as before

2. Wrap the run and its stages. The names below are your own. The AGENT span is the run root, and each stage function gets its own span underneath. LlamaIndex’s retriever and model spans nest inside whichever stage calls them.

import rius
from rius import SpanKind
 
@rius.observe(kind=SpanKind.AGENT, agent_name="research-agent")
def run_research(question: str) -> str:
    queries = plan_queries(question)
    notes = gather(queries)
    return write_report(question, notes)
 
@rius.observe
def plan_queries(question: str) -> list[str]:
    ...
 
@rius.observe
def gather(queries: list[str]) -> list[str]:
    ...
 
@rius.observe
def write_report(question: str, notes: list[str]) -> str:
    ...

3. Trace retrieval and tools that LlamaIndex does not record. Inside gather, a search of your own gets a retriever span. set_retrieved_documents() stores ids and scores, not document text, so it stays visible when content is stripped. Put retrieved text in set_output() if you want it captured as content.

with rius.start_as_current_span(
    kind=SpanKind.RETRIEVER, data_source_id="papers", top_k=5
) as obs:
    obs.set_input(query)
    docs = search(query)
    obs.set_retrieved_documents([{"id": d.id, "score": d.score} for d in docs])

For a tool function that no integration records, pass the TOOL kind:

@rius.observe(kind=SpanKind.TOOL, tool_name="fetch_page")
def fetch_page(url: str) -> str:
    ...

4. Set the session and the user around each run. A session groups the traces of one conversation or run, and the user is the person the run served. Use your own ids, and prefer an opaque user id to an email. In a server that handles many users, scope each request rather than setting a process-wide session, or every user’s traces merge into one session. observe takes no user argument, so scope the user around the call.

with rius.session(conversation_id), rius.user(request.user_id):
    report = run_research(question)

Then run one question and open the trace in the console: you see the agent span, the stage spans, the retriever spans and the model calls with tokens and cost.

What you get

• One trace per research run, with the run root, a span per stage and LlamaIndex’s retrieval, embedding and model-call spans nested inside.
• Tokens and a list-price cost estimate in US dollars on each model call, computed by Rius.
• Retrieved document ids and scores on retriever spans, so you can check what each step was given.
• Runs grouped by session.id, and each end user’s traces, failures and cost on the Users page through user.id.
• Your own tools visible by kind in the Spans tab, and counted by the Tool Loop alert.

When something else fits

The LlamaIndex integration is Python only. For LlamaIndex.TS, the OpenAI integration still traces provider calls, and manual generations cover the pipeline around them. For the other frameworks and libraries Rius integrates with, see Integrations.

Related