What should I trace in a data analysis agent?

Rius traces a data analysis agent best as one root span per run, with child spans for model calls, tools and steps, plus a session and a user on every span. Each question a user asks becomes one trace, so you can read the run from the question to the final answer. Rius computes cost from the token counts on the model-call spans, so you never send it. For the wider picture, see Monitor AI agents in production.

How it works

A trace is one run, and it is a tree of spans. For a data analysis agent, the tree usually looks like this:

• Root span, one per run. Wrap the function that handles one question with observe and kind AGENT. The root’s duration is the run’s latency, and a failed root means a failed run.
• Generations for model calls. Every call that plans the analysis, writes a query or explains a result gets its own span of kind LLM. It records the model, token usage and finish reason, which is what Rius prices and what the model analytics read. With a global init(), installed integrations create these spans for you. For a call no integration records, open one yourself with start_as_current_generation. Around a call an integration already records, it counts the call twice.
• Tool spans. Each tool the model can call, such as running SQL, calling a dataframe helper or rendering a chart, gets a span of kind TOOL. Pass the tool’s own name, and tool_type such as datastore for a database call. When you know the id the model put on the tool call, open the tool span with rius.start_as_current_span(kind=SpanKind.TOOL, tool_name=..., tool_call_id=...). That ties it back to the generation that asked for it.
• Retriever spans. If the agent looks up table schemas, metric definitions or documentation first, trace that as kind RETRIEVER with a data_source_id. It records document ids and scores, not text.
• Step spans. Load, clean, join and summarize stages that are not model or tool calls fit the default CHAIN kind. Wrapping a function with observe is enough.

Spans nest through the normal OpenTelemetry parent context, so the console shows the whole run as a waterfall. The path from the root to a failing span is highlighted. Details on each helper are in Tracing your code.

Set it up

1. Initialize the SDK once at startup with rius.init(api_key="ri_...", service_name="data-analyst"). The service_name value is yours.
2. Wrap the run, the tools and the steps. Keep the SQL result out of the trace if it holds sensitive rows, by turning off output capture on that tool:

import rius
from rius import SpanKind
 
@rius.observe(kind=SpanKind.AGENT, agent_name="data-analyst")
def run_analysis(question: str) -> str:
    schema = lookup_schema(question)
    sql = draft_query(question, schema)
    rows = run_sql(sql)
    return summarize(question, rows)
 
@rius.observe(
    kind=SpanKind.TOOL,
    tool_name="run_sql",
    tool_type="datastore",
    capture_output=False,
)
def run_sql(sql: str) -> list[dict]:
    ...
 
@rius.observe
def summarize(question: str, rows: list[dict]) -> str:
    ...

The names lookup_schema, draft_query, run_sql and summarize are your own functions.

3. Trace the model calls. If you use a supported client library, auto-instrumentation creates the generation spans. If you call a model yourself, open a generation and pass token usage with set_usage(), never a cost.
4. Set the session and user around each request:

with rius.session(conversation_id):
    with rius.user(request.user_id):
        answer = run_analysis(question)

Use the conversation or job id you already have as the session, so every follow-up question lands in the same session. Each question is still its own trace. Use an opaque id for the user rather than an email, because the value is stored on every span. Set both scopes on every request, not as a process-wide default, if the agent serves several people.

The rius.session() and rius.user() scopes are described in Tracing your code. The TypeScript SDK has withSession and withUser for the same job.

What you get

• One trace per question, with the root span’s duration and status as the run’s latency and outcome.
• Cost and tokens per model call, as list-price estimates in US dollars, so you can see whether the query-writing step or the explanation step is the expensive one.
• A waterfall that shows which tool call, such as a slow or failing query, held up or broke a run.
• Tool spans joined to the generation that requested them, so you can see which model output led to which query.
• Sessions that group the follow-up questions of one analysis, and a Users view that shows what each person asked, failed on and cost. Without a session id, each trace counts as its own session. Without a user id, spans are anonymous. See Concepts for how the model fits together.
• Metadata that always flows. Setting capture_output=False drops a span’s return value and capture_input=False its arguments. Turning off content capture strips prompts, results and arguments everywhere. Timings, token counts and cost remain.

When something else fits

• If your analysis runs as a plain batch job with no model calls, one root span with a few CHAIN steps is enough, and the generation and tool spans do not apply.
• If you need scores, evals or datasets to judge whether an answer is correct, Rius does not offer them. It traces, prices and alerts on agent runs.

Related