How do I trace an agent built with LlamaIndex and deployed to AWS Lambda?

Rius traces a LlamaIndex agent on AWS Lambda when you call init once at module level, flush spans at the end of every invocation, and set a service name for the agent. The LlamaIndex integration is Python only, so this setup applies to Python Lambda functions. Turn the heartbeat off, because Lambda freezes the process between invocations.

How it works

The LlamaIndex integration turns each query into a nested trace of typed spans: retrieval steps, embedding calls and the model calls that synthesize the answer. Rius computes cost from the model and token counts on those spans. The LlamaIndex integration page covers what lands in the trace.

Lambda changes three things.

• Where init runs. Call rius.init() at module level, before you build indexes and query engines. It then runs once per cold start, and every invocation on that environment reuses it. The integration patches libraries process-wide, so it must be in place before your first query.
• How spans leave. Spans are queued in memory and exported from a background thread every few seconds. Lambda freezes the process between invocations, so a frozen thread exports nothing and the exit hook may never run. You force the export yourself at the end of every invocation.
• What the heartbeat does. The heartbeat is a liveness ping for long-lived processes. In a frozen runtime the pings stop, and the instance flaps between alive and stale or gone on the Agents page. Turn it off.

Set it up

1. Add the instrumentation and LlamaIndex to your function’s dependencies. The extra installs the instrumentation, not llama-index itself. Integrations lists the other frameworks and libraries.

   pip install "glassflow-rius[llama-index]" llama-index

2. Call init at module level, outside the handler, and name the agent. The service_name is the name the agent is grouped under in Rius. Set heartbeat=False for Lambda.

3. Wrap the handler body in try and finally, and flush in the finally block so the export runs even when the query raises.

   import logging
   import rius
    
   client = rius.init(
       api_key="ri_...",
       service_name="support-agent",
       heartbeat=False,
   )
   # build indexes and query engines exactly as before
    
   def handler(event, context):
       try:
           return {"answer": str(query_engine.query(event["question"]))}
       finally:
           if not client.flush(timeout_millis=5_000):
               logging.warning("rius: span export did not complete")

   Here handler and query_engine are your own names. Pick a timeout_millis that fits inside your function’s timeout, because flush() blocks until the export completes. Call it only at the end of the invocation, never inside the query path.

4. Keep the key and name out of the code if you prefer. The SDK also reads RIUS_API_KEY and RIUS_SERVICE_NAME, and RIUS_HEARTBEAT=false turns the heartbeat off, so you can set all three as Lambda environment variables. RIUS_AGENT_NAME overrides the grouping name and defaults to the service name.

5. Invoke the function once and open the trace in the console. The query should appear as nested spans, with the model call carrying model, tokens and cost.

If the openai extra is installed too, the model call also shows as a second span under the LlamaIndex span. Tokens and cost are still counted once. Details on the flush and the retry rules are on the reliability page.

What you get

• One trace per query, with retrieval, embedding and model-call spans nested under it, so you can see what the retriever handed to the model.
• Cost per model call, computed by Rius as a list-price estimate in US dollars.
• One agent entry in Rius under the name you chose, marked no heartbeat with no Instances section, instead of instances flapping between stale and gone.
• A signal when an export fails. flush() returns False on a timeout or when the most recent export was rejected, such as a 401 on a bad key. Your log line catches that at the end of the invocation.
• Retrieved document content is captured as content attributes. The privacy controls cover it, and document ids and scores stay visible even with content stripped.

When something else fits

• If your Lambda function is written in TypeScript, there is no LlamaIndex integration for it. The OpenAI integration still traces the provider calls, and manual generations cover the pipeline around them.
• A flush does not guarantee delivery. A rejected or timed-out batch is dropped and not retried later, so a 402 or a sustained 429 loses spans. If an invocation can crash before it reaches the flush, enable partial spans, which leave a content-free snapshot of each span at start.

Related