SpringAI instrumentation enhancements

braintrust-java-agent/instrumentation/springai_1_0_0/build.gradle

@@ -0,0 +1,70 @@
+def springAiVersion = '1.0.0'
+
+muzzle {
+  pass {
+    group = 'org.springframework.ai'
+    module = 'spring-ai-openai'
+    versions = "[${springAiVersion},)"
+    ignoredInstrumentation = ["dev.braintrust.instrumentation.springai.v1_0_0.auto.SpringAIAnthropicInstrumentationModule"]
+    extraDependency 'com.fasterxml.jackson.datatype:jackson-datatype-jsr310'
+    extraDependency 'com.fasterxml.jackson.datatype:jackson-datatype-jdk8'
+  }
+  pass {
+    group = 'org.springframework.ai'
+    module = 'spring-ai-anthropic'
+    versions = "[${springAiVersion},)"
+    ignoredInstrumentation = ["dev.braintrust.instrumentation.springai.v1_0_0.auto.SpringAIOpenAIInstrumentationModule"]
+    extraDependency 'com.fasterxml.jackson.datatype:jackson-datatype-jsr310'
+    extraDependency 'com.fasterxml.jackson.datatype:jackson-datatype-jdk8'
+  }
+}
+
+dependencies {
+  implementation project(':braintrust-java-agent:instrumenter')
+  implementation "io.opentelemetry:opentelemetry-api:${otelVersion}"
+  implementation 'com.google.code.findbugs:jsr305:3.0.2'
+  implementation "org.slf4j:slf4j-api:${slf4jVersion}"
+  implementation project(':braintrust-sdk')
+
+  // AutoService for SPI registration
+  compileOnly 'com.google.auto.service:auto-service-annotations:1.1.1'
+  annotationProcessor 'com.google.auto.service:auto-service:1.1.1'
+
+  // ByteBuddy for ElementMatcher types used in instrumentation definitions
+  compileOnly 'net.bytebuddy:byte-buddy:1.17.5'
+
+  // Target libraries — compileOnly because they will be on the app classpath at runtime
+  compileOnly "org.springframework.ai:spring-ai-model:${springAiVersion}"
+  compileOnly "org.springframework.ai:spring-ai-openai:${springAiVersion}"
+  compileOnly "org.springframework.ai:spring-ai-anthropic:${springAiVersion}"
+
+  // Test dependencies
+  testImplementation(testFixtures(project(":test-harness")))
+  testImplementation project(':braintrust-java-agent:instrumenter')
+  testImplementation "org.junit.jupiter:junit-jupiter:${junitVersion}"
+  testRuntimeOnly 'org.junit.platform:junit-platform-launcher'
+  testImplementation 'net.bytebuddy:byte-buddy-agent:1.17.5'
+  testRuntimeOnly "org.slf4j:slf4j-simple:${slf4jVersion}"
+  testImplementation "org.springframework.ai:spring-ai-model:${springAiVersion}"
+  testImplementation "org.springframework.ai:spring-ai-openai:${springAiVersion}"
+  testImplementation "org.springframework.ai:spring-ai-anthropic:${springAiVersion}"
+  // spring-ai-openai and spring-ai-anthropic require spring-webflux at runtime for WebClient
+  testRuntimeOnly 'org.springframework:spring-webflux:6.2.3'
+  // WireMock 3.x bundles Jetty 11, but Spring WebFlux 6.2 requires Jetty 12 for its
+  // JettyClientHttpConnector. Adding reactor-netty-http gives WebFlux a Netty connector to
+  // use for streaming instead, sidestepping the Jetty version conflict entirely.
+  testRuntimeOnly 'io.projectreactor.netty:reactor-netty-http:1.2.3'
+  // Force httpclient5 version to match what spring-ai expects (WireMock pulls in an older one)
+  testImplementation 'org.apache.httpcomponents.client5:httpclient5:5.3.1'
+  testImplementation 'org.apache.httpcomponents.core5:httpcore5:5.2.4'
+}
+
+test {
+  useJUnitPlatform()
+  workingDir = rootProject.projectDir
+  testLogging {
+    events "passed", "skipped", "failed"
+    showStandardStreams = true
+    exceptionFormat "full"
+  }
+}

braintrust-java-agent/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/AnthropicBuilderWrapper.java

