How do I trace a research agent built with AutoGen?

Rius traces an AutoGen research agent with the openai integration and AutoGen’s own spans, and you add a run root, a retrieval span and a span per stage with observe. AutoGen’s create_agent, invoke_agent and execute_tool spans arrive on their own, and each model call arrives as a generation with tokens and cost. You then set the session and user so you can follow one conversation or one person across runs.

How it works

AutoGen gets coverage from two sources, described on the AutoGen integration page:

• Framework spans. AutoGen emits create_agent, invoke_agent and execute_tool spans through the global tracer provider that rius.init() installs. Agent and tool structure arrives with no wrappers.
• Model calls. AutoGen calls the model through the OpenAI client, so the openai integration records each call as a ChatCompletion generation with model, tokens and cost. Rius computes the cost from the model and token counts.

A research agent usually does more than the agent loop, such as planning, searching a knowledge base and writing up findings. AutoGen does not know about those stages, so you add them with the helpers on the tracing your code page:

• A run root. Wrap the entrypoint with observe and kind=SpanKind.AGENT. Agent creation and invocation then group into one trace per research question.
• A step per stage. A bare observe on a stage function such as planning makes a span named after that function. Name stages you want to find later.
• A retrieval span. Open a span of kind RETRIEVER around your own search code, with data_source_id for the index and top_k for how many documents you asked for. set_retrieved_documents() records ids and scores, not document text. Put the text in set_output() if you want it captured as content.

Do not wrap AutoGen tool functions with observe. The framework runs tools in a detached context, so the wrapper’s span cannot find the run and splits off as its own single-span trace. The framework’s execute_tool spans already cover your tools, including a search the model calls as a tool.

Set it up

1. Install the SDK with the openai extra, plus AutoGen. The extra installs the instrumentation, not the OpenAI library itself.

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

2. Initialize Rius once at startup, then add the run root, the stage steps and the retrieval span. Here search, draft_outline and download stand for your own code, and the key and service name are placeholders:

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="research-agent")
 
def fetch_page(url: str) -> str:
    """Fetch a web page."""
    return download(url)
 
@rius.observe
def plan(question: str) -> str:
    return draft_outline(question)
 
def retrieve(question: str) -> list:
    with rius.start_as_current_span(
        kind=SpanKind.RETRIEVER, data_source_id="papers", 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 docs
 
@rius.observe(name="research-run", kind=SpanKind.AGENT)
async def research(question: str) -> str:
    outline = plan(question)
    docs = retrieve(question)
    agent = AssistantAgent(
        name="researcher",
        model_client=OpenAIChatCompletionClient(model="gpt-5-mini"),
        tools=[fetch_page],
        system_message="Write a sourced answer from the outline and documents.",
    )
    sources = "\n\n".join(d.text for d in docs)
    result = await agent.run(task=f"{question}\n\nOutline:\n{outline}\n\nSources:\n{sources}")
    return str(result.messages[-1].content)
 
async def handle_request(question: str, conversation_id: str, user_id: str) -> str:
    with rius.session(conversation_id), rius.user(user_id):
        return await research(question)

3. Set the session and user outside the run root, as handle_request does, so the root span carries them too. Use your own conversation id for the session, so one conversation is one session across its runs. Use an opaque user id rather than an email, because the value is stored on every span and shown in the console.

4. Run one question that makes the agent call a tool, then open the trace in the console. You should see the root span, your stage and retrieval spans, the framework’s invoke_agent and execute_tool spans, and the ChatCompletion generation with model, tokens and cost.

What you get

• One trace per research question, rooted at your research-run span, with the agent and tool structure from AutoGen nested beneath it.
• A named step for each stage you wrapped, so you can see whether time or failures sit in planning, retrieval or the agent loop.
• A retrieval span that shows which index was searched, how many documents were requested, and the id and score of each document returned.
• Tokens and a list-price cost estimate in US dollars on every model call.
• Traces grouped by session, so a multi-turn research conversation reads as one unit, and by user, so you can see what one person’s runs cost.

Some details of the AutoGen spans are worth knowing. Framework span statuses arrive as UNSET on success, though the trace status still reads Ok. Generation spans are all named ChatCompletion, so tell them apart by their position in the trace.

When something else fits

• The AutoGen integration is Python only. For other languages, trace model calls with manual generations as described in Tracing your code.
• If your retrieval runs as a tool the model calls, the framework’s execute_tool span already records it. Add a retrieval span only for search you run in your own code.
• If you want a single-purpose run root with no stages, the observe entrypoint alone from the AutoGen page is enough.

Related