How do I trace a RAG chatbot built with LlamaIndex?

Rius traces a LlamaIndex RAG chatbot once you install the LlamaIndex extra and call rius.init(): retrieval, embedding and model calls arrive as typed spans with tokens and cost. You then add a few observe wrappers for the run root and each stage, and scope the session and user around each chat turn. The LlamaIndex integration is Python only.

How it works

The llama-index instrumentation turns each query into a nested trace of typed spans:

• Retrieval steps appear as RETRIEVER spans. The retriever span’s output shows what the model was given to work with, which is where RAG debugging usually starts.
• Embedding calls appear as EMBEDDING spans.
• Synthesis model calls appear as LLM spans with model, token usage and messages. Rius computes the cost from those.

The integration covers what happens inside LlamaIndex. The structure around it comes from you: a root span for each chat turn, and one span for each stage of your own code, such as rewriting the question or post-processing the answer. Wrap those functions with observe. By default each one becomes a span named after the function, and the LlamaIndex spans nest underneath it.

If part of your retrieval runs outside LlamaIndex, for example a direct vector store lookup, open a RETRIEVER span yourself with rius.start_as_current_span(kind=SpanKind.RETRIEVER, data_source_id="product-kb", top_k=5). Record what came back with set_retrieved_documents(). It stores document ids and scores, not text, so it stays visible when content capture is off. Put the retrieved text in set_output() if you want it captured as content.

A chatbot spans many traces, one for each turn. A session groups the turns of one conversation under a session.id, and a user ties together everything one person did across conversations under a user.id. You supply both. With no session, each trace counts as its own session, and with no user, spans are anonymous.

If the openai extra is installed too, the model call is traced as a second LLM span under the LlamaIndex span. From glassflow-rius 1.1.1 the call is counted once in tokens and cost, and the LlamaIndex span still shows the total. Calls from embedding models, retrievers and MultiModalLLM models are not nested: their OpenAI span starts a trace of its own. If that matters, limit integrations with instruments=["llama-index"].

Set it up

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

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

2. Call rius.init() once at startup, before you build indexes and query engines. Build them exactly as before. The LlamaIndex integration page has the full setup and what gets captured.

3. Wrap the run root and each stage of your own code with observe.

4. Scope the session and the user around each chat turn, using the conversation id and the user id you already have. Prefer an opaque user id to an email.

import rius
from rius import SpanKind
 
rius.init(api_key="ri_...", service_name="my-agent")
 
@rius.observe
def rewrite_question(question: str) -> str:
    ...  # your own stage, for example resolving follow-up questions
 
@rius.observe(name="answer-turn", kind=SpanKind.AGENT)
def answer_turn(question: str) -> str:
    standalone = rewrite_question(question)
    response = query_engine.query(standalone)  # your LlamaIndex query engine
    return str(response)
 
def handle_message(conversation_id: str, user_id: str, message: str) -> str:
    with rius.session(conversation_id), rius.user(user_id):
        return answer_turn(message)

query_engine, rewrite_question and the ids are your own names. Every span of the turn, including the ones LlamaIndex creates, carries session.id and user.id.

5. Run one question, then open the trace in the console. You see answer_turn as the root, your stage spans, and the retriever, embedding and model spans nested inside.

Do not set a process-wide session with RIUS_SESSION_ID in a chatbot server that handles many users, because every user’s traces would merge into one session. Scope each request instead. There is no environment variable for the user.

What you get

• One trace per chat turn, with your root span at the top and your own stages and LlamaIndex’s retrieval, embedding and model spans nested beneath it.
• Tokens and a list-price cost estimate in US dollars for each model call, computed by Rius from the model and token counts.
• The retriever span’s output, so you can check what context the model received when an answer looks wrong.
• Turns grouped by conversation through the session, and each person’s traces, failures and cost through the user.
• Document ids and scores that stay visible even when privacy controls strip content.

For the other frameworks and client libraries Rius supports, see Integrations.

When something else fits

• If your chatbot is built with LlamaIndex.TS, there is no LlamaIndex integration in TypeScript. Trace the pipeline with the manual spans and generations on the tracing page, and the OpenAI integration still traces the provider calls it makes.
• If you use another framework or client library, open its page from the Integrations list instead.

Related