Skip to content

Commit bfab183

Browse files
realarkrealark
committed
enhance SpringAI instrumentation
1 parent 15b6c5f commit bfab183

File tree

14 files changed

+1095
-9
lines changed

14 files changed

+1095
-9
lines changed

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

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
def springAiVersion = '1.0.0'
2+
3+
muzzle {
4+
pass {
5+
group = 'org.springframework.ai'
6+
module = 'spring-ai-openai'
7+
versions = "[${springAiVersion},)"
8+
ignoredInstrumentation = ["dev.braintrust.instrumentation.springai.v1_0_0.auto.SpringAIAnthropicInstrumentationModule"]
9+
}
10+
pass {
11+
group = 'org.springframework.ai'
12+
module = 'spring-ai-anthropic'
13+
versions = "[${springAiVersion},)"
14+
ignoredInstrumentation = ["dev.braintrust.instrumentation.springai.v1_0_0.auto.SpringAIOpenAIInstrumentationModule"]
15+
}
16+
}
17+
18+
dependencies {
19+
implementation project(':braintrust-java-agent:instrumenter')
20+
implementation "io.opentelemetry:opentelemetry-api:${otelVersion}"
21+
implementation 'com.google.code.findbugs:jsr305:3.0.2'
22+
implementation "org.slf4j:slf4j-api:${slf4jVersion}"
23+
implementation project(':braintrust-sdk')
24+
25+
// AutoService for SPI registration
26+
compileOnly 'com.google.auto.service:auto-service-annotations:1.1.1'
27+
annotationProcessor 'com.google.auto.service:auto-service:1.1.1'
28+
29+
// ByteBuddy for ElementMatcher types used in instrumentation definitions
30+
compileOnly 'net.bytebuddy:byte-buddy:1.17.5'
31+
32+
// Target libraries — compileOnly because they will be on the app classpath at runtime
33+
compileOnly "org.springframework.ai:spring-ai-model:${springAiVersion}"
34+
compileOnly "org.springframework.ai:spring-ai-openai:${springAiVersion}"
35+
compileOnly "org.springframework.ai:spring-ai-anthropic:${springAiVersion}"
36+
37+
// Test dependencies
38+
testImplementation(testFixtures(project(":test-harness")))
39+
testImplementation project(':braintrust-java-agent:instrumenter')
40+
testImplementation "org.junit.jupiter:junit-jupiter:${junitVersion}"
41+
testRuntimeOnly 'org.junit.platform:junit-platform-launcher'
42+
testImplementation 'net.bytebuddy:byte-buddy-agent:1.17.5'
43+
testRuntimeOnly "org.slf4j:slf4j-simple:${slf4jVersion}"
44+
testImplementation "org.springframework.ai:spring-ai-model:${springAiVersion}"
45+
testImplementation "org.springframework.ai:spring-ai-openai:${springAiVersion}"
46+
testImplementation "org.springframework.ai:spring-ai-anthropic:${springAiVersion}"
47+
// spring-ai-openai and spring-ai-anthropic require spring-webflux at runtime for WebClient
48+
testRuntimeOnly 'org.springframework:spring-webflux:6.2.3'
49+
// WireMock 3.x bundles Jetty 11, but Spring WebFlux 6.2 requires Jetty 12 for its
50+
// JettyClientHttpConnector. Adding reactor-netty-http gives WebFlux a Netty connector to
51+
// use for streaming instead, sidestepping the Jetty version conflict entirely.
52+
testRuntimeOnly 'io.projectreactor.netty:reactor-netty-http:1.2.3'
53+
// Force httpclient5 version to match what spring-ai expects (WireMock pulls in an older one)
54+
testImplementation 'org.apache.httpcomponents.client5:httpclient5:5.3.1'
55+
testImplementation 'org.apache.httpcomponents.core5:httpcore5:5.2.4'
56+
}
57+
58+
test {
59+
useJUnitPlatform()
60+
workingDir = rootProject.projectDir
61+
testLogging {
62+
events "passed", "skipped", "failed"
63+
showStandardStreams = true
64+
exceptionFormat "full"
65+
}
66+
}

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

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
1+
package dev.braintrust.instrumentation.springai.v1_0_0;
2+
3+
import com.fasterxml.jackson.databind.node.ArrayNode;
4+
import com.fasterxml.jackson.databind.node.ObjectNode;
5+
import dev.braintrust.instrumentation.InstrumentationSemConv;
6+
import dev.braintrust.json.BraintrustJsonMapper;
7+
import io.micrometer.observation.ObservationRegistry;
8+
import io.opentelemetry.api.OpenTelemetry;
9+
import io.opentelemetry.api.trace.Span;
10+
import io.opentelemetry.api.trace.Tracer;
11+
import java.lang.reflect.Field;
12+
import java.util.Collections;
13+
import java.util.List;
14+
import java.util.Map;
15+
import java.util.WeakHashMap;
16+
import lombok.extern.slf4j.Slf4j;
17+
import org.springframework.ai.anthropic.AnthropicChatModel;
18+
import org.springframework.ai.chat.messages.Message;
19+
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
20+
import org.springframework.ai.chat.metadata.Usage;
21+
import org.springframework.ai.chat.model.ChatResponse;
22+
23+
/** Braintrust Spring AI Anthropic instrumentation entry point. */
24+
@Slf4j
25+
class AnthropicBuilderWrapper {
26+
private static final String TRACER_NAME = "braintrust-java";
27+
private static final Map<ObservationRegistry, Boolean> REGISTERED_REGISTRIES =
28+
Collections.synchronizedMap(new WeakHashMap<>());
29+
30+
/** Reflection-friendly entry point called from {@link BraintrustSpringAI#wrap}. */
31+
public static void wrap(OpenTelemetry openTelemetry, Object builderObj) {
32+
wrap(openTelemetry, (AnthropicChatModel.Builder) builderObj);
33+
}
34+
35+
/** Instruments an {@link AnthropicChatModel.Builder} in place before {@code build()} runs. */
36+
static void wrap(OpenTelemetry openTelemetry, AnthropicChatModel.Builder builder) {
37+
try {
38+
Tracer tracer = openTelemetry.getTracer(TRACER_NAME);
39+
ObservationRegistry registry = getField(builder, "observationRegistry");
40+
if (registry == null || registry.isNoop()) {
41+
registry = ObservationRegistry.create();
42+
builder.observationRegistry(registry);
43+
}
44+
synchronized (REGISTERED_REGISTRIES) {
45+
if (!REGISTERED_REGISTRIES.containsKey(registry)) {
46+
registry.observationConfig()
47+
.observationHandler(
48+
new BraintrustObservationHandler(
49+
tracer,
50+
AnthropicBuilderWrapper::tagSpanRequest,
51+
AnthropicBuilderWrapper::tagSpanResponse));
52+
REGISTERED_REGISTRIES.put(registry, Boolean.TRUE);
53+
}
54+
}
55+
} catch (Exception e) {
56+
log.error("failed to prepare Spring AI Anthropic builder", e);
57+
}
58+
}
59+
60+
// -------------------------------------------------------------------------
61+
// Span-tagging helpers
62+
// -------------------------------------------------------------------------
63+
64+
@lombok.SneakyThrows
65+
static void tagSpanRequest(Span span, org.springframework.ai.chat.prompt.Prompt prompt) {
66+
ArrayNode messages = BraintrustJsonMapper.get().createArrayNode();
67+
for (Message msg : prompt.getInstructions()) {
68+
ObjectNode msgNode = BraintrustJsonMapper.get().createObjectNode();
69+
msgNode.put("role", msg.getMessageType().getValue().toLowerCase());
70+
msgNode.put("content", msg.getText());
71+
messages.add(msgNode);
72+
}
73+
String model = null;
74+
if (prompt.getOptions() != null && prompt.getOptions().getModel() != null) {
75+
model = prompt.getOptions().getModel().toString();
76+
}
77+
ObjectNode requestBody = BraintrustJsonMapper.get().createObjectNode();
78+
requestBody.set("messages", messages);
79+
if (model != null) requestBody.put("model", model);
80+
81+
InstrumentationSemConv.tagLLMSpanRequest(
82+
span,
83+
InstrumentationSemConv.PROVIDER_NAME_ANTHROPIC,
84+
"https://api.anthropic.com",
85+
List.of("v1", "messages"),
86+
"POST",
87+
BraintrustJsonMapper.toJson(requestBody));
88+
}
89+
90+
@lombok.SneakyThrows
91+
static void tagSpanResponse(Span span, ChatResponse chatResponse) {
92+
ArrayNode content = BraintrustJsonMapper.get().createArrayNode();
93+
for (var generation : chatResponse.getResults()) {
94+
ObjectNode block = BraintrustJsonMapper.get().createObjectNode();
95+
block.put("type", "text");
96+
block.put("text", generation.getOutput().getText());
97+
content.add(block);
98+
}
99+
ObjectNode responseBody = BraintrustJsonMapper.get().createObjectNode();
100+
responseBody.set("content", content);
101+
102+
ChatResponseMetadata metadata = chatResponse.getMetadata();
103+
if (metadata != null && metadata.getUsage() != null) {
104+
Usage usage = metadata.getUsage();
105+
Integer promptTokens = usage.getPromptTokens();
106+
Integer completionTokens = usage.getCompletionTokens();
107+
ObjectNode usageNode = BraintrustJsonMapper.get().createObjectNode();
108+
if (promptTokens != null) usageNode.put("input_tokens", promptTokens);
109+
if (completionTokens != null) usageNode.put("output_tokens", completionTokens);
110+
responseBody.set("usage", usageNode);
111+
}
112+
113+
InstrumentationSemConv.tagLLMSpanResponse(
114+
span,
115+
InstrumentationSemConv.PROVIDER_NAME_ANTHROPIC,
116+
BraintrustJsonMapper.toJson(responseBody));
117+
}
118+
119+
// -------------------------------------------------------------------------
120+
// Internal helpers
121+
// -------------------------------------------------------------------------
122+
123+
@SuppressWarnings("unchecked")
124+
private static <T> T getField(Object obj, String fieldName)
125+
throws ReflectiveOperationException {
126+
Class<?> clazz = obj.getClass();
127+
while (clazz != null) {
128+
try {
129+
Field field = clazz.getDeclaredField(fieldName);
130+
field.setAccessible(true);
131+
return (T) field.get(obj);
132+
} catch (NoSuchFieldException e) {
133+
clazz = clazz.getSuperclass();
134+
}
135+
}
136+
throw new NoSuchFieldException(
137+
"Field '" + fieldName + "' not found on " + obj.getClass().getName());
138+
}
139+
140+
private AnthropicBuilderWrapper() {}
141+
}

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

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
package dev.braintrust.instrumentation.springai.v1_0_0;
2+
3+
import dev.braintrust.instrumentation.InstrumentationSemConv;
4+
import io.micrometer.observation.Observation;
5+
import io.micrometer.observation.ObservationHandler;
6+
import io.opentelemetry.api.trace.Span;
7+
import io.opentelemetry.api.trace.Tracer;
8+
import java.util.function.BiConsumer;
9+
import org.springframework.ai.chat.model.ChatResponse;
10+
import org.springframework.ai.chat.observation.ChatModelObservationContext;
11+
import org.springframework.ai.chat.prompt.Prompt;
12+
13+
/**
14+
* Provider-agnostic Micrometer observation handler for Spring AI chat model calls.
15+
*
16+
* <p>Starts an OTel span on observation start and ends it on stop/error. Provider-specific request
17+
* and response tagging is delegated to the supplied {@code tagRequest} and {@code tagResponse}
18+
* callbacks so that OpenAI and Anthropic can each supply the correct format.
19+
*/
20+
final class BraintrustObservationHandler
21+
implements ObservationHandler<ChatModelObservationContext> {
22+
private static final String OBSERVATION_SPAN_KEY =
23+
BraintrustObservationHandler.class.getName() + ".span";
24+
25+
private final Tracer tracer;
26+
private final BiConsumer<Span, Prompt> tagRequest;
27+
private final BiConsumer<Span, ChatResponse> tagResponse;
28+
29+
BraintrustObservationHandler(
30+
Tracer tracer,
31+
BiConsumer<Span, Prompt> tagRequest,
32+
BiConsumer<Span, ChatResponse> tagResponse) {
33+
this.tracer = tracer;
34+
this.tagRequest = tagRequest;
35+
this.tagResponse = tagResponse;
36+
}
37+
38+
@Override
39+
public boolean supportsContext(Observation.Context context) {
40+
return context instanceof ChatModelObservationContext;
41+
}
42+
43+
@Override
44+
public void onStart(ChatModelObservationContext context) {
45+
Span span = tracer.spanBuilder(InstrumentationSemConv.UNSET_LLM_SPAN_NAME).startSpan();
46+
context.put(OBSERVATION_SPAN_KEY, span);
47+
Prompt prompt = context.getRequest();
48+
if (prompt != null) {
49+
tagRequest.accept(span, prompt);
50+
}
51+
}
52+
53+
@Override
54+
public void onError(ChatModelObservationContext context) {
55+
Span span = context.get(OBSERVATION_SPAN_KEY);
56+
if (span != null && context.getError() != null) {
57+
InstrumentationSemConv.tagLLMSpanResponse(span, context.getError());
58+
}
59+
}
60+
61+
@Override
62+
public void onStop(ChatModelObservationContext context) {
63+
Span span = context.get(OBSERVATION_SPAN_KEY);
64+
if (span == null) {
65+
return;
66+
}
67+
try {
68+
ChatResponse response = context.getResponse();
69+
if (response != null) {
70+
tagResponse.accept(span, response);
71+
}
72+
} finally {
73+
span.end();
74+
}
75+
}
76+
}

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

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
package dev.braintrust.instrumentation.springai.v1_0_0;
2+
3+
import io.opentelemetry.api.OpenTelemetry;
4+
import java.lang.reflect.Method;
5+
import lombok.extern.slf4j.Slf4j;
6+
7+
/**
8+
* Braintrust Spring AI instrumentation entry point.
9+
*
10+
* <p>Accepts any Spring AI chat-model builder and instruments it in place before {@code build()}
11+
* runs. Provider-specific logic lives in {@link OpenAIBuilderWrapper} and {@link
12+
* AnthropicBuilderWrapper}, which are only referenced here by string class name so that muzzle does
13+
* not follow the reference when a given provider library is absent from the classpath.
14+
*/
15+
@Slf4j
16+
public class BraintrustSpringAI {
17+
private static final String OPENAI_BUILDER_CLASS =
18+
"org.springframework.ai.openai.OpenAiChatModel$Builder";
19+
private static final String ANTHROPIC_BUILDER_CLASS =
20+
"org.springframework.ai.anthropic.AnthropicChatModel$Builder";
21+
22+
private static final String OPENAI_WRAPPER_CLASS =
23+
"dev.braintrust.instrumentation.springai.v1_0_0.OpenAIBuilderWrapper";
24+
private static final String ANTHROPIC_WRAPPER_CLASS =
25+
"dev.braintrust.instrumentation.springai.v1_0_0.AnthropicBuilderWrapper";
26+
27+
/**
28+
* Instruments a Spring AI chat-model builder in place.
29+
*
30+
* <p>Dispatches to the appropriate provider wrapper based on the builder's runtime class name.
31+
* Uses reflection so that neither this class nor its callers create a compile-time reference to
32+
* provider-specific types — allowing muzzle to verify each provider independently.
33+
*/
34+
public static <T> T wrap(OpenTelemetry openTelemetry, T builder) {
35+
try {
36+
String builderClassName = builder.getClass().getName();
37+
String wrapperClass;
38+
if (OPENAI_BUILDER_CLASS.equals(builderClassName)) {
39+
wrapperClass = OPENAI_WRAPPER_CLASS;
40+
} else if (ANTHROPIC_BUILDER_CLASS.equals(builderClassName)) {
41+
wrapperClass = ANTHROPIC_WRAPPER_CLASS;
42+
} else {
43+
log.debug(
44+
"BraintrustSpringAI.wrap: unrecognised builder type {}", builderClassName);
45+
return builder;
46+
}
47+
Class<?> wrapper = builder.getClass().getClassLoader().loadClass(wrapperClass);
48+
Method wrapMethod =
49+
wrapper.getDeclaredMethod("wrap", OpenTelemetry.class, Object.class);
50+
wrapMethod.invoke(null, openTelemetry, builder);
51+
} catch (Exception e) {
52+
log.error("failed to apply spring ai instrumentation", e);
53+
}
54+
return builder;
55+
}
56+
57+
private BraintrustSpringAI() {}
58+
}

0 commit comments

Comments
 (0)