Skip to content

Commit b16f3fe

Browse files
realarkrealark
committed
support web search in openai-responses + anthropic
1 parent 16d0d97 commit b16f3fe

13 files changed

Lines changed: 1404 additions & 52 deletions

File tree

‎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: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
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+
private static final AttributeKey<String> OUTPUT_JSON =
39+
AttributeKey.stringKey("braintrust.output_json");
40+
41+
@BeforeAll
42+
public static void beforeAll() {
43+
var instrumentation = ByteBuddyAgent.install();
44+
Instrumenter.install(
45+
instrumentation, BraintrustAnthropicServerSideSpansTest.class.getClassLoader());
46+
}
47+
48+
private TestHarness testHarness;
49+
50+
@BeforeEach
51+
void beforeEach() {
52+
testHarness = TestHarness.setup();
53+
}
54+
55+
private static MessageCreateParams webSearchRequest() {
56+
return MessageCreateParams.builder()
57+
.model(Model.of(TEST_MODEL))
58+
.maxTokens(1024)
59+
.addUserMessage(
60+
"Search the web for one recent AI news headline and answer in one"
61+
+ " sentence.")
62+
.addTool(WebSearchTool20250305.builder().maxUses(3).build())
63+
.build();
64+
}
65+
66+
@Test
67+
@SneakyThrows
68+
void testWebSearch() {
69+
AnthropicClient client =
70+
AnthropicOkHttpClient.builder()
71+
.baseUrl(testHarness.anthropicBaseUrl())
72+
.apiKey(testHarness.anthropicApiKey())
73+
.build();
74+
75+
var response = client.messages().create(webSearchRequest());
76+
assertNotNull(response);
77+
78+
assertWebSearch(testHarness.awaitExportedSpans(2));
79+
}
80+
81+
@Test
82+
@SneakyThrows
83+
void testWebSearchStreaming() {
84+
AnthropicClient client =
85+
AnthropicOkHttpClient.builder()
86+
.baseUrl(testHarness.anthropicBaseUrl())
87+
.apiKey(testHarness.anthropicApiKey())
88+
.build();
89+
90+
try (var stream = client.messages().createStreaming(webSearchRequest())) {
91+
stream.stream().forEach(event -> {});
92+
}
93+
94+
assertWebSearch(testHarness.awaitExportedSpans(2));
95+
}
96+
97+
@SneakyThrows
98+
private static void assertWebSearch(List<SpanData> spans) {
99+
var llmSpans = spans.stream().filter(s -> isType(s, "llm")).toList();
100+
assertEquals(1, llmSpans.size(), "expected a single LLM span");
101+
var llm = llmSpans.get(0);
102+
103+
// Cost metric on the LLM span.
104+
JsonNode metrics = JSON_MAPPER.readTree(llm.getAttributes().get(METRICS));
105+
assertTrue(
106+
metrics.has("server_tool_use_web_search_requests"),
107+
"expected server_tool_use_web_search_requests metric, got: " + metrics);
108+
assertTrue(metrics.get("server_tool_use_web_search_requests").asDouble() >= 1.0);
109+
110+
// The opaque per-result blob (multi-KB of base64, ~80% of a real web-search response) is
111+
// stripped from the LLM span output, not just from the tool span. Covers both the raw-bytes
112+
// path and the streaming path, where the body is a reassembled Message.
113+
JsonNode llmOutput = JSON_MAPPER.readTree(llm.getAttributes().get(OUTPUT_JSON));
114+
JsonNode searchResults =
115+
llmOutput.path("content").findValue("content").path(0); // first web_search_result
116+
assertEquals(
117+
"<redacted>",
118+
searchResults.path("encrypted_content").asText(),
119+
"expected encrypted_content redacted in LLM span output, got: " + searchResults);
120+
assertFalse(
121+
searchResults.path("title").asText().isEmpty(),
122+
"redaction must preserve the rest of the result: " + searchResults);
123+
124+
// At least one web_search tool span, parented to the LLM span.
125+
var toolSpans =
126+
spans.stream()
127+
.filter(s -> isType(s, "tool"))
128+
.filter(s -> "web_search".equals(s.getName()))
129+
.toList();
130+
assertFalse(
131+
toolSpans.isEmpty(),
132+
"expected at least one web_search tool span, got: "
133+
+ spans.stream().map(SpanData::getName).toList());
134+
135+
for (var tool : toolSpans) {
136+
assertEquals(
137+
llm.getSpanId(),
138+
tool.getParentSpanId(),
139+
"web_search tool span must be a child of the LLM span");
140+
JsonNode metadata = JSON_MAPPER.readTree(tool.getAttributes().get(METADATA));
141+
assertEquals("server_tool_use", metadata.path("tool_call_type").asText());
142+
assertEquals("web_search_tool_result", metadata.path("tool_result_type").asText());
143+
assertFalse(metadata.path("tool_use_id").asText().isEmpty());
144+
}
145+
}
146+
147+
@SneakyThrows
148+
private static boolean isType(SpanData span, String type) {
149+
String attr = span.getAttributes().get(SPAN_ATTRIBUTES);
150+
return attr != null && type.equals(JSON_MAPPER.readTree(attr).path("type").asText());
151+
}
152+
}

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