How do I trace an agent built with LlamaIndex and deployed to a Celery worker?

Rius traces a LlamaIndex agent in a Celery worker once you call init in each worker process with the llama-index extra installed. Query engines, retrievers, embeddings and model calls arrive as nested spans with tokens and cost. Start Rius from Celery’s worker_process_init signal, shut it down from worker_process_shutdown, and the name you pass to init is the name the agent appears under.

How it works

• Init once per worker process. Call rius.init in each Celery worker process after it forks, not at module import in the parent. With a default global init, every installed integration turns on automatically, including LlamaIndex. See LlamaIndex for what lands in the trace.
• Spans export in the background. Spans queue in memory and export in batches from a background thread, so a task never waits on the network.
• A query becomes a nested trace. Retrieval steps are RETRIEVER spans, embedding calls are EMBEDDING spans, and synthesis model calls are LLM spans with model, token usage and messages. Rius computes cost on the server from the model and token counts.
• One instance per worker process. Each process that calls init gets its own heartbeat instance ID, so each worker process shows as its own instance on the Agents page.
• The agent name comes from the service name. service_name in rius.init, or RIUS_SERVICE_NAME, sets the service name. RIUS_AGENT_NAME sets the name the agent is grouped under and defaults to the service name.

Set it up

1. Install the SDK with the LlamaIndex extra. The extra installs the instrumentation, not LlamaIndex itself, so name llama-index for a fresh environment. Without it, init() logs a DependencyConflict and the integration stays off.

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

2. Start Rius in each worker process and stop it when that process exits. Celery’s prefork pool sends worker_process_init in every child after the fork, and worker_process_shutdown just before the child exits. Children end through os._exit, which skips Python’s exit hook, so shut the client down yourself. Pick a service name that says what the agent is, since this is the name you look for in the console.

   import rius
   from celery import Celery
   from celery.signals import worker_process_init, worker_process_shutdown
    
   app = Celery("agents", broker="redis://localhost:6379/0")
   client = None
    
   @worker_process_init.connect
   def start_rius(**kwargs):
       global client
       client = rius.init(api_key="ri_...", service_name="my-agent")
    
   @worker_process_shutdown.connect
   def stop_rius(**kwargs):
       if client is not None:
           client.flush(timeout_millis=10_000)  # force an export
           client.shutdown()  # drain, stop the thread, send the stopped heartbeat

   Build indexes and query engines exactly as before, inside your tasks or lazily after init. Keep rius.init out of module scope, because the parent imports your tasks module too. The threads, gevent and eventlet pools do not send worker_process_init, so with those pools call init once when the worker starts.

3. If you want the agent grouped under a different name than the service, set it in the worker’s environment:

   export RIUS_AGENT_NAME=support-agent

4. Know what a hard kill costs. Spans export in batches every few seconds, so a worker process killed with SIGKILL, by a hard time limit or the out-of-memory killer, loses the spans still queued and sends no stopped heartbeat, and its instance reports gone. Do not call flush inside a task, because it blocks until the export completes. Details are on Reliability.

5. Run one task that queries an index, then open the trace in the console. You see the pipeline as nested spans, with the LLM leaf carrying model, tokens and cost.

What you get

• A nested trace for every query your worker runs: retrieval, embedding and synthesis steps, each with timing.
• Tokens and a list-price cost estimate in US dollars on each model call. If the openai extra is also installed, the provider call is counted once in tokens and cost, and the LlamaIndex span still shows the call’s total.
• The agent listed on the Agents page under the name you chose, with an instance per worker process while heartbeats are on. Heartbeats need credentials, so an agent with no API key shows as no heartbeat and lists no instances.
• Spans that reach Rius without your tasks waiting on the network. After a clean stop that calls shutdown, the instance shows a clean exit rather than a crash.
• Retrieved document content and messages arrive as content attributes. Document ids and scores stay visible even if you strip content with the privacy controls.

When something else fits

• If you run the solo pool, for example while testing locally, worker_process_init fires once in the main process, so the same handlers work.
• If your pipeline is in TypeScript with LlamaIndex.TS, there is no LlamaIndex integration for it. See Integrations for what each language covers.
• If you run in a Lambda-style runtime rather than a long-lived worker, set the heartbeat off and flush at the end of every invocation.

Related