Skip to content

Add support for ConverseStream to bedrock instrumentation #13410

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Mar 12, 2025
Merged

Add support for ConverseStream to bedrock instrumentation #13410

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
b02d108
Add support for ConverseStream to bedrock instrumentation
anuraaga Feb 27, 2025
24e9b30
Cleanup
anuraaga Feb 28, 2025
be043a3
Empty
anuraaga Mar 5, 2025
cf7b818
Remove Nullable
anuraaga Mar 6, 2025
cbd0649
Check empty
anuraaga Mar 6, 2025
6c57d34
not a pon
anuraaga Mar 6, 2025
73fdb57
no-nullable
anuraaga Mar 10, 2025
5648989
Merge branch 'main' of github.com:open-telemetry/opentelemetry-java-i…
anuraaga Mar 11, 2025
e849d2f
Merge branch 'main' of github.com:open-telemetry/opentelemetry-java-i…
anuraaga Mar 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions instrumentation/aws-sdk/aws-sdk-2.2/javaagent/build.gradle.kts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,11 @@ dependencies {
library("software.amazon.awssdk:aws-core:2.2.0")
library("software.amazon.awssdk:sqs:2.2.0")

// Don't use library to make sure base test is run with the floor version.
// bedrock runtime is tested separately in testBedrockRuntime.
// First release with Converse API
compileOnly("software.amazon.awssdk:bedrockruntime:2.25.63")

testImplementation(project(":instrumentation:aws-sdk:aws-sdk-2.2:testing"))
testImplementation("io.opentelemetry.contrib:opentelemetry-aws-xray-propagator")

Expand Down
10 changes: 10 additions & 0 deletions ...ntelemetry/javaagent/instrumentation/awssdk/v2_2/BedrockRuntimeInstrumentationModule.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,10 @@
import com.google.auto.service.AutoService;
import io.opentelemetry.instrumentation.awssdk.v2_2.internal.BedrockRuntimeAdviceBridge;
import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import java.util.ArrayList;
import java.util.List;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.matcher.ElementMatcher;

Expand All @@ -27,6 +30,13 @@ public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
return hasClassesNamed("software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeClient");
}

@Override
public List<TypeInstrumentation> typeInstrumentations() {
List<TypeInstrumentation> instrumentations = new ArrayList<>(super.typeInstrumentations());
instrumentations.add(new DefaultBedrockRuntimeAsyncClientBuilderInstrumentation());
return instrumentations;
}

