How do I trace a document processing pipeline built with AutoGen?

Rius traces an AutoGen document pipeline by combining AutoGen’s own spans and the openai integration with a few observe and span wrappers of your own. The framework gives you agent, tool and model-call spans with tokens and cost. Your wrappers add a root span for each document, a span for each stage and a retrieval span, and you set the session and user around each run. Rius supports AutoGen in Python.

How it works

Two layers produce the trace.

• The framework layer. AutoGen emits create_agent, invoke_agent and execute_tool spans through the global tracer provider that rius.init() installs. Model calls go through the OpenAI client, so the openai integration records each one as a ChatCompletion generation with model, tokens and cost. You wrap nothing for this layer. The AutoGen integration page covers it in full.
• Your layer. AutoGen knows nothing about the stages around your agent, such as reading a file, searching a knowledge base or writing results. You add those with the wrappers described in Tracing your code:
  • observe with kind=SpanKind.AGENT on the entrypoint makes the root span, so agent creation and invocation group into one trace per document.
  • observe on a stage function makes a span named after the function (or the name you pass).
  • rius.start_as_current_span with kind=SpanKind.RETRIEVER makes a retrieval span. data_source_id names the index searched and top_k records how many documents you asked for. set_retrieved_documents() stores document ids and scores, not document text.

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 a one-span trace. The execute_tool spans already cover your tools. Run retrieval and other stages in your own pipeline code, outside the tools.

Your wrapper spans carry no tokens or cost of their own, so totals come from the model calls. Costs are list-price estimates in US dollars.

Set it up

1. Install the SDK with the openai extra, plus the AutoGen packages. The integrations overview lists every extra.

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

2. Initialize Rius once at startup, then wrap your pipeline. This example extracts text, retrieves context and runs an AutoGen agent. read_document, search and fetch_vendor stand in 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="doc-pipeline")
 
def lookup_vendor(name: str) -> str:
    """Fetch a vendor record by name."""
    return fetch_vendor(name)  # no observe here: AutoGen traces tools
 
@rius.observe(name="extract-text")
def extract_text(path: str) -> str:
    return read_document(path)
 
def retrieve_context(text: str) -> list[str]:
    with rius.start_as_current_span(
        kind=SpanKind.RETRIEVER, data_source_id="policy-kb", top_k=5
    ) as obs:
        obs.set_input(text[:500])
        docs = search(text)
        obs.set_retrieved_documents([{"id": d.id, "score": d.score} for d in docs])
        return [d.text for d in docs]
 
@rius.observe(name="process-document", kind=SpanKind.AGENT)
async def process_document(path: str) -> str:
    text = extract_text(path)
    context = retrieve_context(text)
    agent = AssistantAgent(
        name="doc_agent",
        model_client=OpenAIChatCompletionClient(model="gpt-5-mini"),
        tools=[lookup_vendor],
        system_message="Summarize the document and check vendors with the tool.",
    )
    result = await agent.run(task=text + "\n\nContext:\n" + "\n".join(context))
    return str(result.messages[-1].content)
 
async def handle_upload(path: str, document_id: str, user_id: str) -> str:
    with rius.session(document_id), rius.user(user_id):
        return await process_document(path)

3. Choose the session id. rius.session() takes an id you mint, so pass one that matches how you group work, such as a document id, a batch job id or a conversation id. With no session scope, each trace is its own session.
4. Choose the user id. rius.user() stamps user.id on every span in scope, including spans from auto-instrumentation. Prefer an opaque id to an email. A span outside any user scope is anonymous.
5. Run one document that triggers a tool call and open the trace in the console. You should see your root span, the extract-text and retrieval spans, AutoGen’s invoke_agent and execute_tool spans, and a ChatCompletion generation with model, tokens and cost.

What you get

• One trace per document, rooted at your process-document span, with each stage as a child span you can open and time.
• A retrieval span that shows the data source, how many documents you asked for, and the ids and scores of what came back.
• AutoGen’s agent and tool structure, and every model call with tokens and a list-price cost estimate, with no extra wrappers.
• Traces grouped by your session id, and per-user views from user.id, so you can follow one document or one person across runs.

A few AutoGen details are worth knowing. Framework spans arrive with an UNSET status rather than OK on success, although the trace status still reads Ok. Generation spans are all named ChatCompletion, so tell them apart by position in the tree.

When something else fits

• If your pipeline is TypeScript, there is no AutoGen integration. Trace those model calls with manual generations from Tracing your code.
• If your stages call a model outside AutoGen, the integrations overview shows which other libraries Rius records automatically.
• If a stage is a plain function with no model call, observe alone is enough and you need no retrieval span.

Related