Tracing embedding, retrieval, and generation as one span

Most RAG stacks log embedding, retrieval, and generation separately, leaving you blind when a bad answer traces back to a silent retrieval miss. Tracing embedding retrieval generation as a single distributed span closes that gap: you see latency, token cost, and vector hit rate in one timeline.

Why a single span matters

A RAG request is three dependent phases: embed the query, fetch candidate documents, then synthesize a response. If you emit independent logs for each, reconstructing the causal chain requires manual correlation by timestamp or request ID. That works until you have concurrency, retries, or multiple simultaneous users.

A distributed trace models this as a tree of spans under one trace ID. The root span represents the whole query; child spans represent embedding, retrieval, and generation. When you adopt tracing embedding retrieval generation as one trace, you can answer questions like “did the slow answer come from a cold vector index or from model inference?” in one query to your observability backend.

The alternative is three dashboards and a grep across log files. In incident response, that delay is the difference between a five-minute fix and a rollback.

Step 1: Start a root span at the boundary

Instrument the outermost handler. In a service, that’s your HTTP endpoint or queue consumer. Use OpenTelemetry’s Python SDK to create a provider and a batch exporter so you don’t block the request path:

from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, OTLPExporter

trace.set_tracer_provider(TracerProvider())
trace.get_tracer_provider().add_span_processor(
    BatchSpanProcessor(OTLPExporter(endpoint="http://localhost:4317"))
)
tracer = trace.get_tracer("rag.pipeline")

def handle_query(user_id: str, query: str):
    with tracer.start_as_current_span("rag_query") as span:
        span.set_attribute("user.id", user_id)
        span.set_attribute("query.hash", hash(query) & 0xffffffff)
        # proceed to embedding, retrieval, generation

The root span must be active (start_as_current_span) so child spans automatically nest. Do not create it with start_span and forget to activate; you’ll get orphan spans that never attach to the parent. If you run multiple worker processes, ensure each initializes the provider once at startup, not per request.

Step 2: Embed the query as a child span

Nested context is easiest when the calls are in-process. Create a child span for the embedding call and record model identity and vector size.

def embed_query(text: str):
    with tracer.start_as_current_span("embedding") as span:
        span.set_attribute("embedding.model", "text-embedding-3-small")
        span.set_attribute("embedding.dim", 1536)
        vec = embeddings_client.embed(text)
        span.set_attribute("embedding.latency_ms", round(vec.latency * 1000, 1))
        return vec

If embedding runs in a separate service, propagate the OTel context via headers using opentelemetry.propagate.inject. The receiving service extracts it and continues the trace. Skipping this is the most common reason tracing embedding retrieval generation breaks across microservices. In Python, the contextvars backing keeps the active span across async/await, but only if you don’t launch detached threads.

Step 3: Tag the retrieval span with hit metadata

Retrieval is where silent failures hide. A span that only records “called index” is useless. Capture the number of hits, top score, and index name.

def retrieve(vec, top_k=5):
    with tracer.start_as_current_span("retrieval") as span:
        span.set_attribute("retrieval.index", "docs-2024")
        span.set_attribute("retrieval.top_k", top_k)
        results = vector_db.search(vec, top_k)
        span.set_attribute("retrieval.hits", len(results))
        if results:
            span.set_attribute("retrieval.top_score", results[0].score)
        else:
            span.set_status(trace.Status(trace.StatusCode.ERROR, "no hits"))
        return results

Mark empty results as an error status. A generation span that proceeds with zero context should be visible as a degraded path, not a success. Also record whether the search used a cached index snapshot; stale snapshots produce plausible but wrong answers.

Step 4: Link generation to the same trace

Generation calls an LLM. If you use an OpenAI-compatible client, inject the current trace context so the inference gateway can correlate. When routing through n4n.ai, its OpenAI-compatible endpoint honors client routing directives and forwards provider cache-control hints, so the same trace parent travels to the upstream provider and per-token usage returns in the response.