@@ -0,0 +1,162 @@
+package dev.braintrust.instrumentation.springai.v1_0_0;
+
+import com.fasterxml.jackson.databind.node.ArrayNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import dev.braintrust.instrumentation.InstrumentationSemConv;
+import dev.braintrust.json.BraintrustJsonMapper;
+import io.micrometer.observation.ObservationRegistry;
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.Tracer;
+import java.lang.reflect.Field;
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
+import java.util.WeakHashMap;
+import lombok.SneakyThrows;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.ai.anthropic.AnthropicChatModel;
+import org.springframework.ai.chat.messages.Message;
+import org.springframework.ai.chat.metadata.ChatResponseMetadata;
+import org.springframework.ai.chat.metadata.Usage;
+import org.springframework.ai.chat.model.ChatResponse;
+
+@Slf4j
+class AnthropicBuilderWrapper {
+    private static final String TRACER_NAME = "braintrust-java";
+    private static final Set<ObservationRegistry> REGISTERED_REGISTRIES =
+            Collections.synchronizedSet(Collections.newSetFromMap(new WeakHashMap<>()));
+
+    /** Reflection-friendly entry point called from {@link BraintrustSpringAI#wrap}. */
+    static void wrap(OpenTelemetry openTelemetry, Object builderObj) {
+        wrap(openTelemetry, (AnthropicChatModel.Builder) builderObj);
+    }
+
+    /** Instruments an {@link AnthropicChatModel.Builder} in place before {@code build()} runs. */
+    static void wrap(OpenTelemetry openTelemetry, AnthropicChatModel.Builder builder) {
+        try {
+            Tracer tracer = openTelemetry.getTracer(TRACER_NAME);
+            ObservationRegistry registry = getField(builder, "observationRegistry");
+            if (registry == null || registry.isNoop()) {
+                registry = ObservationRegistry.create();
+                builder.observationRegistry(registry);
+            }
+            synchronized (REGISTERED_REGISTRIES) {
+                if (!REGISTERED_REGISTRIES.contains(registry)) {
+                    registry.observationConfig()
+                            .observationHandler(
+                                    new BraintrustObservationHandler(
+                                            tracer,
+                                            extractBaseUrl(builder),
+                                            AnthropicBuilderWrapper::tagSpanRequest,
+                                            AnthropicBuilderWrapper::tagSpanResponse));
+                    REGISTERED_REGISTRIES.add(registry);
+                }
+            }
+        } catch (Exception e) {
+            log.error("failed to prepare Spring AI Anthropic builder", e);
+        }
+    }
+
+    // -------------------------------------------------------------------------
+    // Span-tagging helpers
+    // -------------------------------------------------------------------------
+
+    @SneakyThrows
+    static void tagSpanRequest(
+            BraintrustObservationHandler observationHandler,
+            Span span,
+            org.springframework.ai.chat.prompt.Prompt prompt) {
+        ArrayNode messages = BraintrustJsonMapper.get().createArrayNode();
+        for (Message msg : prompt.getInstructions()) {
+            ObjectNode msgNode = BraintrustJsonMapper.get().createObjectNode();
+            msgNode.put("role", msg.getMessageType().getValue().toLowerCase());
+            msgNode.put("content", msg.getText());
+            messages.add(msgNode);
+        }
+        String model = null;
+        if (prompt.getOptions() != null && prompt.getOptions().getModel() != null) {
+            model = prompt.getOptions().getModel().toString();
+        }
+        ObjectNode requestBody = BraintrustJsonMapper.get().createObjectNode();
+        requestBody.set("messages", messages);
+        if (model != null) requestBody.put("model", model);
+
+        InstrumentationSemConv.tagLLMSpanRequest(
+                span,
+                InstrumentationSemConv.PROVIDER_NAME_ANTHROPIC,
+                observationHandler.getBaseUrl(),
+                List.of("v1", "messages"),
+                "POST",
+                BraintrustJsonMapper.toJson(requestBody));
+    }
+
+    @SneakyThrows
+    static void tagSpanResponse(
+            BraintrustObservationHandler observationHandler, Span span, ChatResponse chatResponse) {
+        ArrayNode content = BraintrustJsonMapper.get().createArrayNode();
+        for (var generation : chatResponse.getResults()) {
+            ObjectNode block = BraintrustJsonMapper.get().createObjectNode();
+            block.put("type", "text");
+            block.put("text", generation.getOutput().getText());
+            content.add(block);
+        }
+        ObjectNode responseBody = BraintrustJsonMapper.get().createObjectNode();
+        responseBody.set("content", content);
+
+        ChatResponseMetadata metadata = chatResponse.getMetadata();
+        if (metadata != null && metadata.getUsage() != null) {
+            Usage usage = metadata.getUsage();
+            Integer promptTokens = usage.getPromptTokens();
+            Integer completionTokens = usage.getCompletionTokens();
+            ObjectNode usageNode = BraintrustJsonMapper.get().createObjectNode();
+            if (promptTokens != null) usageNode.put("input_tokens", promptTokens);
+            if (completionTokens != null) usageNode.put("output_tokens", completionTokens);
+            responseBody.set("usage", usageNode);
+        }
+
+        InstrumentationSemConv.tagLLMSpanResponse(
+                span,
+                InstrumentationSemConv.PROVIDER_NAME_ANTHROPIC,
+                BraintrustJsonMapper.toJson(responseBody));
+    }
+
+    private static String extractBaseUrl(AnthropicChatModel.Builder builder) {
+        try {
+            // AnthropicApi doesn't store baseUrl directly; dig through:
+            // builder.anthropicApi -> restClient (DefaultRestClient)
+            //   -> uriBuilderFactory (DefaultUriBuilderFactory) -> baseUri (UriComponentsBuilder)
+            Object anthropicApi = getField(builder, "anthropicApi");
+            Object restClient = getField(anthropicApi, "restClient");
+            Object uriBuilderFactory = getField(restClient, "uriBuilderFactory");
+            Object baseUri = getField(uriBuilderFactory, "baseUri");
+            return (String) baseUri.getClass().getMethod("toUriString").invoke(baseUri);
+        } catch (Exception e) {
+            log.warn("Failed to extract baseUrl from builder", e);
+            return "https://api.anthropic.com";
+        }
+    }
+
+    // -------------------------------------------------------------------------
+    // Internal helpers
+    // -------------------------------------------------------------------------
+
+    @SuppressWarnings("unchecked")
+    private static <T> T getField(Object obj, String fieldName)
+            throws ReflectiveOperationException {
+        Class<?> clazz = obj.getClass();
+        while (clazz != null) {
+            try {
+                Field field = clazz.getDeclaredField(fieldName);
+                field.setAccessible(true);
+                return (T) field.get(obj);
+            } catch (NoSuchFieldException e) {
+                clazz = clazz.getSuperclass();
+            }
+        }
+        throw new NoSuchFieldException(
+                "Field '" + fieldName + "' not found on " + obj.getClass().getName());
+    }
+
+    private AnthropicBuilderWrapper() {}
+}

