spring-ai/observability/src/main/java/com/ankurm/observability/UsageCostObservationHandler.java

package com.ankurm.observability;

import io.micrometer.core.instrument.Counter;
import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.observation.Observation;
import io.micrometer.observation.ObservationHandler;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.observation.ChatClientObservationContext;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.ai.chat.model.ChatResponse;

/**
 * When a ChatClient call finishes, reads the token usage the provider reported and records two
 * counters tagged by endpoint and model: {@code app.ai.tokens} and {@code app.ai.cost.usd}. Cost is
 * tokens times the configured price. If the provider reported no usage, or the model has no price,
 * it counts the call under {@code app.ai.unpriced} instead of silently recording zero. A response
 * with no usage block arrives as a usage of 0 tokens, not as null, and a real call always has at
 * least one prompt token, so zero is treated as "no usage".
 */
public class UsageCostObservationHandler implements ObservationHandler<ChatClientObservationContext> {

    private final MeterRegistry registry;

    private final Pricing pricing;

    public UsageCostObservationHandler(MeterRegistry registry, Pricing pricing) {
        this.registry = registry;
        this.pricing = pricing;
    }

    @Override
    public boolean supportsContext(Observation.Context context) {
        return context instanceof ChatClientObservationContext;
    }

    @Override
    public void onStop(ChatClientObservationContext context) {
        String endpoint = String.valueOf(context.getRequest().context().getOrDefault(EndpointTagConvention.KEY, "none"));
        ChatClientResponse response = context.getResponse();
        ChatResponse chat = response == null ? null : response.chatResponse();
        Usage usage = chat == null || chat.getMetadata() == null ? null : chat.getMetadata().getUsage();
        String model = chat == null || chat.getMetadata() == null ? null : chat.getMetadata().getModel();
        Pricing.Price price = pricing.forModel(model);
        boolean noUsage = usage == null || usage.getPromptTokens() == null || usage.getPromptTokens() == 0;
        if (noUsage || price == null) {
            Counter.builder("app.ai.unpriced").tag("endpoint", endpoint).tag("reason", noUsage ? "no-usage" : "no-price")
                    .register(registry).increment();
            return;
        }
        long in = usage.getPromptTokens();
        long out = usage.getCompletionTokens() == null ? 0 : usage.getCompletionTokens();
        Counter.builder("app.ai.tokens").tag("endpoint", endpoint).tag("model", model).tag("type", "input").register(registry).increment(in);
        Counter.builder("app.ai.tokens").tag("endpoint", endpoint).tag("model", model).tag("type", "output").register(registry).increment(out);
        Counter.builder("app.ai.cost.usd").tag("endpoint", endpoint).tag("model", model).baseUnit("usd").register(registry)
                .increment(in * price.inputPerMillion() / 1_000_000d + out * price.outputPerMillion() / 1_000_000d);
    }
}