Skip to content

Commit bfbc715

Browse files
realarkrealark
committed
basic bedrock instrumentation
1 parent 2b6e5e3 commit bfbc715

9 files changed

Lines changed: 610 additions & 5 deletions

File tree

‎braintrust-sdk/instrumentation/aws_bedrock_2_30_0/build.gradle‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
def awsBedrockVersion = '2.30.0'
2+
3+
muzzle {
4+
pass {
5+
group = 'software.amazon.awssdk'
6+
module = 'bedrockruntime'
7+
versions = "[${awsBedrockVersion},)"
8+
}
9+
}
10+
11+
dependencies {
12+
compileOnly project(':braintrust-java-agent:instrumenter')
13+
implementation "io.opentelemetry:opentelemetry-api:${otelVersion}"
14+
implementation 'com.google.code.findbugs:jsr305:3.0.2'
15+
implementation "org.slf4j:slf4j-api:${slf4jVersion}"
16+
implementation project(':braintrust-sdk')
17+
18+
// ByteBuddy for ElementMatcher types used in instrumentation definitions
19+
compileOnly 'net.bytebuddy:byte-buddy:1.17.5'
20+
21+
// Target library — compileOnly because it will be on the app classpath at runtime
22+
compileOnly "software.amazon.awssdk:bedrockruntime:${awsBedrockVersion}"
23+
24+
// Test dependencies
25+
testImplementation(testFixtures(project(":test-harness")))
26+
testImplementation project(':braintrust-java-agent:instrumenter')
27+
testImplementation "org.junit.jupiter:junit-jupiter:${junitVersion}"
28+
testRuntimeOnly 'org.junit.platform:junit-platform-launcher'
29+
testImplementation 'net.bytebuddy:byte-buddy-agent:1.17.5'
30+
testRuntimeOnly "org.slf4j:slf4j-simple:${slf4jVersion}"
31+
testImplementation "software.amazon.awssdk:bedrockruntime:${awsBedrockVersion}"
32+
testImplementation 'com.fasterxml.jackson.core:jackson-databind:2.17.0'
33+
}
34+
35+
test {
36+
useJUnitPlatform()
37+
workingDir = rootProject.projectDir
38+
testLogging {
39+
events "passed", "skipped", "failed"
40+
showStandardStreams = true
41+
exceptionFormat "full"
42+
}
43+
}

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

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
package dev.braintrust.instrumentation.awsbedrock.v2_30_0;
2+
3+
import io.opentelemetry.api.OpenTelemetry;
4+
import lombok.extern.slf4j.Slf4j;
5+
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
6+
import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeClientBuilder;
7+
8+
/** Braintrust instrumentation for the AWS Bedrock Runtime client. */
9+
@Slf4j
10+
public class BraintrustAWSBedrock {
11+
12+
/**
13+
* Wraps a {@link BedrockRuntimeClientBuilder} so that every {@code converse} and {@code
14+
* converseStream} call made through the resulting client is traced via OpenTelemetry.
15+
*
16+
* <p>Call this method after applying all custom builder settings
17+
*
18+
* @param openTelemetry the OpenTelemetry instance to use for tracing
19+
* @param builder the client builder to instrument
20+
* @return the same builder (for fluent chaining)
21+
*/
22+
public static BedrockRuntimeClientBuilder wrap(
23+
OpenTelemetry openTelemetry, BedrockRuntimeClientBuilder builder) {
24+
try {
25+
// Read existing config so we don't clobber user-registered interceptors
26+
ClientOverrideConfiguration existing = builder.overrideConfiguration();
27+
ClientOverrideConfiguration.Builder configBuilder = existing.toBuilder();
28+
for (var interceptor : configBuilder.executionInterceptors()) {
29+
if (interceptor instanceof BraintrustBedrockInterceptor) {
30+
log.info("builder already wrapped. Skipping");
31+
return builder;
32+
}
33+
}
34+
configBuilder.addExecutionInterceptor(new BraintrustBedrockInterceptor(openTelemetry));
35+
builder.overrideConfiguration(configBuilder.build());
36+
} catch (Exception e) {
37+
log.warn("Failed to apply Bedrock instrumentation", e);
38+
}
39+
return builder;
40+
}
41+
}

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

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
package dev.braintrust.instrumentation.awsbedrock.v2_30_0;
2+
3+
import dev.braintrust.instrumentation.InstrumentationSemConv;
4+
import io.opentelemetry.api.OpenTelemetry;
5+
import io.opentelemetry.api.trace.Span;
6+
import io.opentelemetry.api.trace.Tracer;
7+
import io.opentelemetry.context.Context;
8+
import java.io.ByteArrayInputStream;
9+
import java.io.InputStream;
10+
import java.nio.charset.StandardCharsets;
11+
import java.util.Arrays;
12+
import java.util.List;
13+
import java.util.Optional;
14+
import lombok.extern.slf4j.Slf4j;
15+
import software.amazon.awssdk.core.SdkRequest;
16+
import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
17+
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
18+
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
19+
import software.amazon.awssdk.core.interceptor.SdkExecutionAttribute;
20+
import software.amazon.awssdk.http.SdkHttpRequest;
21+
import software.amazon.awssdk.services.bedrockruntime.model.ConverseRequest;
22+
23+
/**
24+
* AWS SDK ExecutionInterceptor that creates OpenTelemetry spans for Bedrock Converse calls,
25+
* capturing the raw request and response bodies via {@link InstrumentationSemConv}.
26+
*/
27+
@Slf4j
28+
class BraintrustBedrockInterceptor implements ExecutionInterceptor {
29+
private static final String INSTRUMENTATION_NAME = "braintrust-aws-bedrock";
30+
31+
private static final ExecutionAttribute<Span> SPAN_ATTRIBUTE =
32+
new ExecutionAttribute<>("braintrust.span");
33+
private static final ExecutionAttribute<String> MODEL_ID_ATTRIBUTE =
34+
new ExecutionAttribute<>("braintrust.modelId");
35+
36+
private final Tracer tracer;
37+
38+
BraintrustBedrockInterceptor(OpenTelemetry openTelemetry) {
39+
this.tracer = openTelemetry.getTracer(INSTRUMENTATION_NAME);
40+
}
41+
42+
@Override
43+
public void beforeExecution(
44+
software.amazon.awssdk.core.interceptor.Context.BeforeExecution context,
45+
ExecutionAttributes executionAttributes) {
46+
String operationName =
47+
executionAttributes.getAttribute(SdkExecutionAttribute.OPERATION_NAME);
48+
if (operationName == null) {
49+
operationName = "BedrockRuntime";
50+
}
51+
52+
SdkRequest sdkRequest = context.request();
53+
String modelId = extractModelId(sdkRequest);
54+
55+
Span span = tracer.spanBuilder(operationName).setParent(Context.current()).startSpan();
56+
executionAttributes.putAttribute(SPAN_ATTRIBUTE, span);
57+
if (modelId != null) {
58+
executionAttributes.putAttribute(MODEL_ID_ATTRIBUTE, modelId);
59+
}
60+
}
61+
62+
@Override
63+
public SdkHttpRequest modifyHttpRequest(
64+
software.amazon.awssdk.core.interceptor.Context.ModifyHttpRequest context,
65+
ExecutionAttributes executionAttributes) {
66+
Span span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
67+
if (span == null) {
68+
return context.httpRequest();
69+
}
70+
71+
SdkHttpRequest httpRequest = context.httpRequest();
72+
String modelId = executionAttributes.getAttribute(MODEL_ID_ATTRIBUTE);
73+
74+
// Extract model ID from URL path if we didn't get it from the SDK request
75+
if (modelId == null) {
76+
modelId = extractModelIdFromPath(httpRequest.encodedPath());
77+
if (modelId != null) {
78+
executionAttributes.putAttribute(MODEL_ID_ATTRIBUTE, modelId);
79+
}
80+
}
81+
82+
// Capture the raw request body
83+
String requestBody = null;
84+
if (context.requestBody().isPresent()) {
85+
try (InputStream is = context.requestBody().get().contentStreamProvider().newStream()) {
86+
requestBody = new String(is.readAllBytes(), StandardCharsets.UTF_8);
87+
} catch (Exception e) {
88+
log.debug("Failed to capture request body", e);
89+
}
90+
}
91+
92+
String baseUrl = httpRequest.protocol() + "://" + httpRequest.host();
93+
List<String> pathSegments =
94+
Arrays.stream(httpRequest.encodedPath().split("/"))
95+
.filter(s -> !s.isEmpty())
96+
.toList();
97+
98+
InstrumentationSemConv.tagLLMSpanRequest(
99+
span,
100+
InstrumentationSemConv.PROVIDER_NAME_BEDROCK,
101+
baseUrl,
102+
pathSegments,
103+
"POST",
104+
requestBody,
105+
modelId);
106+
107+
return httpRequest;
108+
}
109+
110+
@Override
111+
public Optional<InputStream> modifyHttpResponseContent(
112+
software.amazon.awssdk.core.interceptor.Context.ModifyHttpResponse context,
113+
ExecutionAttributes executionAttributes) {
114+
Span span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
115+
if (span == null) {
116+
return context.responseBody();
117+
}
118+
119+
Optional<InputStream> body = context.responseBody();
120+
if (body.isPresent()) {
121+
try {
122+
byte[] bytes = body.get().readAllBytes();
123+
String responseBodyStr = new String(bytes, StandardCharsets.UTF_8);
124+
InstrumentationSemConv.tagLLMSpanResponse(
125+
span, InstrumentationSemConv.PROVIDER_NAME_BEDROCK, responseBodyStr);
126+
return Optional.of(new ByteArrayInputStream(bytes));
127+
} catch (Exception e) {
128+
log.debug("Failed to capture response body", e);
129+
}
130+
}
131+
return body;
132+
}
133+
134+
@Override
135+
public void afterExecution(
136+
software.amazon.awssdk.core.interceptor.Context.AfterExecution context,
137+
ExecutionAttributes executionAttributes) {
138+
endSpan(executionAttributes, null);
139+
}
140+
141+
@Override
142+
public void onExecutionFailure(
143+
software.amazon.awssdk.core.interceptor.Context.FailedExecution context,
144+
ExecutionAttributes executionAttributes) {
145+
endSpan(executionAttributes, context.exception());
146+
}
147+
148+
private static void endSpan(
149+
ExecutionAttributes executionAttributes, @javax.annotation.Nullable Throwable error) {
150+
Span span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
151+
if (span == null) {
152+
return;
153+
}
154+
if (error != null) {
155+
InstrumentationSemConv.tagLLMSpanResponse(span, error);
156+
}
157+
span.end();
158+
}
159+
160+
/** Extract the model ID from the SDK request. */
161+
private static String extractModelId(SdkRequest request) {
162+
if (request instanceof ConverseRequest converseRequest) {
163+
return converseRequest.modelId();
164+
}
165+
return null;
166+
}
167+
168+
/** Extract model ID from URL path like /model/{modelId}/converse */
169+
private static String extractModelIdFromPath(String path) {
170+
if (path != null && path.startsWith("/model/")) {
171+
int start = "/model/".length();
172+
int end = path.indexOf("/", start);
173+
if (end > start) {
174+
return java.net.URLDecoder.decode(
175+
path.substring(start, end), StandardCharsets.UTF_8);
176+
}
177+
}
178+
return null;
179+
}
180+
}

0 commit comments

Comments
 (0)