braintrust-java-agent/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/BraintrustObservationHandler.java

@@ -0,0 +1,95 @@
+package dev.braintrust.instrumentation.springai.v1_0_0;
+
+import dev.braintrust.instrumentation.InstrumentationSemConv;
+import io.micrometer.observation.Observation;
+import io.micrometer.observation.ObservationHandler;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.Tracer;
+import javax.annotation.Nonnull;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.observation.ChatModelObservationContext;
+import org.springframework.ai.chat.prompt.Prompt;
+
+/**
+ * Provider-agnostic Micrometer observation handler for Spring AI chat model calls.
+ *
+ * <p>Starts an OTel span on observation start and ends it on stop/error. Provider-specific request
+ * and response tagging is delegated to the supplied {@code tagRequest} and {@code tagResponse}
+ * callbacks so that OpenAI and Anthropic can each supply the correct format.
+ */
+@Slf4j
+final class BraintrustObservationHandler
+        implements ObservationHandler<ChatModelObservationContext> {
+    private static final String OBSERVATION_SPAN_KEY =
+            BraintrustObservationHandler.class.getName() + ".span";
+
+    private final Tracer tracer;
+    private final TriConsumer<BraintrustObservationHandler, Span, Prompt> tagRequest;
+    private final TriConsumer<BraintrustObservationHandler, Span, ChatResponse> tagResponse;
+    private final String baseUrl;
+
+    BraintrustObservationHandler(
+            Tracer tracer,
+            String baseUrl,
+            TriConsumer<BraintrustObservationHandler, Span, Prompt> tagRequest,
+            TriConsumer<BraintrustObservationHandler, Span, ChatResponse> tagResponse) {
+        this.tracer = tracer;
+        this.baseUrl = baseUrl;
+        this.tagRequest = tagRequest;
+        this.tagResponse = tagResponse;
+    }
+
+    String getBaseUrl() {
+        return this.baseUrl;
+    }
+
+    @Override
+    public boolean supportsContext(@Nonnull Observation.Context context) {
+        return context instanceof ChatModelObservationContext;
+    }
+
+    @Override
+    public void onStart(@Nonnull ChatModelObservationContext context) {
+        try {
+            Span span = tracer.spanBuilder(InstrumentationSemConv.UNSET_LLM_SPAN_NAME).startSpan();
+            context.put(OBSERVATION_SPAN_KEY, span);
+            Prompt prompt = context.getRequest();
+            tagRequest.accept(this, span, prompt);
+        } catch (Exception e) {
+            log.debug("instrumentation error", e);
+        }
+    }
+
+    @Override
+    public void onError(@Nonnull ChatModelObservationContext context) {
+        try {
+            Span span = context.get(OBSERVATION_SPAN_KEY);
+            if (span != null && context.getError() != null) {
+                InstrumentationSemConv.tagLLMSpanResponse(span, context.getError());
+            }
+        } catch (Exception e) {
+            log.debug("instrumentation error", e);
+        }
+    }
+
+    @Override
+    public void onStop(@Nonnull ChatModelObservationContext context) {
+        try {
+            Span span = context.get(OBSERVATION_SPAN_KEY);
+            if (span == null) {
+                return;
+            }
+            try {
+                ChatResponse response = context.getResponse();
+                if (response != null) {
+                    tagResponse.accept(this, span, response);
+                }
+            } finally {
+                span.end();
+            }
+        } catch (Exception e) {
+            log.debug("instrumentation error", e);
+        }
+    }
+}

