Skip to content

Commit 7797207

Browse files
realarkrealark
committed
enhance SpringAI instrumentation
1 parent 22507f1 commit 7797207

File tree

11 files changed

+1133
-7
lines changed

11 files changed

+1133
-7
lines changed

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

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
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+
// Force httpclient5 version to match what spring-ai expects (WireMock pulls in an older one)
50+
testImplementation 'org.apache.httpcomponents.client5:httpclient5:5.3.1'
51+
testImplementation 'org.apache.httpcomponents.core5:httpcore5:5.2.4'
52+
}
53+
54+
test {
55+
useJUnitPlatform()
56+
workingDir = rootProject.projectDir
57+
testLogging {
58+
events "passed", "skipped", "failed"
59+
showStandardStreams = true
60+
exceptionFormat "full"
61+
}
62+
}

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

Lines changed: 208 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,208 @@
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.Observation;
8+
import io.micrometer.observation.ObservationHandler;
9+
import io.micrometer.observation.ObservationRegistry;
10+
import io.opentelemetry.api.OpenTelemetry;
11+
import io.opentelemetry.api.trace.Span;
12+
import io.opentelemetry.api.trace.Tracer;
13+
import java.lang.reflect.Field;
14+
import java.util.Collections;
15+
import java.util.List;
16+
import java.util.Map;
17+
import java.util.WeakHashMap;
18+
import lombok.SneakyThrows;
19+
import lombok.extern.slf4j.Slf4j;
20+
import org.springframework.ai.anthropic.AnthropicChatModel;
21+
import org.springframework.ai.chat.messages.Message;
22+
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
23+
import org.springframework.ai.chat.metadata.Usage;
24+
import org.springframework.ai.chat.model.ChatResponse;
25+
import org.springframework.ai.chat.observation.ChatModelObservationContext;
26+
import org.springframework.ai.chat.prompt.Prompt;
27+
28+
/** Braintrust Spring AI Anthropic instrumentation entry point. */
29+
@Slf4j
30+
class AnthropicBuilderWrapper {
31+
private static final String TRACER_NAME = "braintrust-java";
32+
private static final Map<ObservationRegistry, Boolean> REGISTERED_REGISTRIES =
33+
Collections.synchronizedMap(new WeakHashMap<>());
34+
35+
/** Reflection-friendly entry point called from {@link BraintrustSpringAI#wrap}. */
36+
public static void wrap(OpenTelemetry openTelemetry, Object builderObj) {
37+
wrap(openTelemetry, (AnthropicChatModel.Builder) builderObj);
38+
}
39+
40+
/** Instruments an {@link AnthropicChatModel.Builder} in place before {@code build()} runs. */
41+
static void wrap(OpenTelemetry openTelemetry, AnthropicChatModel.Builder builder) {
42+
try {
43+
Tracer tracer = openTelemetry.getTracer(TRACER_NAME);
44+
ObservationRegistry registry = getField(builder, "observationRegistry");
45+
if (registry == null || registry.isNoop()) {
46+
registry = ObservationRegistry.create();
47+
builder.observationRegistry(registry);
48+
}
49+
synchronized (REGISTERED_REGISTRIES) {
50+
if (!REGISTERED_REGISTRIES.containsKey(registry)) {
51+
registry.observationConfig()
52+
.observationHandler(new BraintrustAnthropicObservationHandler(tracer));
53+
REGISTERED_REGISTRIES.put(registry, Boolean.TRUE);
54+
}
55+
}
56+
} catch (Exception e) {
57+
log.error("failed to prepare Spring AI Anthropic builder", e);
58+
}
59+
}
60+
61+
// -------------------------------------------------------------------------
62+
// Shared span-tagging helpers
63+
// -------------------------------------------------------------------------
64+
65+
@SneakyThrows
66+
static void tagSpanRequest(Span span, Prompt prompt) {
67+
ArrayNode messages = BraintrustJsonMapper.get().createArrayNode();
68+
for (Message msg : prompt.getInstructions()) {
69+
ObjectNode msgNode = BraintrustJsonMapper.get().createObjectNode();
70+
msgNode.put("role", msg.getMessageType().getValue().toLowerCase());
71+
msgNode.put("content", msg.getText());
72+
messages.add(msgNode);
73+
}
74+
75+
String model = null;
76+
if (prompt.getOptions() != null) {
77+
Object modelOpt = prompt.getOptions().getModel();
78+
if (modelOpt != null) {
79+
model = modelOpt.toString();
80+
}
81+
}
82+
83+
ObjectNode requestBody = BraintrustJsonMapper.get().createObjectNode();
84+
requestBody.set("messages", messages);
85+
if (model != null) {
86+
requestBody.put("model", model);
87+
}
88+
89+
InstrumentationSemConv.tagLLMSpanRequest(
90+
span,
91+
InstrumentationSemConv.PROVIDER_NAME_ANTHROPIC,
92+
"https://api.anthropic.com",
93+
List.of("v1", "messages"),
94+
"POST",
95+
BraintrustJsonMapper.toJson(requestBody));
96+
}
97+
98+
@SneakyThrows
99+
static void tagSpanResponse(Span span, ChatResponse chatResponse) {
100+
// Build an Anthropic-shaped response body so InstrumentationSemConv.tagAnthropicResponse
101+
// can parse it: {"content":[{"type":"text","text":"..."}],
102+
// "usage":{"input_tokens":N,"output_tokens":N}}
103+
ArrayNode content = BraintrustJsonMapper.get().createArrayNode();
104+
for (var generation : chatResponse.getResults()) {
105+
ObjectNode block = BraintrustJsonMapper.get().createObjectNode();
106+
block.put("type", "text");
107+
block.put("text", generation.getOutput().getText());
108+
content.add(block);
109+
}
110+
111+
ObjectNode responseBody = BraintrustJsonMapper.get().createObjectNode();
112+
responseBody.set("content", content);
113+
114+
ChatResponseMetadata metadata = chatResponse.getMetadata();
115+
if (metadata != null && metadata.getUsage() != null) {
116+
Usage usage = metadata.getUsage();
117+
Integer promptTokens = usage.getPromptTokens();
118+
Integer completionTokens = usage.getCompletionTokens();
119+
ObjectNode usageNode = BraintrustJsonMapper.get().createObjectNode();
120+
// Anthropic wire format uses input_tokens / output_tokens
121+
if (promptTokens != null) usageNode.put("input_tokens", promptTokens);
122+
if (completionTokens != null) usageNode.put("output_tokens", completionTokens);
123+
responseBody.set("usage", usageNode);
124+
}
125+
126+
InstrumentationSemConv.tagLLMSpanResponse(
127+
span,
128+
InstrumentationSemConv.PROVIDER_NAME_ANTHROPIC,
129+
BraintrustJsonMapper.toJson(responseBody));
130+
}
131+
132+
// -------------------------------------------------------------------------
133+
// Internal helpers
134+
// -------------------------------------------------------------------------
135+
136+
@SuppressWarnings("unchecked")
137+
private static <T> T getField(Object obj, String fieldName)
138+
throws ReflectiveOperationException {
139+
Class<?> clazz = obj.getClass();
140+
while (clazz != null) {
141+
try {
142+
Field field = clazz.getDeclaredField(fieldName);
143+
field.setAccessible(true);
144+
return (T) field.get(obj);
145+
} catch (NoSuchFieldException e) {
146+
clazz = clazz.getSuperclass();
147+
}
148+
}
149+
throw new NoSuchFieldException(
150+
"Field '" + fieldName + "' not found on " + obj.getClass().getName());
151+
}
152+
153+
static final class BraintrustAnthropicObservationHandler
154+
implements ObservationHandler<ChatModelObservationContext> {
155+
private static final String OBSERVATION_SPAN_KEY =
156+
BraintrustAnthropicObservationHandler.class.getName() + ".span";
157+
158+
private final Tracer tracer;
159+
160+
BraintrustAnthropicObservationHandler(Tracer tracer) {
161+
this.tracer = tracer;
162+
}
163+
164+
@Override
165+
public boolean supportsContext(Observation.Context context) {
166+
return context instanceof ChatModelObservationContext;
167+
}
168+
169+
@Override
170+
public void onStart(ChatModelObservationContext context) {
171+
Span span = tracer.spanBuilder(InstrumentationSemConv.UNSET_LLM_SPAN_NAME).startSpan();
172+
context.put(OBSERVATION_SPAN_KEY, span);
173+
Prompt prompt = context.getRequest();
174+
if (prompt != null) {
175+
tagSpanRequest(span, prompt);
176+
}
177+
}
178+
179+
@Override
180+
public void onError(ChatModelObservationContext context) {
181+
Span span = context.get(OBSERVATION_SPAN_KEY);
182+
if (span != null) {
183+
Throwable error = context.getError();
184+
if (error != null) {
185+
InstrumentationSemConv.tagLLMSpanResponse(span, error);
186+
}
187+
}
188+
}
189+
190+
@Override
191+
public void onStop(ChatModelObservationContext context) {
192+
Span span = context.get(OBSERVATION_SPAN_KEY);
193+
if (span == null) {
194+
return;
195+
}
196+
try {
197+
ChatResponse response = context.getResponse();
198+
if (response != null) {
199+
tagSpanResponse(span, response);
200+
}
201+
} finally {
202+
span.end();
203+
}
204+
}
205+
}
206+
207+
private AnthropicBuilderWrapper() {}
208+
}

‎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)