What should I trace in a research agent?

Rius gives a research agent a clear span tree when you trace one root span per run, a generation for each model call, and a span for each search, tool call and step. Set a session id and a user id around the run so the traces group by job and by person. You then see where a run spent its time and cost, and which search or tool failed.

How it works

A trace is one research run, and it has one root span. The root’s duration is the run’s latency, and a failed root means a failed run. A tree that fits a research agent looks like this:

• Root span, kind AGENT. One per question or job. Wrap your top-level function.
  • Plan step, kind CHAIN. Where the agent breaks the question into sub-questions.
  • Search spans, kind RETRIEVER. One per search or index lookup, with the data source and how many results you asked for.
  • Tool spans, kind TOOL. One per page fetch, file read or API call.
  • Generations, kind LLM. One per model call: planning, reading, summarizing, writing the answer.
  • Rerank span, kind RERANKER. Optional, for a step that reorders results by relevance.

Each kind answers a different question. Generations carry the model, token counts and finish reasons, and Rius computes cost from them. Retriever spans record gen_ai.data_source.id and gen_ai.retrieval.top_k, and set_retrieved_documents() adds document ids and scores. That is metadata, so it survives content stripping. Tool spans carry the tool name, so a failing fetch groups with other failures of the same tool.

Research runs are long and loop through search, read and summarize. Giving every model call its own generation and every lookup its own span lets the waterfall show which loop iteration was slow or expensive. Send token usage and never a cost value. Use the exact spelling anthropic as the provider for Claude calls, so cached tokens are counted correctly.

The session and the user are separate from the tree. A session groups the traces of one research job or conversation across turns. A user is the person the run served, and it persists across sessions. The SDK stamps both on every span inside the scope. See Concepts for the full data model.

Set it up

1. Install and initialize the SDK once. Tracing your code and the installation page cover this in full.

   pip install glassflow-rius

2. Wrap the run as the root span, and give tools, searches and steps their own kind. The function names below are your own.

   import rius
   from rius import SpanKind
    
   rius.init(api_key="ri_...", service_name="research-agent")
    
   @rius.observe(kind=SpanKind.TOOL)
   def fetch_page(url: str) -> str:
       ...
    
   @rius.observe
   def plan(question: str) -> list[str]:
       ...
    
   @rius.observe(kind=SpanKind.AGENT)
   def research(question: str) -> str:
       for sub_question in plan(question):
           with rius.start_as_current_span(
               kind=SpanKind.RETRIEVER, data_source_id="web-search", top_k=5
           ) as obs:
               obs.set_input(sub_question)
               docs = search(sub_question)
               obs.set_retrieved_documents([{"id": d.id, "score": d.score} for d in docs])
           for doc in docs:
               fetch_page(doc.url)
       ...

3. Add a generation around each model call you make by hand. Pass tokens, not cost. If the openai extra is installed, leave openai out of rius.init(instruments=[...]), or the call is counted twice.

   with rius.start_as_current_generation(model="gpt-4o", provider="openai") as gen:
       gen.set_input(messages)
       response = client.chat.completions.create(model="gpt-4o", messages=messages)
       gen.set_output(response.choices[0].message.content)
       gen.set_usage(
           input_tokens=response.usage.prompt_tokens,
           output_tokens=response.usage.completion_tokens,
       )

   If you use a supported client library or framework, auto-instrumentation creates these generations for you.

4. Scope the session and the user around each run. Use your own job or conversation id, and an opaque user id rather than an email.

   with rius.session(job_id):
       with rius.user(request.user_id):
           research(question)

   observe takes no user argument, because a user belongs to each request. Do not set a process-wide session in a server that handles many users, or all their traces merge into one session.

5. For long runs, turn on partial spans with RIUS_PARTIAL_SPANS. Rius then records a content-free snapshot when each span starts, so a run still going is visible and a crashed run leaves a record. In TypeScript, name your wrapped functions and await shutdown() before a short script exits.

What you get

• A waterfall for each run that shows the plan, every search, every fetch and every model call in order, with the path to the failing span highlighted.
• Cost per model call, as a list-price estimate in US dollars, so you see which step of the research loop spends the most.
• Retrieval spans with data source, result count and document scores, so you can compare what a search returned with what the agent then read.
• Sessions that tie together the traces of one research job, and a Users view of each person’s traces, failures and cost.
• Error grouping by tool and by exception type, because failures are recorded with error.type on the span that raised them.

For more on reading these views across all your agents, see Monitor AI agents in production.

When something else fits

If your research agent is a single model call with no search or tools, one generation under a root span is enough. If you only need model calls and use a supported client library, auto-instrumentation gives you generations without manual spans, and you can add retriever and tool spans later where you want more detail.

Related