AI Observability

// ❌ BAD — logs full prompts (PII), synchronous audit, hardcoded pricing, no error metrics,
//          @Async via self-invocation (silently runs synchronously)

@Component
@Slf4j
public class AiAuditAdvisor implements CallAdvisor {

    @Autowired                                           // field injection
    private AuditLogRepository auditLogRepository;

    @Override
    public int getOrder() { return 0; }

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

    @Override
    public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {

        // ❌ Logs full prompt content — may contain PII, secrets, customer data
        log.info("AI request prompt: {}", request.prompt().getContents());

        long start = System.currentTimeMillis();
        ChatClientResponse response = chain.nextCall(request);  // no try/catch — errors untracked
        long duration = System.currentTimeMillis() - start;

        // ❌ Hardcoded token pricing — changes when provider updates pricing
        Usage usage = response.chatResponse().getMetadata().getUsage();
        double cost = usage.getPromptTokens() * 0.003 / 1000
            + usage.getCompletionTokens() * 0.015 / 1000;

        // ❌ Synchronous audit save — blocks the main request thread
        // ❌ this.persistAsync(...) would be just as broken: @Async on a method of the
        //    same bean bypasses the proxy and runs inline (self-invocation pitfall)
        auditLogRepository.save(new AiAuditLog(
            request.prompt().getContents(),                              // stores full prompt (PII risk)
            response.chatResponse().getResult().getOutput().getText(),  // stores full response
            cost,
            duration));

        // ❌ No Micrometer metrics — can't alert on latency spikes or error rates
        // ❌ No request ID for tracing
        // ❌ No error handling — if chain throws, no metrics or audit recorded
        // ❌ Only counters, no histograms for latency distribution

        return response;
    }
}

// ✅ GOOD — GA CallAdvisor API, Micrometer metrics, async audit via separate bean, no PII in logs

@Component
@RequiredArgsConstructor
@Slf4j
public class AiAuditAdvisor implements CallAdvisor {

    private final MeterRegistry meterRegistry;
    private final AiAuditRecorder auditRecorder;   // separate bean — @Async needs the proxy

    @Override
    public int getOrder() {
        return 0;
    }

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

    @Override
    public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
        String requestId = UUID.randomUUID().toString();
        Timer.Sample sample = Timer.start(meterRegistry);

        try {
            ChatClientResponse response = chain.nextCall(request);
            ChatResponse chatResponse = response.chatResponse();

            // Record metrics
            sample.stop(Timer.builder("ai.request.duration")
                .tag("status", "success")
                .register(meterRegistry));

            // Track token usage from response metadata
            if (chatResponse != null && chatResponse.getMetadata() != null
                    && chatResponse.getMetadata().getUsage() != null) {
                var usage = chatResponse.getMetadata().getUsage();
                meterRegistry.counter("ai.tokens.input").increment(usage.getPromptTokens());
                meterRegistry.counter("ai.tokens.output").increment(usage.getCompletionTokens()); // GA name
            }

            // Async audit — separate bean, so the @Async proxy is actually used
            auditRecorder.record(requestId, "SUCCESS", null);

            // Log without PII — only metadata
            log.info("AI request {} completed", requestId);

            return response;
        } catch (Exception e) {
            sample.stop(Timer.builder("ai.request.duration")
                .tag("status", "error")
                .register(meterRegistry));

            meterRegistry.counter("ai.request.errors").increment();
            auditRecorder.record(requestId, "ERROR", e.getMessage());

            throw e;
        }
    }
}

// Separate bean: calling an @Async method on `this` bypasses the Spring proxy and
// runs synchronously — the same self-invocation pitfall as @Transactional.
@Component
@RequiredArgsConstructor
class AiAuditRecorder {

    private final AuditLogRepository auditLogRepository;

    @Async
    public void record(String requestId, String status, String errorMessage) {
        auditLogRepository.save(new AiAuditLog(requestId, status, errorMessage, Instant.now()));
    }
}

package com.example.ai.observability;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
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.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Component;

import java.time.Instant;
import java.util.UUID;

