Skip to content

Commit 9b860f0

Browse files
realarkrealark
committed
enhance SpringAI instrumentation
1 parent f5bfdec commit 9b860f0

13 files changed

Lines changed: 1005 additions & 9 deletions

File tree

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

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

‎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.Set;
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 Set<ObservationRegistry> REGISTERED_REGISTRIES =
28+
Collections.synchronizedSet(Collections.newSetFromMap(new WeakHashMap<>()));
29+
30+
/** Reflection-friendly entry point called from {@link BraintrustSpringAI#wrap}. */
31+
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.contains(registry)) {
46+
registry.observationConfig()
47+
.observationHandler(
48+
new BraintrustObservationHandler(
49+
tracer,
50+
AnthropicBuilderWrapper::tagSpanRequest,
51+
AnthropicBuilderWrapper::tagSpanResponse));
52+
REGISTERED_REGISTRIES.add(registry);
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: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
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+
/** Instruments a Spring AI chat-model builder in place. */
28+
public static <T> T wrap(OpenTelemetry openTelemetry, T chatModelBuilder) {
29+
try {
30+
String builderClassName = chatModelBuilder.getClass().getName();
31+
String wrapperClass;
32+
if (OPENAI_BUILDER_CLASS.equals(builderClassName)) {
33+
wrapperClass = OPENAI_WRAPPER_CLASS;
34+
} else if (ANTHROPIC_BUILDER_CLASS.equals(builderClassName)) {
35+
wrapperClass = ANTHROPIC_WRAPPER_CLASS;
36+
} else {
37+
log.info("BraintrustSpringAI.wrap: unrecognised builder type {}", builderClassName);
38+
return chatModelBuilder;
39+
}
40+
Class<?> wrapper = chatModelBuilder.getClass().getClassLoader().loadClass(wrapperClass);
41+
Method wrapMethod =
42+
wrapper.getDeclaredMethod("wrap", OpenTelemetry.class, Object.class);
43+
wrapMethod.invoke(null, openTelemetry, chatModelBuilder);
44+
} catch (Exception e) {
45+
log.error("failed to apply spring ai instrumentation", e);
46+
}
47+
return chatModelBuilder;
48+
}
49+
50+
private BraintrustSpringAI() {}
51+
}

0 commit comments

Comments
 (0)