Tracing a RAG Pipeline

Find which stage failed

A retrieval-augmented generation (RAG) answer passes through several stages: rewrite or embed the query, search a vector store, optionally rerank, build the prompt, then generate. A wrong answer can come from any of them, and the fix is completely different. If the right document never appeared in the retrieved chunks, the problem is retrieval: chunking, embeddings, filters or top_k. If the right chunk was retrieved but the answer ignored or contradicted it, the problem is generation: prompt wording, context ordering or the model. Without tracing you guess. With a span per stage you open the trace, read the retrieved chunks, and decide in seconds. Record for the retrieval span: the query actually sent (after rewriting), the filters, the returned chunk IDs, scores and text. Record for generation: the final assembled prompt. These traces also become your evaluation dataset later.

Spans in a RAG request

Each RAG stage gets its own span so a trace shows exactly where an answer went wrong.

One span per stage (Langfuse)

Stub functions keep the example short; swap in your real embedder, store and model client. Requires the langfuse package.

from langfuse import observe

@observe()
def rewrite(q: str) -> str:
    return q.strip().lower()

@observe()
def search(q: str, top_k: int = 4) -> list[dict]:
    # return IDs and scores too, so the trace shows *why* chunks were chosen
    return [{"id": "faq-12", "score": 0.83, "text": "Refunds within 7 days."}]

@observe()
def generate(q: str, chunks: list[dict]) -> str:
    prompt = "Context:\n" + "\n".join(c["text"] for c in chunks) + f"\n\nQ: {q}"
    return "No, refunds are allowed only within 7 days."  # model call here

@observe()
def rag_answer(question: str) -> str:
    q = rewrite(question)
    return generate(q, search(q))

Log chunk IDs, not only text

If you store only chunk text, you cannot tell which source document or version it came from. Log the chunk ID, source path and score. When a user reports a wrong answer, you can fix the exact document and re-run the trace's query.

Quick check: A trace shows the correct policy chunk in the retrieval span, yet the answer contradicts it. Where should you focus first?

• Generation: the prompt, context ordering or model choice
• Retrieval: the embedding model
• Chunking: the chunk size is too large
• The vector database's index type