braintrust-java-agent/instrumentation/springai_1_0_0/src/main/java/dev/braintrust/instrumentation/springai/v1_0_0/BraintrustSpringAI.java

@@ -0,0 +1,51 @@
+package dev.braintrust.instrumentation.springai.v1_0_0;
+
+import io.opentelemetry.api.OpenTelemetry;
+import java.lang.reflect.Method;
+import lombok.extern.slf4j.Slf4j;
+
+/**
+ * Braintrust Spring AI instrumentation entry point.
+ *
+ * <p>Accepts any Spring AI chat-model builder and instruments it in place before {@code build()}
+ * runs. Provider-specific logic lives in {@link OpenAIBuilderWrapper} and {@link
+ * AnthropicBuilderWrapper}, which are only referenced here by string class name so that muzzle does
+ * not follow the reference when a given provider library is absent from the classpath.
+ */
+@Slf4j
+public class BraintrustSpringAI {
+    private static final String OPENAI_BUILDER_CLASS =
+            "org.springframework.ai.openai.OpenAiChatModel$Builder";
+    private static final String ANTHROPIC_BUILDER_CLASS =
+            "org.springframework.ai.anthropic.AnthropicChatModel$Builder";
+
+    private static final String OPENAI_WRAPPER_CLASS =
+            "dev.braintrust.instrumentation.springai.v1_0_0.OpenAIBuilderWrapper";
+    private static final String ANTHROPIC_WRAPPER_CLASS =
+            "dev.braintrust.instrumentation.springai.v1_0_0.AnthropicBuilderWrapper";
+
+    /** Instruments a Spring AI chat-model builder in place. */
+    public static <T> T wrap(OpenTelemetry openTelemetry, T chatModelBuilder) {
+        try {
+            String builderClassName = chatModelBuilder.getClass().getName();
+            String wrapperClass;
+            if (OPENAI_BUILDER_CLASS.equals(builderClassName)) {
+                wrapperClass = OPENAI_WRAPPER_CLASS;
+            } else if (ANTHROPIC_BUILDER_CLASS.equals(builderClassName)) {
+                wrapperClass = ANTHROPIC_WRAPPER_CLASS;
+            } else {
+                log.info("BraintrustSpringAI.wrap: unrecognised builder type {}", builderClassName);
+                return chatModelBuilder;
+            }
+            Class<?> wrapper = chatModelBuilder.getClass().getClassLoader().loadClass(wrapperClass);
+            Method wrapMethod =
+                    wrapper.getDeclaredMethod("wrap", OpenTelemetry.class, Object.class);
+            wrapMethod.invoke(null, openTelemetry, chatModelBuilder);
+        } catch (Exception e) {
+            log.error("failed to apply spring ai instrumentation", e);
+        }
+        return chatModelBuilder;
+    }
+
+    private BraintrustSpringAI() {}
+}