/**
 * Production audit advisor for Spring AI 2.0.
 * Tracks latency, token usage, and errors without logging PII.
 *
 * API notes:
 * - CallAroundAdvisor / AdvisedRequest / AdvisedResponse are pre-GA APIs.
 *   The current interface is CallAdvisor with adviseCall(ChatClientRequest, CallAdvisorChain).
 * - Use Usage.getCompletionTokens(); do not call the old getGenerationTokens().
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class AiAuditAdvisor implements CallAdvisor {

    private final AiMetrics metrics;
    private final AiAuditRecorder auditRecorder;   // separate bean — see note on @Async below

    @Override
    public int getOrder() {
        return 0; // runs first in advisor chain
    }

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

    @Override
    public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
        String requestId = UUID.randomUUID().toString();
        Instant startTime = Instant.now();

        String requestedModel = requestedModel(request);

        try {
            ChatClientResponse response = metrics.track("chat", requestedModel, () -> chain.nextCall(request));
            ChatResponse chatResponse = response.chatResponse();

            // Extract token usage if available
            if (chatResponse != null && chatResponse.getMetadata() != null
                    && chatResponse.getMetadata().getUsage() != null) {
                var usage = chatResponse.getMetadata().getUsage();
                long inputTokens = usage.getPromptTokens();
                long outputTokens = usage.getCompletionTokens();
                // The response reports the model that actually served the call.
                String servedModel = chatResponse.getMetadata().getModel();
                metrics.recordTokens(servedModel != null ? servedModel : requestedModel, inputTokens, outputTokens);

                // Log metadata only — never log prompt content (PII risk)
                log.info("AI request completed: requestId={}, inputTokens={}, outputTokens={}",
                    requestId, inputTokens, outputTokens);
            }

            // Persist audit asynchronously — don't block the response
            recordBestEffort(requestId, "SUCCESS", null, startTime);

            return response;
        } catch (Exception e) {
            log.error("AI request failed: requestId={}, failureType={}", requestId, e.getClass().getSimpleName());
            recordBestEffort(requestId, "ERROR", e.getClass().getSimpleName(), startTime);
            throw e;
        }
    }

    private void recordBestEffort(String requestId, String status, String errorType, Instant startTime) {
        try {
            auditRecorder.record(requestId, status, errorType, startTime);
        } catch (RuntimeException failure) {
            // Submission can fail before the @Async method body runs.
            log.warn("AI audit submission failed: requestId={}, failureType={}",
                requestId, failure.getClass().getSimpleName());
        }
    }

    private static String requestedModel(ChatClientRequest request) {
        ChatOptions options = request.prompt().getOptions();
        return options != null && options.getModel() != null ? options.getModel() : "default";
    }
}

/**
 * Async persistence lives in its OWN bean on purpose: an @Async method invoked on
 * `this` bypasses the Spring proxy and runs synchronously (self-invocation pitfall,
 * same as @Transactional). A separate component goes through the proxy.
 */
@Component
@RequiredArgsConstructor
class AiAuditRecorder {

    private final AiAuditLogRepository auditRepository;

    @Async
    public void record(String requestId, String status, String errorMessage, Instant startTime) {
        try {
            auditRepository.save(new AiAuditLog(
                UUID.fromString(requestId),
                status,
                errorMessage,
                startTime,
                Instant.now()
            ));
        } catch (Exception e) {
            // Never let audit failures break the AI call path
            org.slf4j.LoggerFactory.getLogger(AiAuditRecorder.class)
                .warn("Failed to persist AI audit log: requestId={}, failureType={}",
                    requestId, e.getClass().getSimpleName());
        }
    }
}

package com.example.ai.observability;

import io.micrometer.core.instrument.Counter;
import io.micrometer.core.instrument.MeterRegistry;
import io.micrometer.core.instrument.Timer;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Component;

import java.util.function.Supplier;

/**
 * Centralized AI metrics using Micrometer.
 * Records latency, token usage, errors, and cost estimates.
 */
@Component
@RequiredArgsConstructor
public class AiMetrics {

    private final MeterRegistry registry;

    /**
     * Time an AI operation and record success/failure.
     */
    public <T> T track(String operation, String model, Supplier<T> action) {
        Timer.Sample sample = Timer.start(registry);
        try {
            T result = action.get();
            sample.stop(timer(operation, model, "success"));
            return result;
        } catch (Exception e) {
            sample.stop(timer(operation, model, "error"));
            errorCounter(operation, model).increment();
            throw e;
        }
    }

    /**
     * Record token usage from a completed AI response.
     */
    public void recordTokens(String model, long inputTokens, long outputTokens) {
        counter("ai.tokens", "model", model, "type", "input").increment(inputTokens);
        counter("ai.tokens", "model", model, "type", "output").increment(outputTokens);
        counter("ai.tokens.total", "model", model).increment(inputTokens + outputTokens);
    }

    /**
     * Record estimated cost (pass rate from configuration, not hardcoded).
     */
    public void recordCost(String model, double estimatedCost) {
        registry.counter("ai.cost.estimated", "model", model).increment(estimatedCost);
    }

    private Timer timer(String operation, String model, String status) {
        return Timer.builder("ai.request.duration")
            .tag("operation", operation)
            .tag("model", model)
            .tag("status", status)
            .description("AI request latency")
            .register(registry);
    }

    private Counter errorCounter(String operation, String model) {
        return counter("ai.request.errors", "operation", operation, "model", model);
    }

    private Counter counter(String name, String... tags) {
        return Counter.builder(name).tags(tags).register(registry);
    }
}