Add support for ConverseStream to bedrock instrumentation

Author: anuraaga

instrumentation/aws-sdk/aws-sdk-2.2/javaagent/build.gradle.kts

@@ -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")
 

instrumentation/aws-sdk/aws-sdk-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awssdk/v2_2/BedrockRuntimeInstrumentationModule.java

@@ -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;
 
@@ -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(

instrumentation/aws-sdk/aws-sdk-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awssdk/v2_2/DefaultBedrockRuntimeAsyncClientBuilderInstrumentation.java

@@ -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);
+    }
+  }
+}

instrumentation/aws-sdk/aws-sdk-2.2/javaagent/src/testBedrockRuntime/java/io/opentelemetry/instrumentation/awssdk/v2_2/internal/Aws2BedrockRuntimeTest.java

@@ -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
@@ -24,4 +25,10 @@ protected InstrumentationExtension getTesting() {
   protected ClientOverrideConfiguration.Builder createOverrideConfigurationBuilder() {
     return ClientOverrideConfiguration.builder();
   }
+
+  @Override
+  protected BedrockRuntimeAsyncClient configureBedrockRuntimeClient(
+      BedrockRuntimeAsyncClient client) {
+    return client;
+  }
 }

instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/AwsSdkTelemetry.java

@@ -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;
@@ -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;
 
@@ -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()
@@ -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(
+      BedrockRuntimeAsyncClient bedrockClient) {
+    return BedrockRuntimeImpl.wrap(bedrockClient);
+  }
 }

instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/internal/BedrockRuntimeAccess.java

@@ -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() {}
@@ -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;
   }
 }

instrumentation/aws-sdk/aws-sdk-2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/internal/BedrockRuntimeAttributesGetter.java

@@ -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;
@@ -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;
@@ -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;
     }
-    return Arrays.asList(stopReason);
+    return stopReasons;
   }
 
   @Nullable
@@ -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);
   }
 }