Example: LangChain RAG

Prerequisites: 02-tracing-observations.md and 05-tracing-frameworks.md.

A retrieval-augmented generation (RAG) pipeline: given a user query, retrieve relevant documents from a vector store, then generate an answer using an LLM with the retrieved context.

"""
Tracing annotations:
- setup_tracing() makes Langfuse the global trace backend.
- @step with as_type="retriever" marks the retrieval step so Langfuse shows it
  with the retriever icon.
- @step with as_type="generation" on the LLM call enables token/cost tracking.
- get_langchain_callback() nests LangChain's internal spans (prompt template,
  model call, output parser) under the enclosing @step observation automatically.
- event() marks edge cases (no documents found, empty context fallback).
- update_current_step() attaches computed metadata after execution.
"""

from semantic_tracing import (
    event,
    flush,
    get_langchain_callback,
    setup_tracing,
    step,
    update_current_step,
)
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI


# --- tracing init ---
setup_tracing(service_name="rag_pipeline")


# --- retrieval step ---

retriever = build_vector_retriever()  # application-specific setup

@step("retrieve_context", as_type="retriever")
async def retrieve(query: str) -> list[str]:
    docs = await retriever.ainvoke(query)
    if not docs:
        event("no_documents_found", metadata={"query": query})
        return []
    update_current_step(metadata={"doc_count": len(docs), "top_score": docs[0].metadata.get("score")})
    return [doc.page_content for doc in docs]


# --- generation step ---

@step("generate_answer", as_type="generation")
async def generate(query: str, context: list[str]) -> str:
    if not context:
        event("empty_context_fallback")
        return "I don't have enough information to answer."

    handler = get_langchain_callback()
    prompt = ChatPromptTemplate.from_messages([
        ("system", "Answer based on context:\n{context}"),
        ("human", "{query}"),
    ])
    chain = prompt | ChatOpenAI(model="gpt-4o") | StrOutputParser()
    answer = await chain.ainvoke(
        {"query": query, "context": "\n".join(context)},
        config={"callbacks": [handler]},
    )
    update_current_step(metadata={"context_length": len(context)})
    return answer


# --- pipeline entry ---

@step("rag_pipeline")
async def run_pipeline(query: str) -> str:
    context = await retrieve(query)
    return await generate(query, context)


# --- execution ---

async def main():
    answer = await run_pipeline("What is our refund policy?")
    print(answer)
    flush()  # ensure all buffered observations are sent to Langfuse

Resulting Langfuse trace:

trace (rag_pipeline)
└── rag_pipeline (span)
      input: {"query": "What is our refund policy?"}
      output: "Our refund policy allows returns within 30 days..."
      ├── retrieve_context (retriever)
      │     input: {"query": "What is our refund policy?"}
      │     output: ["Refund policy document...", "Returns FAQ..."]
      │     metadata: {"doc_count": 2, "top_score": 0.92}
      └── generate_answer (generation)
            input: {"query": "...", "context": ["...", "..."]}
            output: "Our refund policy allows returns within 30 days..."
            metadata: {"context_length": 2}
            └── ChatOpenAI (LangChain generation)
                  input: {messages: [...]}
                  output: "Our refund policy..."
                  usage: {prompt_tokens: 312, completion_tokens: 45}