Observability for Agents

from langchain.agents import create_agent
from langchain.agents.middleware import wrap_tool_call
from langchain_ollama import ChatOllama
from opentelemetry import trace

tracer = trace.get_tracer("docsassistant")
AGENT_NAME = "docs_assistant"


@wrap_tool_call
def trace_tool_calls(request, handler):
    call = request.tool_call
    with tracer.start_as_current_span(
        f"execute_tool {call['name']}",
        attributes={
            "gen_ai.operation.name": "execute_tool",
            "gen_ai.tool.name": call["name"],
            "gen_ai.tool.call.id": call["id"],
        },
    ) as span:
        result = handler(request)
        if getattr(result, "status", None) == "error":
            span.set_status(trace.StatusCode.ERROR)
        return result


agent = create_agent(ChatOllama(model="llama3.1:8b", temperature=0),
                     tools=[docs_search], middleware=[trace_tool_calls], name=AGENT_NAME)


def ask(question: str, conversation_id: str) -> str:
    with tracer.start_as_current_span(
        f"invoke_agent {AGENT_NAME}",
        attributes={"gen_ai.operation.name": "invoke_agent",
                    "gen_ai.agent.name": AGENT_NAME,
                    "gen_ai.conversation.id": conversation_id},
    ):
        result = agent.invoke({"messages": [("user", question)]})
        return result["messages"][-1].content

# Langfuse: a callback handler per run
from langfuse.langchain import CallbackHandler

agent.invoke({"messages": [("user", "What does ef_search control?")]},
             config={"callbacks": [CallbackHandler()]})

# Phoenix: OpenTelemetry-based auto-instrumentation, once at startup
from openinference.instrumentation.langchain import LangChainInstrumentor
from phoenix.otel import register

tracer_provider = register(project_name="docsassistant")
LangChainInstrumentor().instrument(tracer_provider=tracer_provider)

management.tracing.sampling.probability=1.0
spring.ai.chat.client.observations.log-prompt=false
spring.ai.chat.client.observations.log-completion=false
spring.ai.tools.observations.include-content=true

import io.micrometer.observation.Observation;
import io.micrometer.observation.ObservationHandler;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.observation.ChatModelObservationContext;
import org.springframework.stereotype.Component;

@Component
class TokenUsageHandler implements ObservationHandler<ChatModelObservationContext> {

    @Override
    public boolean supportsContext(Observation.Context context) {
        return context instanceof ChatModelObservationContext;      // only gen_ai.client.operation
    }

    @Override
    public void onStop(ChatModelObservationContext context) {
        ChatResponse response = context.getResponse();
        if (response != null) {
            Usage usage = response.getMetadata().getUsage();
            System.out.printf("model=%s in=%d out=%d%n", response.getMetadata().getModel(),
                    usage.getPromptTokens(), usage.getCompletionTokens());
        }
    }
}

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.chat.response.ChatResponse;
import dev.langchain4j.model.ollama.OllamaChatModel;
import java.util.List;

class TracingListener implements ChatModelListener {

    @Override
    public void onRequest(ChatModelRequestContext ctx) {
        ctx.attributes().put("start", System.nanoTime());
    }

    @Override
    public void onResponse(ChatModelResponseContext ctx) {
        long ms = (System.nanoTime() - (long) ctx.attributes().get("start")) / 1_000_000;
        ChatResponse response = ctx.chatResponse();
        System.out.printf("model=%s ms=%d in=%d out=%d%n", response.modelName(), ms,
                response.tokenUsage().inputTokenCount(), response.tokenUsage().outputTokenCount());
    }

    @Override
    public void onError(ChatModelErrorContext ctx) {
        System.out.printf("model call failed: %s%n", ctx.error().getMessage());
    }
}

OllamaChatModel model = OllamaChatModel.builder()
        .baseUrl("http://localhost:11434")
        .modelName("llama3.1:8b")
        .listeners(List.of(new TracingListener()))
        .build();