from opentelemetry.propagate import inject
from openai import OpenAI

def generate(question: str, contexts: list[str]):
    with tracer.start_as_current_span("generation") as span:
        headers = {}
        inject(headers)  # adds traceparent
        client = OpenAI(base_url="https://api.n4n.ai/v1", api_key="KEY")
        resp = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "system", "content": " ".join(contexts)},
                      {"role": "user", "content": question}],
            extra_headers=headers
        )
        span.set_attribute("generation.tokens_in", resp.usage.prompt_tokens)
        span.set_attribute("generation.tokens_out", resp.usage.completion_tokens)
        return resp.choices[0].message.content

If you don’t use a gateway, still inject traceparent into your HTTP call to the model provider; many providers accept it for their own internal tracing. The key point for tracing embedding retrieval generation is that the generation span stays a child of the root, not a new trace.

Handling fallback and caching

Automatic fallback when a provider is degraded is a gateway concern. Your span should record which model actually served the request, not just the requested one. Read the response header or usage field that indicates the resolved route. If the gateway forwards cache-control hints, log whether the prompt hit a provider-side cache to explain token billing anomalies.

Step 5: Export and query

Use a batch exporter to avoid blocking request paths. Send to an OTLP-compatible collector (Jaeger, Tempo, Honeycomb). A minimal query in Jaeger looks like:

trace.duration > 2s and span.name = "retrieval"

This finds slow retrievals across all RAG traces. Set a sampling policy. Full sampling in production with high QPS will blow up storage. Use head sampling at 10% plus always-sample on error statuses. That way, every failed RAG run is captured, but happy-path noise is bounded.

What a unified trace looks like

A simplified JSON representation of one trace helps confirm your instrumentation:

{
  "traceId": "a1b2c3d4",
  "spans": [
    {"name": "rag_query", "parent": null, "dur_ms": 1200},
    {"name": "embedding", "parent": "rag_query", "dur_ms": 200},
    {"name": "retrieval", "parent": "rag_query", "dur_ms": 300},
    {"name": "generation", "parent": "rag_query", "dur_ms": 700}
  ]
}

When you see retrieval.hits: 0 under a rag_query that still returned text, you immediately know the model hallucinated around missing context. That is the payoff of tracing embedding retrieval generation as one object.

Common pitfalls and tradeoffs

Pitfall: async context loss. If you await embedding in a separate task without contextvars propagation, the child span detaches. Use opentelemetry-instrumentation-asyncio or pass the context explicitly via context.attach.

Pitfall: over-instrumentation. Recording full document text in span attributes seems helpful but inflates cardinality and may leak PII. Store only IDs, scores, and counts.

Pitfall: ignoring embedding errors. If embedding throws, but you catch and return a zero vector, the retrieval span will show “hits: 0” with no root cause. Let the exception propagate to mark the root span as error.

Pitfall: clock skew. Across services, unsynced clocks make span ordering nonsense. Run NTP and use the collector’s ingestion time as fallback.

Tradeoff: latency vs. visibility. Synchronous export adds milliseconds. Batch exporter mitigates this but can lose spans on crash. For most RAG apps, losing a few traces on restart is acceptable.

Tradeoff: vendor lock-in. Using gateway-specific headers like routing directives ties your tracing to that gateway. Keep a thin wrapper so you can switch providers without rewriting span logic.

Minimal end-to-end snippet

def rag_pipeline(user_id, query):
    with tracer.start_as_current_span("rag_query") as root:
        root.set_attribute("user.id", user_id)
        vec = embed_query(query)
        docs = retrieve(vec)
        answer = generate(query, [d.text for d in docs])
        return answer

That function, with the instrumented children shown earlier, gives you one trace per query. Tracing embedding retrieval generation this way turns an opaque pipeline into a debuggable system.

When you land in on-call at 3am because answers went stale, the span waterfall will show whether the embedding service timed out, the index returned empty, or the generation token stream stalled. Build the trace before you need it.