Observability & Tracing

import time
from contextlib import contextmanager

from langchain_core.documents import Document
from langchain_core.messages import AIMessage
from opentelemetry import trace

tracer = trace.get_tracer("docsassistant")


@contextmanager
def retrieve_span(index: str, top_k: int, filter_hash: str):
    with tracer.start_as_current_span(f"retrieve {index}") as span:
        span.set_attribute("docsassistant.retrieval.top_k", top_k)
        span.set_attribute("docsassistant.retrieval.filter_hash", filter_hash)
        result: dict = {}
        yield result                                   # the caller stores the documents in result["docs"]
        docs: list[Document] = result.get("docs", [])
        span.set_attribute("docsassistant.retrieval.returned", len(docs))
        span.set_attribute("docsassistant.retrieval.chunk_ids", [d.metadata.get("chunk_id", "") for d in docs])
        if not docs:
            span.add_event("empty_retrieval")


@contextmanager
def generate_span(model: str):
    with tracer.start_as_current_span(f"chat {model}") as span:
        span.set_attribute("gen_ai.operation.name", "chat")
        span.set_attribute("gen_ai.request.model", model)
        state = {"start": time.perf_counter(), "first": None}
        yield state                                    # the caller sets state["first"] on the first token
        if state["first"] is not None:
            span.set_attribute("docsassistant.ttft_ms", round((state["first"] - state["start"]) * 1000, 1))
        final: AIMessage | None = state.get("message")
        if final is not None and final.usage_metadata:
            span.set_attribute("gen_ai.usage.input_tokens", final.usage_metadata["input_tokens"])
            span.set_attribute("gen_ai.usage.output_tokens", final.usage_metadata["output_tokens"])

package com.example.docsassistant.rag.observability;

import java.util.List;

import io.micrometer.core.instrument.DistributionSummary;
import io.micrometer.core.instrument.MeterRegistry;
import org.springframework.ai.chat.client.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.core.Ordered;

public class RagOutcomeMetricsAdvisor implements CallAdvisor {

    private final MeterRegistry meters;
    private final DistributionSummary retrieved;

    public RagOutcomeMetricsAdvisor(MeterRegistry meters) {
        this.meters = meters;
        this.retrieved = DistributionSummary.builder("docsassistant.rag.documents").register(meters);
    }

    @Override
    public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
        ChatClientResponse response = chain.nextCall(request);
        if (response.chatResponse() != null) {
            List<Document> docs = response.chatResponse().getMetadata().get(RetrievalAugmentationAdvisor.DOCUMENT_CONTEXT);
            int count = docs == null ? 0 : docs.size();
            retrieved.record(count);
            if (count == 0) {
                meters.counter("docsassistant.rag.empty_retrieval").increment();
            }
        }
        return response;
    }

    @Override
    public String getName() {
        return "ragOutcomeMetrics";
    }

    @Override
    public int getOrder() {
        return Ordered.HIGHEST_PRECEDENCE + 500;       // outside the RAG advisor, so it sees its result
    }
}

package com.example.docsassistant.l4j.observability;

import dev.langchain4j.model.chat.listener.ChatModelErrorContext;
import dev.langchain4j.model.chat.listener.ChatModelListener;
import dev.langchain4j.model.chat.listener.ChatModelRequestContext;
import dev.langchain4j.model.chat.listener.ChatModelResponseContext;
import dev.langchain4j.model.output.TokenUsage;
import dev.langchain4j.store.embedding.listener.EmbeddingStoreListener;
import dev.langchain4j.store.embedding.listener.EmbeddingStoreResponseContext;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.SpanKind;
import io.opentelemetry.api.trace.StatusCode;
import io.opentelemetry.api.trace.Tracer;

public final class OtelListeners {

    private OtelListeners() {
    }

    public static ChatModelListener chat(Tracer tracer) {
        return new ChatModelListener() {
            @Override
            public void onRequest(ChatModelRequestContext ctx) {
                String model = ctx.chatRequest().modelName() == null ? "unknown" : ctx.chatRequest().modelName();
                Span span = tracer.spanBuilder("chat " + model).setSpanKind(SpanKind.CLIENT).startSpan();
                span.setAttribute("gen_ai.operation.name", "chat");
                span.setAttribute("gen_ai.request.model", model);
                ctx.attributes().put(Span.class, span);          // attributes are shared by the three callbacks
            }

            @Override
            public void onResponse(ChatModelResponseContext ctx) {
                Span span = (Span) ctx.attributes().get(Span.class);
                TokenUsage usage = ctx.chatResponse().tokenUsage();
                if (usage != null && usage.inputTokenCount() != null) {
                    span.setAttribute("gen_ai.usage.input_tokens", usage.inputTokenCount());
                    span.setAttribute("gen_ai.usage.output_tokens", usage.outputTokenCount());
                }
                if (ctx.chatResponse().modelName() != null) {
                    span.setAttribute("gen_ai.response.model", ctx.chatResponse().modelName());
                }
                span.end();
            }

            @Override
            public void onError(ChatModelErrorContext ctx) {
                Span span = (Span) ctx.attributes().get(Span.class);
                if (span != null) {
                    span.setStatus(StatusCode.ERROR, ctx.error().getClass().getSimpleName());
                    span.end();
                }
            }
        };
    }

    public static EmbeddingStoreListener retrieval() {
        return new EmbeddingStoreListener() {
            @Override
            public void onResponse(EmbeddingStoreResponseContext<?> ctx) {
                if (ctx instanceof EmbeddingStoreResponseContext.Search<?> search) {
                    int returned = search.searchResult().matches().size();
                    Span.current().setAttribute("docsassistant.retrieval.returned", returned);
                    if (returned == 0) {
                        Span.current().addEvent("empty_retrieval");
                    }
                }
            }
        };
    }
}