How do I trace a RAG chatbot built with AutoGen?

Rius traces an AutoGen RAG chatbot with the AutoGen integration, one root span per turn, a retrieval span, and a session and user set around each turn. AutoGen’s own spans give you the agent, tool and model-call structure, and the openai integration adds model, tokens and cost. You add only the parts AutoGen does not emit: the run root, the retrieval step and the identifiers.

How it works

The AutoGen integration has two layers, both on once you call rius.init() (Python only):

• Framework spans. AutoGen emits create_agent, invoke_agent and execute_tool spans through the global tracer provider that rius.init() installs.
• Generations. Model calls go through the OpenAI client from autogen-ext[openai], so the openai integration records each one as a ChatCompletion generation with model, tokens and cost. Rius computes the cost, so you never send one.

For a RAG chatbot, three things sit on top, using observe and span helpers:

• A run root. One observe wrapper on your turn handler, with kind AGENT, groups agent creation and invocation into a single trace per turn.
• A retrieval span. A span of kind RETRIEVER records the knowledge base searched (data_source_id), how many documents you asked for (top_k) and the ids and scores of what came back. It carries identifiers and scores, not document text. Pass retrieved text to set_output() if you want it captured as content.
• A step per stage. A reranking function you write gets observe with kind RERANKER, so it shows as its own step.

Do not wrap AutoGen tool functions with observe. AutoGen runs tools in a detached context, so the wrapper’s span cannot find the run and splits off as its own single-span trace. If your agent retrieves through a tool, the framework’s execute_tool span already covers it. Run retrieval you want a RETRIEVER span for in your turn handler, before the agent runs.

Set it up

1. Install the SDK with the openai extra, plus AutoGen. The extra installs the instrumentation, not the libraries:

pip install "glassflow-rius[openai]" "autogen-agentchat" "autogen-ext[openai]"

2. Initialize Rius once, and wrap the turn handler. search, rerank, the document fields and the names are your own:

import rius
from autogen_agentchat.agents import AssistantAgent
from autogen_ext.models.openai import OpenAIChatCompletionClient
from rius import SpanKind
 
rius.init(api_key="ri_...", service_name="rag-chatbot")
 
@rius.observe(name="rerank", kind=SpanKind.RERANKER, capture_output=False)
def rerank(docs):
    return sorted(docs, key=lambda d: d.score, reverse=True)  # your reordering logic
 
@rius.observe(name="handle-turn", kind=SpanKind.AGENT)
async def handle_turn(question: str) -> str:
    with rius.start_as_current_span(
        kind=SpanKind.RETRIEVER, data_source_id="product-kb", top_k=5
    ) as obs:
        obs.set_input(question)
        docs = search(question)  # your retrieval call
        obs.set_retrieved_documents([{"id": d.id, "score": d.score} for d in docs])
 
    context = "\n\n".join(d.text for d in rerank(docs))
    agent = AssistantAgent(
        name="rag_agent",
        model_client=OpenAIChatCompletionClient(model="gpt-5-mini"),
        system_message="Answer from the context you are given.",
    )
    result = await agent.run(task=f"Context:\n{context}\n\nQuestion: {question}")
    return str(result.messages[-1].content)

3. Set the session and the user around each turn. A session is one conversation, so reuse your conversation id. A user is the person served, so pass an opaque id rather than an email:

async def answer(question: str, conversation_id: str, user_id: str) -> str:
    with rius.session(conversation_id), rius.user(user_id):
        return await handle_turn(question)

Every span of the turn carries session.id and user.id, including the spans AutoGen and the openai integration create. Call this in your request handler, not once for the process, so each user and conversation stays separate.

4. Send one question that retrieves, then open the trace in the console and check for the root span, the retrieval span, the invoke_agent span and the ChatCompletion generation with model, tokens and cost.

What you get

One trace per turn that reads in the order the chatbot works:

handle-turn                          AGENT
├─ retrieval product-kb              RETRIEVER
├─ rerank                            RERANKER
├─ create_agent rag_agent
└─ invoke_agent rag_agent
   └─ ChatCompletion                 LLM

• Tokens and a list-price cost estimate in US dollars on each model call.
• The ids and scores of the documents each question retrieved, so you can see which turns pulled weak matches.
• Turns grouped into one session per conversation, and traces, failures and cost grouped per user.
• If your agent calls a tool, an execute_tool span for it, from the framework.

Two AutoGen quirks to expect: framework spans arrive with status UNSET on success, though the trace status still reads Ok, and generation spans are named ChatCompletion.

When something else fits

• If your model calls do not go through the OpenAI client, the openai integration does not capture them. Trace them as manual generations and set their usage yourself, as the tracing guide shows.
• AutoGen has no TypeScript integration today. For other frameworks and languages, see the integrations overview.

Related