Skip to content

Commit c0ffa18

Browse files
realarkrealark
authored
Merge pull request #161 from braintrustdata/ark/sdk-113-web-search
add web search server-side tool spans to anthropic and openai instrumentation
2 parents 50bd1b4 + 1591cad commit c0ffa18

44 files changed

Lines changed: 2692 additions & 59 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎braintrust-sdk/instrumentation/anthropic_2_2_0/src/main/java/dev/braintrust/instrumentation/anthropic/v2_2_0/TracingHttpClient.java‎

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,7 @@ public void close() {
128128
inputJson);
129129

130130
var response = underlying.execute(bufferedRequest, requestOptions);
131-
return new TeeingStreamHttpResponse(response, span);
131+
return new TeeingStreamHttpResponse(response, span, tracer);
132132
} catch (Exception e) {
133133
InstrumentationSemConv.tagLLMSpanResponse(span, e);
134134
span.end();
@@ -157,7 +157,9 @@ public void close() {
157157
return underlying
158158
.executeAsync(bufferedRequest, requestOptions)
159159
.thenApply(
160-
response -> (HttpResponse) new TeeingStreamHttpResponse(response, span))
160+
response ->
161+
(HttpResponse)
162+
new TeeingStreamHttpResponse(response, span, tracer))
161163
.whenComplete(
162164
(response, t) -> {
163165
if (t != null) {
@@ -237,14 +239,16 @@ private static String readBodyAsString(HttpRequestBody body) {
237239
private static final class TeeingStreamHttpResponse implements HttpResponse {
238240
private final HttpResponse delegate;
239241
private final Span span;
242+
private final Tracer tracer;
240243
private final long spanStartNanos = System.nanoTime();
241244
private final AtomicLong timeToFirstTokenNanos = new AtomicLong();
242245
private final ByteArrayOutputStream teeBuffer = new ByteArrayOutputStream();
243246
private final InputStream teeStream;
244247

245-
TeeingStreamHttpResponse(HttpResponse delegate, Span span) {
248+
TeeingStreamHttpResponse(HttpResponse delegate, Span span, Tracer tracer) {
246249
this.delegate = delegate;
247250
this.span = span;
251+
this.tracer = tracer;
248252
this.teeStream =
249253
new TeeInputStream(
250254
delegate.body(), teeBuffer, this::onFirstByte, this::onStreamClosed);
@@ -260,7 +264,9 @@ private void onStreamClosed() {
260264
synchronized (teeBuffer) {
261265
bytes = teeBuffer.toByteArray();
262266
}
263-
tagSpanFromBuffer(span, bytes, timeToFirstTokenNanos.get());
267+
// tagLLMSpanResponse also emits child spans for any server-side tool calls (web
268+
// search, etc.) nested under the LLM span while it is still live.
269+
tagSpanFromBuffer(tracer, span, bytes, timeToFirstTokenNanos.get());
264270
} finally {
265271
span.end();
266272
}
@@ -354,7 +360,8 @@ private void notifyClosed() {
354360
// Span tagging from buffered bytes
355361
// -------------------------------------------------------------------------
356362

357-
private static void tagSpanFromBuffer(Span span, byte[] bytes, Long timeToFirstTokenNanos) {
363+
private static void tagSpanFromBuffer(
364+
Tracer tracer, Span span, byte[] bytes, Long timeToFirstTokenNanos) {
358365
if (bytes.length == 0) return;
359366
try {
360367
String firstLine = firstNonEmptyLine(bytes);
@@ -364,13 +371,15 @@ private static void tagSpanFromBuffer(Span span, byte[] bytes, Long timeToFirstT
364371
firstLine != null
365372
&& (firstLine.startsWith("data:") || firstLine.startsWith("event:"));
366373
if (isSse) {
367-
tagSpanFromSseBytes(span, bytes, timeToFirstTokenNanos);
374+
tagSpanFromSseBytes(tracer, span, bytes, timeToFirstTokenNanos);
368375
} else {
369376
// Non-streaming: plain Message JSON — pass it whole, no time_to_first_token
377+
String responseJson = new String(bytes, StandardCharsets.UTF_8);
370378
InstrumentationSemConv.tagLLMSpanResponse(
379+
tracer,
371380
span,
372381
InstrumentationSemConv.PROVIDER_NAME_ANTHROPIC,
373-
new String(bytes, StandardCharsets.UTF_8),
382+
responseJson,
374383
null);
375384
}
376385
} catch (Exception e) {
@@ -406,7 +415,7 @@ private static String firstNonEmptyLine(byte[] bytes) {
406415
* assembled {@link com.anthropic.models.messages.Message} for the span.
407416
*/
408417
private static void tagSpanFromSseBytes(
409-
Span span, byte[] sseBytes, Long timeToFirstTokenNanos) {
418+
Tracer tracer, Span span, byte[] sseBytes, Long timeToFirstTokenNanos) {
410419
try {
411420
var mapper = BraintrustJsonMapper.get();
412421
var reader =
@@ -427,6 +436,7 @@ private static void tagSpanFromSseBytes(
427436
}
428437
String assembledMessageJson = BraintrustJsonMapper.toJson(accumulator.message());
429438
InstrumentationSemConv.tagLLMSpanResponse(
439+
tracer,
430440
span,
431441
InstrumentationSemConv.PROVIDER_NAME_ANTHROPIC,
432442
assembledMessageJson,

‎braintrust-sdk/instrumentation/anthropic_2_2_0/src/test/java/dev/braintrust/instrumentation/anthropic/v2_2_0/BraintrustAnthropicServerSideSpansTest.java‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
package dev.braintrust.instrumentation.anthropic.v2_2_0;
2+
3+
import static org.junit.jupiter.api.Assertions.*;
4+
5+
import com.anthropic.client.AnthropicClient;
6+
import com.anthropic.client.okhttp.AnthropicOkHttpClient;
7+
import com.anthropic.models.messages.MessageCreateParams;
8+
import com.anthropic.models.messages.Model;
9+
import com.anthropic.models.messages.WebSearchTool20250305;
10+
import com.fasterxml.jackson.databind.JsonNode;
11+
import com.fasterxml.jackson.databind.ObjectMapper;
12+
import dev.braintrust.TestHarness;
13+
import dev.braintrust.instrumentation.Instrumenter;
14+
import io.opentelemetry.api.common.AttributeKey;
15+
import io.opentelemetry.sdk.trace.data.SpanData;
16+
import java.util.List;
17+
import lombok.SneakyThrows;
18+
import net.bytebuddy.agent.ByteBuddyAgent;
19+
import org.junit.jupiter.api.BeforeAll;
20+
import org.junit.jupiter.api.BeforeEach;
21+
import org.junit.jupiter.api.Test;
22+
23+
/**
24+
* Verifies that Anthropic server-side tool calls are captured as both a cost metric on the LLM span
25+
* and a child {@code type:"tool"} span parented to it, giving each call its own cost/latency
26+
* visibility on the trace timeline. Web search ({@code server_tool_use_web_search_requests}) is the
27+
* case exercised here.
28+
*/
29+
public class BraintrustAnthropicServerSideSpansTest {
30+
private static final String TEST_MODEL = "claude-sonnet-4-5-20250929";
31+
private static final ObjectMapper JSON_MAPPER = new ObjectMapper();
32+
private static final AttributeKey<String> SPAN_ATTRIBUTES =
33+
AttributeKey.stringKey("braintrust.span_attributes");
34+
private static final AttributeKey<String> METADATA =
35+
AttributeKey.stringKey("braintrust.metadata");
36+
private static final AttributeKey<String> METRICS =
37+
AttributeKey.stringKey("braintrust.metrics");
38+
39+
@BeforeAll
40+
public static void beforeAll() {
41+
var instrumentation = ByteBuddyAgent.install();
42+
Instrumenter.install(
43+
instrumentation, BraintrustAnthropicServerSideSpansTest.class.getClassLoader());
44+
}
45+
46+
private TestHarness testHarness;
47+
48+
@BeforeEach
49+
void beforeEach() {
50+
testHarness = TestHarness.setup();
51+
}
52+
53+
private static MessageCreateParams webSearchRequest() {
54+
return MessageCreateParams.builder()
55+
.model(Model.of(TEST_MODEL))
56+
.maxTokens(1024)
57+
.addUserMessage(
58+
"Search the web for one recent AI news headline and answer in one"
59+
+ " sentence.")
60+
.addTool(WebSearchTool20250305.builder().maxUses(3).build())
61+
.build();
62+
}
63+
64+
@Test
65+
@SneakyThrows
66+
void testWebSearch() {
67+
AnthropicClient client =
68+
AnthropicOkHttpClient.builder()
69+
.baseUrl(testHarness.anthropicBaseUrl())
70+
.apiKey(testHarness.anthropicApiKey())
71+
.build();
72+
73+
var response = client.messages().create(webSearchRequest());
74+
assertNotNull(response);
75+
76+
assertWebSearch(testHarness.awaitExportedSpans(2));
77+
}
78+
79+
@Test
80+
@SneakyThrows
81+
void testWebSearchStreaming() {
82+
AnthropicClient client =
83+
AnthropicOkHttpClient.builder()
84+
.baseUrl(testHarness.anthropicBaseUrl())
85+
.apiKey(testHarness.anthropicApiKey())
86+
.build();
87+
88+
try (var stream = client.messages().createStreaming(webSearchRequest())) {
89+
stream.stream().forEach(event -> {});
90+
}
91+
92+
assertWebSearch(testHarness.awaitExportedSpans(2));
93+
}
94+
95+
@SneakyThrows
96+
private static void assertWebSearch(List<SpanData> spans) {
97+
var llmSpans = spans.stream().filter(s -> isType(s, "llm")).toList();
98+
assertEquals(1, llmSpans.size(), "expected a single LLM span");
99+
var llm = llmSpans.get(0);
100+
101+
// Cost metric on the LLM span.
102+
JsonNode metrics = JSON_MAPPER.readTree(llm.getAttributes().get(METRICS));
103+
assertTrue(
104+
metrics.has("server_tool_use_web_search_requests"),
105+
"expected server_tool_use_web_search_requests metric, got: " + metrics);
106+
assertTrue(metrics.get("server_tool_use_web_search_requests").asDouble() >= 1.0);
107+
108+
// At least one web_search tool span, parented to the LLM span.
109+
var toolSpans =
110+
spans.stream()
111+
.filter(s -> isType(s, "tool"))
112+
.filter(s -> "web_search".equals(s.getName()))
113+
.toList();
114+
assertFalse(
115+
toolSpans.isEmpty(),
116+
"expected at least one web_search tool span, got: "
117+
+ spans.stream().map(SpanData::getName).toList());
118+
119+
for (var tool : toolSpans) {
120+
assertEquals(
121+
llm.getSpanId(),
122+
tool.getParentSpanId(),
123+
"web_search tool span must be a child of the LLM span");
124+
JsonNode metadata = JSON_MAPPER.readTree(tool.getAttributes().get(METADATA));
125+
assertEquals("server_tool_use", metadata.path("tool_call_type").asText());
126+
assertEquals("web_search_tool_result", metadata.path("tool_result_type").asText());
127+
assertFalse(metadata.path("tool_use_id").asText().isEmpty());
128+
}
129+
}
130+
131+
@SneakyThrows
132+
private static boolean isType(SpanData span, String type) {
133+
String attr = span.getAttributes().get(SPAN_ATTRIBUTES);
134+
return attr != null && type.equals(JSON_MAPPER.readTree(attr).path("type").asText());
135+
}
136+
}

‎braintrust-sdk/instrumentation/aws_bedrock_2_30_0/src/main/java/dev/braintrust/instrumentation/awsbedrock/v2_30_0/BraintrustBedrockInterceptor.java‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,10 @@ public Optional<InputStream> modifyHttpResponseContent(
146146
try {
147147
String responseBodyStr = new String(bytes, StandardCharsets.UTF_8);
148148
InstrumentationSemConv.tagLLMSpanResponse(
149-
span, InstrumentationSemConv.PROVIDER_NAME_BEDROCK, responseBodyStr);
149+
tracer,
150+
span,
151+
InstrumentationSemConv.PROVIDER_NAME_BEDROCK,
152+
responseBodyStr);
150153
} catch (Exception e) {
151154
log.debug("Failed to capture response body", e);
152155
}
@@ -182,7 +185,7 @@ public Optional<Publisher<ByteBuffer>> modifyAsyncHttpResponseContent(
182185

183186
Publisher<ByteBuffer> original = publisherOpt.get();
184187
Publisher<ByteBuffer> teed =
185-
subscriber -> original.subscribe(new TeeingSubscriber(subscriber, span));
188+
subscriber -> original.subscribe(new TeeingSubscriber(subscriber, span, tracer));
186189
return Optional.of(teed);
187190
}
188191

@@ -238,6 +241,7 @@ private static String extractModelIdFromPath(String path) {
238241
private static class TeeingSubscriber implements Subscriber<ByteBuffer> {
239242
private final Subscriber<? super ByteBuffer> downstream;
240243
private final Span span;
244+
private final Tracer tracer;
241245
private final MessageDecoder decoder = new MessageDecoder();
242246

243247
// Accumulated incrementally in onNext — no message list retained.
@@ -248,9 +252,10 @@ private static class TeeingSubscriber implements Subscriber<ByteBuffer> {
248252
private long startNanos;
249253
private Long timeToFirstTokenNanos = null;
250254

251-
TeeingSubscriber(Subscriber<? super ByteBuffer> downstream, Span span) {
255+
TeeingSubscriber(Subscriber<? super ByteBuffer> downstream, Span span, Tracer tracer) {
252256
this.downstream = downstream;
253257
this.span = span;
258+
this.tracer = tracer;
254259
}
255260

256261
@Override
@@ -304,6 +309,7 @@ public void onError(Throwable t) {
304309
public void onComplete() {
305310
try {
306311
InstrumentationSemConv.tagLLMSpanResponse(
312+
tracer,
307313
span,
308314
InstrumentationSemConv.PROVIDER_NAME_BEDROCK,
309315
buildConverseJson(text.toString(), stopReason, inputTokens, outputTokens),

‎braintrust-sdk/instrumentation/genai_1_18_0/src/main/java/com/google/genai/BraintrustApiClient.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ private void tagSpan(
5555
@Nullable String responseBody) {
5656
try {
5757
Map<String, Object> metadata = new java.util.HashMap<>();
58-
metadata.put("provider", "gemini");
58+
metadata.put("provider", "google");
5959

6060
// Parse request
6161
if (requestBody != null) {

‎braintrust-sdk/instrumentation/genai_1_18_0/src/test/java/dev/braintrust/instrumentation/genai/v1_18_0/BraintrustGenAITest.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ void testWrapGemini() {
6767
span.getAttributes().get(AttributeKey.stringKey("braintrust.metadata"));
6868
assertNotNull(metadataJson, "braintrust.metadata should be set");
6969
var metadata = JSON_MAPPER.readTree(metadataJson);
70-
assertEquals("gemini", metadata.get("provider").asText());
70+
assertEquals("google", metadata.get("provider").asText());
7171
assertEquals(MODEL_ID, metadata.get("model").asText());
7272
assertEquals(0.0, metadata.get("temperature").asDouble());
7373
assertEquals(50, metadata.get("maxOutputTokens").asInt());
@@ -145,7 +145,7 @@ void testWrapGeminiAsync() {
145145
span.getAttributes().get(AttributeKey.stringKey("braintrust.metadata"));
146146
assertNotNull(metadataJson, "braintrust.metadata should be set");
147147
var metadata = JSON_MAPPER.readTree(metadataJson);
148-
assertEquals("gemini", metadata.get("provider").asText());
148+
assertEquals("google", metadata.get("provider").asText());
149149
assertEquals(MODEL_ID, metadata.get("model").asText());
150150
assertEquals(0.0, metadata.get("temperature").asDouble());
151151
assertEquals(50, metadata.get("maxOutputTokens").asInt());

‎braintrust-sdk/instrumentation/langchain_1_8_0/src/main/java/dev/braintrust/instrumentation/langchain/v1_8_0/WrappedHttpClient.java‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ public SuccessfulHttpResponse execute(HttpRequest request)
4949
tagRequest(span, request);
5050
var response = underlying.execute(request);
5151
InstrumentationSemConv.tagLLMSpanResponse(
52-
span, options.providerName(), response.body());
52+
tracer, span, options.providerName(), response.body());
5353
return response;
5454
} catch (Throwable t) {
5555
InstrumentationSemConv.tagLLMSpanResponse(span, t);
@@ -73,7 +73,8 @@ public void execute(HttpRequest request, ServerSentEventListener listener) {
7373
tagRequest(span, request);
7474
underlying.execute(
7575
request,
76-
new WrappedServerSentEventListener(listener, span, options.providerName()));
76+
new WrappedServerSentEventListener(
77+
listener, span, options.providerName(), tracer));
7778
} catch (Throwable t) {
7879
InstrumentationSemConv.tagLLMSpanResponse(span, t);
7980
span.end();
@@ -97,7 +98,8 @@ public void execute(
9798
underlying.execute(
9899
request,
99100
parser,
100-
new WrappedServerSentEventListener(listener, span, options.providerName()));
101+
new WrappedServerSentEventListener(
102+
listener, span, options.providerName(), tracer));
101103
} catch (Throwable t) {
102104
InstrumentationSemConv.tagLLMSpanResponse(span, t);
103105
span.end();
@@ -122,16 +124,18 @@ static class WrappedServerSentEventListener implements ServerSentEventListener {
122124
private final ServerSentEventListener delegate;
123125
private final Span span;
124126
private final String providerName;
127+
private final Tracer tracer;
125128
private final long startNanos = System.nanoTime();
126129
private final AtomicLong timeToFirstTokenNanos = new AtomicLong();
127130
private final SseResponseAccumulator accumulator =
128131
new SseResponseAccumulator(BraintrustJsonMapper.get());
129132

130133
WrappedServerSentEventListener(
131-
ServerSentEventListener delegate, Span span, String providerName) {
134+
ServerSentEventListener delegate, Span span, String providerName, Tracer tracer) {
132135
this.delegate = delegate;
133136
this.span = span;
134137
this.providerName = providerName;
138+
this.tracer = tracer;
135139
}
136140

137141
@Override
@@ -188,8 +192,9 @@ private void accumulateChunk(String data) {
188192
private void finalizeSpan() {
189193
try {
190194
Long ttft = timeToFirstTokenNanos.get();
195+
String responseBody = accumulator.build();
191196
InstrumentationSemConv.tagLLMSpanResponse(
192-
span, providerName, accumulator.build(), ttft);
197+
tracer, span, providerName, responseBody, ttft);
193198
} catch (Exception e) {
194199
log.debug("Failed to finalize streaming span", e);
195200
}

0 commit comments

Comments
 (0)