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

Rius traces an AutoGen data analysis agent with the AutoGen integration, and you add a few spans of your own for retrieval, the run root and each stage. AutoGen emits its own spans for agent creation, agent runs and tool calls, and the openai integration records each model call with tokens and cost. You wrap the entrypoint in observe, open a retriever span around your lookup, and scope each run with a session and a user.

How it works

A data analysis agent is an AutoGen agent plus your own steps around it, so the trace has two layers.

• What AutoGen and the integration give you. After rius.init(), AutoGen’s create_agent, invoke_agent and execute_tool spans arrive on their own. Model calls go through OpenAIChatCompletionClient, so the openai integration records each one as a ChatCompletion generation with model, tokens and cost. Rius computes cost from the model and token counts, so you never send it. The AutoGen integration page covers this layer in full.
• A run root. One observe wrapper with kind=SpanKind.AGENT on your entrypoint groups agent creation and invocation into one trace per question.
• A retrieval span. If the agent looks up schema notes, metric definitions or past reports before it queries, open a RETRIEVER span with data_source_id and top_k. It records document ids and scores, not document text.
• A step per stage. Each stage you write yourself, such as loading context or building the report, gets its own observe wrapper. A bare observe makes a span named after the function.
• Tools stay unwrapped. Do not wrap the tool functions AutoGen runs, such as a SQL query tool. AutoGen executes tools in a detached context, so the wrapper splits off as its own single-span trace. The framework’s execute_tool span already covers each tool call.
• Session and user. rius.session() stamps session.id and rius.user() stamps user.id on every span started in scope. See Tracing your code for both.

Prompts, responses and tool results can contain your data. Turn off capture_content, or set RIUS_CAPTURE_CONTENT to false, to keep that content from leaving the process.

Set it up

1. Install the SDK with the openai extra and the AutoGen packages. The integrations overview lists every other extra.

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

2. Keep the retrieval and each stage in your own code, outside the tools AutoGen runs. Give the retrieval a RETRIEVER span and each stage an observe wrapper.

3. Wrap the entrypoint as the run root, and scope the session and user around the call. search, warehouse and build_report_text stand for your own code.

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="analysis-agent")
 
def run_query(sql: str) -> str:
    """Run a read-only SQL query and return the rows as text."""
    return warehouse.execute(sql)  # a tool: do not wrap it with observe
 
@rius.observe(name="load-context")
def load_context(question: str) -> str:
    with rius.start_as_current_span(
        kind=SpanKind.RETRIEVER, data_source_id="schema-docs", top_k=5
    ) as obs:
        obs.set_input(question)
        docs = search(question)
        obs.set_retrieved_documents([{"id": d.id, "score": d.score} for d in docs])
    return "\n".join(d.text for d in docs)
 
@rius.observe(name="build-report")
def build_report(answer: str) -> str:
    return build_report_text(answer)
 
@rius.observe(name="analyze-data", kind=SpanKind.AGENT)
async def analyze(question: str) -> str:
    context = load_context(question)
    agent = AssistantAgent(
        name="analyst",
        model_client=OpenAIChatCompletionClient(model="gpt-5-mini"),
        tools=[run_query],
        system_message="Answer data questions with the query tool.\n" + context,
    )
    result = await agent.run(task=question)
    return build_report(str(result.messages[-1].content))
 
async def handle(question: str, conversation_id: str, user_id: str) -> str:
    with rius.session(conversation_id):
        with rius.user(user_id):
            return await analyze(question)

4. Pass your own conversation id to rius.session() so one conversation is one session. Use an opaque user id rather than an email.

5. Run one question that calls the query tool, then open the trace in the console and check that every span below is there.

What you get

One trace per question, shaped like this:

analyze-data                         AGENT
├─ load-context                      CHAIN
│  └─ retrieval schema-docs          RETRIEVER
├─ create_agent analyst
├─ invoke_agent analyst
│  ├─ ChatCompletion                 LLM
│  └─ execute_tool run_query
└─ build-report                      CHAIN

• Model, tokens and cost on each ChatCompletion generation, with cost as a list-price estimate in US dollars.
• Each tool call as an execute_tool span, and each retrieval with the ids and scores of the documents it returned.
• Your own stages as named spans, so a slow or failing step is easy to find.
• Traces grouped by conversation through session.id, and by person through user.id.

AutoGen’s own spans arrive with an unset status on success, and the generation spans are all named ChatCompletion. The trace status still reads Ok.

When something else fits

• If your agent calls models through a client other than OpenAI, or is built on another framework, pick the matching page from the integrations overview.
• If one process serves a single job, such as a nightly batch run, you can set a process-wide session id instead of scoping each call. For a server with many users, scope each request.

Related