Add initial gen_ai instrumentation of bedrock InvokeModel (#13547)

Author: anuraaga

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

@@ -12,6 +12,7 @@
 import javax.annotation.Nullable;
 import software.amazon.awssdk.core.SdkRequest;
 import software.amazon.awssdk.core.SdkResponse;
+import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
 
 final class BedrockRuntimeAccess {
   private BedrockRuntimeAccess() {}
@@ -44,72 +45,99 @@ static boolean isBedrockRuntimeResponse(SdkResponse response) {
     return enabled && BedrockRuntimeImpl.isBedrockRuntimeResponse(response);
   }
 
+  @NoMuzzle
+  static void maybeParseInvokeModelRequest(
+      ExecutionAttributes executionAttributes, SdkRequest request) {
+    if (enabled) {
+      BedrockRuntimeImpl.maybeParseInvokeModelRequest(executionAttributes, request);
+    }
+  }
+
+  @NoMuzzle
+  static void maybeParseInvokeModelResponse(
+      ExecutionAttributes executionAttributes, SdkResponse response) {
+    if (enabled) {
+      BedrockRuntimeImpl.maybeParseInvokeModelResponse(executionAttributes, response);
+    }
+  }
+
+  @Nullable
+  @NoMuzzle
+  static String getModelId(ExecutionAttributes executionAttributes) {
+    return enabled ? BedrockRuntimeImpl.getModelId(executionAttributes) : null;
+  }
+
   @Nullable
   @NoMuzzle
-  static String getModelId(SdkRequest request) {
-    return enabled ? BedrockRuntimeImpl.getModelId(request) : null;
+  static String getOperationName(ExecutionAttributes executionAttributes) {
+    return enabled ? BedrockRuntimeImpl.getOperationName(executionAttributes) : null;
   }
 
   @Nullable
   @NoMuzzle
-  static Long getMaxTokens(SdkRequest request) {
-    return enabled ? BedrockRuntimeImpl.getMaxTokens(request) : null;
+  static Long getMaxTokens(ExecutionAttributes executionAttributes) {
+    return enabled ? BedrockRuntimeImpl.getMaxTokens(executionAttributes) : null;
   }
 
   @Nullable
   @NoMuzzle
-  static Double getTemperature(SdkRequest request) {
-    return enabled ? BedrockRuntimeImpl.getTemperature(request) : null;
+  static Double getTemperature(ExecutionAttributes executionAttributes) {
+    return enabled ? BedrockRuntimeImpl.getTemperature(executionAttributes) : null;
   }
 
   @Nullable
   @NoMuzzle
-  static Double getTopP(SdkRequest request) {
-    return enabled ? BedrockRuntimeImpl.getTopP(request) : null;
+  static Double getTopP(ExecutionAttributes executionAttributes) {
+    return enabled ? BedrockRuntimeImpl.getTopP(executionAttributes) : null;
   }
 
   @Nullable
   @NoMuzzle
-  static List<String> getStopSequences(SdkRequest request) {
-    return enabled ? BedrockRuntimeImpl.getStopSequences(request) : null;
+  static List<String> getStopSequences(ExecutionAttributes executionAttributes) {
+    return enabled ? BedrockRuntimeImpl.getStopSequences(executionAttributes) : null;
   }
 
   @Nullable
   @NoMuzzle
-  static List<String> getStopReasons(Response response) {
-    return enabled ? BedrockRuntimeImpl.getStopReasons(response) : null;
+  static List<String> getStopReasons(ExecutionAttributes executionAttributes, Response response) {
+    return enabled ? BedrockRuntimeImpl.getStopReasons(executionAttributes, response) : null;
   }
 
   @Nullable
   @NoMuzzle
-  static Long getUsageInputTokens(Response response) {
-    return enabled ? BedrockRuntimeImpl.getUsageInputTokens(response) : null;
+  static Long getUsageInputTokens(ExecutionAttributes executionAttributes, Response response) {
+    return enabled ? BedrockRuntimeImpl.getUsageInputTokens(executionAttributes, response) : null;
   }
 
   @Nullable
   @NoMuzzle
-  static Long getUsageOutputTokens(Response response) {
-    return enabled ? BedrockRuntimeImpl.getUsageOutputTokens(response) : null;
+  static Long getUsageOutputTokens(ExecutionAttributes executionAttributes, Response response) {
+    return enabled ? BedrockRuntimeImpl.getUsageOutputTokens(executionAttributes, response) : null;
   }
 
   @NoMuzzle
   static void recordRequestEvents(
-      Context otelContext, Logger eventLogger, SdkRequest request, boolean captureMessageContent) {
+      Context otelContext,
+      Logger eventLogger,
+      ExecutionAttributes executionAttributes,
+      SdkRequest request,
+      boolean captureMessageContent) {
     if (enabled) {
       BedrockRuntimeImpl.recordRequestEvents(
-          otelContext, eventLogger, request, captureMessageContent);
+          otelContext, eventLogger, executionAttributes, request, captureMessageContent);
     }
   }
 
   @NoMuzzle
   static void recordResponseEvents(
       Context otelContext,
       Logger eventLogger,
+      ExecutionAttributes executionAttributes,
       SdkResponse response,
       boolean captureMessageContent) {
     if (enabled) {
       BedrockRuntimeImpl.recordResponseEvents(
-          otelContext, eventLogger, response, captureMessageContent);
+          otelContext, eventLogger, executionAttributes, response, captureMessageContent);
     }
   }
 }

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

@@ -5,27 +5,18 @@
 
 package io.opentelemetry.instrumentation.awssdk.v2_2.internal;
 
-import static io.opentelemetry.instrumentation.awssdk.v2_2.internal.TracingExecutionInterceptor.SDK_REQUEST_ATTRIBUTE;
 import static java.util.Collections.emptyList;
 
 import io.opentelemetry.instrumentation.api.incubator.semconv.genai.GenAiAttributesGetter;
 import java.util.Collections;
 import java.util.List;
 import javax.annotation.Nullable;
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
-import software.amazon.awssdk.core.interceptor.SdkExecutionAttribute;
 
 enum BedrockRuntimeAttributesGetter
     implements GenAiAttributesGetter<ExecutionAttributes, Response> {
   INSTANCE;
 
-  // copied from GenAiIncubatingAttributes
-  private static final class GenAiOperationNameIncubatingValues {
-    static final String CHAT = "chat";
-
-    private GenAiOperationNameIncubatingValues() {}
-  }
-
   static final class GenAiSystemIncubatingValues {
     static final String AWS_BEDROCK = "aws.bedrock";
 
@@ -34,18 +25,7 @@ private GenAiSystemIncubatingValues() {}
 
   @Override
   public String getOperationName(ExecutionAttributes executionAttributes) {
-    String operation = executionAttributes.getAttribute(SdkExecutionAttribute.OPERATION_NAME);
-    if (operation != null) {
-      switch (operation) {
-        case "Converse":
-        // fallthrough
-        case "ConverseStream":
-          return GenAiOperationNameIncubatingValues.CHAT;
-        default:
-          return null;
-      }
-    }
-    return null;
+    return BedrockRuntimeAccess.getOperationName(executionAttributes);
   }
 
   @Override
@@ -56,7 +36,7 @@ public String getSystem(ExecutionAttributes executionAttributes) {
   @Nullable
   @Override
   public String getRequestModel(ExecutionAttributes executionAttributes) {
-    return BedrockRuntimeAccess.getModelId(executionAttributes.getAttribute(SDK_REQUEST_ATTRIBUTE));
+    return BedrockRuntimeAccess.getModelId(executionAttributes);
   }
 
   @Nullable
@@ -80,8 +60,7 @@ public Double getRequestFrequencyPenalty(ExecutionAttributes executionAttributes
   @Nullable
   @Override
   public Long getRequestMaxTokens(ExecutionAttributes executionAttributes) {
-    return BedrockRuntimeAccess.getMaxTokens(
-        executionAttributes.getAttribute(SDK_REQUEST_ATTRIBUTE));
+    return BedrockRuntimeAccess.getMaxTokens(executionAttributes);
   }
 
   @Nullable
@@ -93,15 +72,13 @@ public Double getRequestPresencePenalty(ExecutionAttributes executionAttributes)
   @Nullable
   @Override
   public List<String> getRequestStopSequences(ExecutionAttributes executionAttributes) {
-    return BedrockRuntimeAccess.getStopSequences(
-        executionAttributes.getAttribute(SDK_REQUEST_ATTRIBUTE));
+    return BedrockRuntimeAccess.getStopSequences(executionAttributes);
   }
 
   @Nullable
   @Override
   public Double getRequestTemperature(ExecutionAttributes executionAttributes) {
-    return BedrockRuntimeAccess.getTemperature(
-        executionAttributes.getAttribute(SDK_REQUEST_ATTRIBUTE));
+    return BedrockRuntimeAccess.getTemperature(executionAttributes);
   }
 
   @Nullable
@@ -113,7 +90,7 @@ public Double getRequestTopK(ExecutionAttributes executionAttributes) {
   @Nullable
   @Override
   public Double getRequestTopP(ExecutionAttributes executionAttributes) {
-    return BedrockRuntimeAccess.getTopP(executionAttributes.getAttribute(SDK_REQUEST_ATTRIBUTE));
+    return BedrockRuntimeAccess.getTopP(executionAttributes);
   }
 
   @Override
@@ -122,7 +99,7 @@ public List<String> getResponseFinishReasons(
     if (response == null) {
       return emptyList();
     }
-    List<String> stopReasons = BedrockRuntimeAccess.getStopReasons(response);
+    List<String> stopReasons = BedrockRuntimeAccess.getStopReasons(executionAttributes, response);
     if (stopReasons == null) {
       return Collections.emptyList();
     }
@@ -148,7 +125,7 @@ public Long getUsageInputTokens(
     if (response == null) {
       return null;
     }
-    return BedrockRuntimeAccess.getUsageInputTokens(response);
+    return BedrockRuntimeAccess.getUsageInputTokens(executionAttributes, response);
   }
 
   @Nullable
@@ -158,6 +135,6 @@ public Long getUsageOutputTokens(
     if (response == null) {
       return null;
     }
-    return BedrockRuntimeAccess.getUsageOutputTokens(response);
+    return BedrockRuntimeAccess.getUsageOutputTokens(executionAttributes, response);
   }
 }