@Override
public void doTransform(TypeTransformer transformer) {
transformer.applyAdviceToMethod(
Expand Down
41 changes: 41 additions & 0 deletions ...t/instrumentation/awssdk/v2_2/DefaultBedrockRuntimeAsyncClientBuilderInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
/*
* Copyright The OpenTelemetry Authors
* SPDX-License-Identifier: Apache-2.0
*/

package io.opentelemetry.javaagent.instrumentation.awssdk.v2_2;

import static net.bytebuddy.matcher.ElementMatchers.named;

import io.opentelemetry.instrumentation.awssdk.v2_2.autoconfigure.AwsSdkSingletons;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;
import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeAsyncClient;

public class DefaultBedrockRuntimeAsyncClientBuilderInstrumentation implements TypeInstrumentation {

@Override
public ElementMatcher<TypeDescription> typeMatcher() {
return named(
"software.amazon.awssdk.services.bedrockruntime.DefaultBedrockRuntimeAsyncClientBuilder");
}

@Override
public void transform(TypeTransformer transformer) {
transformer.applyAdviceToMethod(
named("buildClient"), this.getClass().getName() + "$BuildClientAdvice");
}

@SuppressWarnings("unused")
public static class BuildClientAdvice {

@Advice.OnMethodExit(suppress = Throwable.class)
public static void methodExit(
@Advice.Return(readOnly = false) BedrockRuntimeAsyncClient client) {
client = AwsSdkSingletons.telemetry().wrapBedrockRuntimeClient(client);
}
}
}
7 changes: 7 additions & 0 deletions ...me/java/io/opentelemetry/instrumentation/awssdk/v2_2/internal/Aws2BedrockRuntimeTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
import org.junit.jupiter.api.extension.RegisterExtension;
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeAsyncClient;

class Aws2BedrockRuntimeTest extends AbstractAws2BedrockRuntimeTest {
@RegisterExtension
Expand All @@ -24,4 +25,10 @@ protected InstrumentationExtension getTesting() {
protected ClientOverrideConfiguration.Builder createOverrideConfigurationBuilder() {
return ClientOverrideConfiguration.builder();
}

@Override
protected BedrockRuntimeAsyncClient configureBedrockRuntimeClient(
BedrockRuntimeAsyncClient client) {
return client;
}
}
20 changes: 20 additions & 0 deletions ...2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/AwsSdkTelemetry.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import io.opentelemetry.context.propagation.TextMapPropagator;
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
import io.opentelemetry.instrumentation.awssdk.v2_2.internal.AwsSdkInstrumenterFactory;
import io.opentelemetry.instrumentation.awssdk.v2_2.internal.BedrockRuntimeImpl;
import io.opentelemetry.instrumentation.awssdk.v2_2.internal.Response;
import io.opentelemetry.instrumentation.awssdk.v2_2.internal.SqsImpl;
import io.opentelemetry.instrumentation.awssdk.v2_2.internal.SqsProcessRequest;
Expand All @@ -20,6 +21,7 @@
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeAsyncClient;
import software.amazon.awssdk.services.sqs.SqsAsyncClient;
import software.amazon.awssdk.services.sqs.SqsClient;

Expand All @@ -28,6 +30,14 @@
* ExecutionInterceptor} returned by {@link #newExecutionInterceptor()} with an SDK client to have
* all requests traced.
*
* <p>Certain services additionally require wrapping the SDK client itself:
*
* <ul>
* <li>SQSClient - {@link #wrap(SqsClient)}
* <li>SQSAsyncClient - {@link #wrap(SqsAsyncClient)}
* <li>BedrockRuntimeAsyncClient - {@link #wrapBedrockRuntimeClient(BedrockRuntimeAsyncClient)}
* </ul>
*
* <pre>{@code
* DynamoDbClient dynamoDb = DynamoDbClient.builder()
* .overrideConfiguration(ClientOverrideConfiguration.builder()
Expand Down Expand Up @@ -126,4 +136,14 @@ public SqsClient wrap(SqsClient sqsClient) {
public SqsAsyncClient wrap(SqsAsyncClient sqsClient) {
return SqsImpl.wrap(sqsClient);
}

/**
* Construct a new tracing-enabled {@link BedrockRuntimeAsyncClient} using the provided {@link
* BedrockRuntimeAsyncClient} instance.
*/
@NoMuzzle
public BedrockRuntimeAsyncClient wrapBedrockRuntimeClient(
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I couldn't use wrap as was probably intended because all argument classes for overloads have to be on the compile classpath, which isn't the case. If the library instrumentation hasn't been marked stable yet, in a separate PR I could deprecate wrap and add wrapSqsClient if it makes sense

BedrockRuntimeAsyncClient bedrockClient) {
return BedrockRuntimeImpl.wrap(bedrockClient);
}
}
9 changes: 4 additions & 5 deletions ...main/java/io/opentelemetry/instrumentation/awssdk/v2_2/internal/BedrockRuntimeAccess.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
import java.util.List;
import javax.annotation.Nullable;
import software.amazon.awssdk.core.SdkRequest;
import software.amazon.awssdk.core.SdkResponse;

final class BedrockRuntimeAccess {
private BedrockRuntimeAccess() {}
Expand Down Expand Up @@ -69,19 +68,19 @@ static List<String> getStopSequences(SdkRequest request) {

@Nullable
@NoMuzzle
static String getStopReason(SdkResponse response) {
return enabled ? BedrockRuntimeImpl.getStopReason(response) : null;
static List<String> getStopReasons(Response response) {
return enabled ? BedrockRuntimeImpl.getStopReasons(response) : null;
}

@Nullable
@NoMuzzle
static Long getUsageInputTokens(SdkResponse response) {
static Long getUsageInputTokens(Response response) {
return enabled ? BedrockRuntimeImpl.getUsageInputTokens(response) : null;
}

@Nullable
@NoMuzzle
static Long getUsageOutputTokens(SdkResponse response) {
static Long getUsageOutputTokens(Response response) {
return enabled ? BedrockRuntimeImpl.getUsageOutputTokens(response) : null;
}
}
13 changes: 7 additions & 6 deletions ...io/opentelemetry/instrumentation/awssdk/v2_2/internal/BedrockRuntimeAttributesGetter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
import static io.opentelemetry.instrumentation.awssdk.v2_2.internal.TracingExecutionInterceptor.SDK_REQUEST_ATTRIBUTE;

import io.opentelemetry.instrumentation.api.incubator.semconv.genai.GenAiAttributesGetter;
import java.util.Arrays;
import java.util.List;
import javax.annotation.Nullable;
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
Expand Down Expand Up @@ -37,6 +36,8 @@ public String getOperationName(ExecutionAttributes executionAttributes) {
if (operation != null) {
switch (operation) {
case "Converse":
// fallthrough
case "ConverseStream":
return GenAiOperationNameIncubatingValues.CHAT;
default:
return null;
Expand Down Expand Up @@ -117,11 +118,11 @@ public Double getRequestTopP(ExecutionAttributes executionAttributes) {
@Override
public List<String> getResponseFinishReasons(
ExecutionAttributes executionAttributes, Response response) {
String stopReason = BedrockRuntimeAccess.getStopReason(response.getSdkResponse());
if (stopReason == null) {
List<String> stopReasons = BedrockRuntimeAccess.getStopReasons(response);
if (stopReasons == null) {
return null;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think that we currently return empty list instead of null from other getters

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks - I think we can't do that for the request stop sequences to differentiate from user setting to empty vs not setting, but for response here empty seems good.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you also remove the @Nullable. When returning empty instead of null you may also need to change the extractor as internalSet only skips null attributes, but not the empty ones. Idk if that matters.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, removed @Nullable and had the extractor check for empty. I left the Getter permitting null anyways since I couldn't see a strong reason to be too strict on it (I noticed MessagingAttributesGetter.getMesssageHeader being strict though)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll leave this to @trask to decide whether it makes sense to keep the @Nullable on the getter.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm ok with @nullable List as long as there's a comment in the code explaining what the difference between null and empty signals to the caller since it's never obvious to me

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks - actually no comment comes to mind so I removed the Nullable but kept the null check to avoid crashing if instrumentation has a bug, realized that's probably closer to the practice of what's in the OTel APIs.

}
return Arrays.asList(stopReason);
return stopReasons;
}

@Nullable
Expand All @@ -139,12 +140,12 @@ public String getResponseModel(ExecutionAttributes executionAttributes, Response
@Nullable
@Override
public Long getUsageInputTokens(ExecutionAttributes executionAttributes, Response response) {
return BedrockRuntimeAccess.getUsageInputTokens(response.getSdkResponse());
return BedrockRuntimeAccess.getUsageInputTokens(response);
}

@Nullable
@Override
public Long getUsageOutputTokens(ExecutionAttributes executionAttributes, Response response) {
return BedrockRuntimeAccess.getUsageOutputTokens(response.getSdkResponse());
return BedrockRuntimeAccess.getUsageOutputTokens(response);
}
}
Loading
Loading