support bedrock and bedrock runtime resources

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

@@ -4,6 +4,11 @@ plugins {
 
 dependencies {
   implementation("io.opentelemetry.contrib:opentelemetry-aws-xray-propagator")
+  implementation("com.fasterxml.jackson.core:jackson-databind") {
+    version {
+      strictly("[2.13.3,)")
+    }
+  }
 
   library("software.amazon.awssdk:aws-core:2.2.0")
   library("software.amazon.awssdk:sqs:2.2.0")

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

@@ -0,0 +1,52 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.awssdk.v2_2;
+
+import static io.opentelemetry.api.common.AttributeKey.stringKey;
+
+import io.opentelemetry.api.common.AttributeKey;
+
+final class AwsExperimentalAttributes {
+  static final AttributeKey<String> AWS_BUCKET_NAME = stringKey("aws.bucket.name");
+  static final AttributeKey<String> AWS_QUEUE_URL = stringKey("aws.queue.url");
+  static final AttributeKey<String> AWS_QUEUE_NAME = stringKey("aws.queue.name");
+  static final AttributeKey<String> AWS_STREAM_NAME = stringKey("aws.stream.name");
+  static final AttributeKey<String> AWS_TABLE_NAME = stringKey("aws.table.name");
+  static final AttributeKey<String> AWS_BEDROCK_GUARDRAIL_ID =
+      stringKey("aws.bedrock.guardrail_id");
+  static final AttributeKey<String> AWS_BEDROCK_AGENT_ID = stringKey("aws.bedrock.agent_id");
+  static final AttributeKey<String> AWS_BEDROCK_DATASOURCE_ID =
+      stringKey("aws.bedrock.datasource_id");
+  static final AttributeKey<String> AWS_BEDROCK_KNOWLEDGEBASE_ID =
+      stringKey("aws.bedrock.knowledgebase_id");
+  static final AttributeKey<String> GEN_AI_FINISH_REASON =
+      stringKey("gen_ai.response.finish_reason");
+  static final AttributeKey<String> GEN_AI_PROMPT_TOKENS = stringKey("gen_ai.usage.prompt_tokens");
+  static final AttributeKey<String> GEN_AI_COMPLETION_TOKENS =
+      stringKey("gen_ai.usage.completion_tokens");
+
+  // OTel GenAI/LLM group has defined gen_ai attributes but not yet add it in
+  // semantic-conventions-java package
+  // https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-spans.md#genai-attributes
+  static final AttributeKey<String> GEN_AI_MODEL = stringKey("gen_ai.request.model");
+  static final AttributeKey<String> GEN_AI_TEMPERATURE = stringKey("gen_ai.request.temperature");
+  static final AttributeKey<String> GEN_AI_TOP_P = stringKey("gen_ai.request.top_p");
+  static final AttributeKey<String> GEN_AI_MAX_TOKENS = stringKey("gen_ai.request.max_tokens");
+  static final AttributeKey<String> GEN_AI_SYSTEM = stringKey("gen_ai.system");
+
+  private AwsExperimentalAttributes() {}
+
+  static boolean isGenAiAttribute(String attributeKey) {
+    return attributeKey.equals(GEN_AI_MODEL.getKey())
+        || attributeKey.equals(GEN_AI_FINISH_REASON.getKey())
+        || attributeKey.equals(GEN_AI_PROMPT_TOKENS.getKey())
+        || attributeKey.equals(GEN_AI_COMPLETION_TOKENS.getKey())
+        || attributeKey.equals(GEN_AI_TEMPERATURE.getKey())
+        || attributeKey.equals(GEN_AI_TOP_P.getKey())
+        || attributeKey.equals(GEN_AI_MAX_TOKENS.getKey())
+        || attributeKey.equals(GEN_AI_SYSTEM.getKey());
+  }
+}

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

@@ -5,6 +5,11 @@
 
 package io.opentelemetry.instrumentation.awssdk.v2_2;
 
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsSdkRequestType.BEDROCK;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsSdkRequestType.BEDROCKAGENTOPERATION;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsSdkRequestType.BEDROCKDATASOURCEOPERATION;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsSdkRequestType.BEDROCKKNOWLEDGEBASEOPERATION;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsSdkRequestType.BEDROCKRUNTIME;
 import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsSdkRequestType.DYNAMODB;
 import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsSdkRequestType.KINESIS;
 import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsSdkRequestType.S3;
@@ -32,6 +37,43 @@ enum AwsSdkRequest {
   S3Request(S3, "S3Request"),
   SqsRequest(SQS, "SqsRequest"),
   KinesisRequest(KINESIS, "KinesisRequest"),
+  BedrockRequest(BEDROCK, "BedrockRequest"),
+  BedrockAgentRuntimeRequest(BEDROCKAGENTOPERATION, "BedrockAgentRuntimeRequest"),
+  BedrockRuntimeRequest(BEDROCKRUNTIME, "BedrockRuntimeRequest"),
+  // BedrockAgent API based requests
+  BedrockCreateAgentActionGroupRequest(BEDROCKAGENTOPERATION, "CreateAgentActionGroupRequest"),
+  BedrockCreateAgentAliasRequest(BEDROCKAGENTOPERATION, "CreateAgentAliasRequest"),
+  BedrockDeleteAgentActionGroupRequest(BEDROCKAGENTOPERATION, "DeleteAgentActionGroupRequest"),
+  BedrockDeleteAgentAliasRequest(BEDROCKAGENTOPERATION, "DeleteAgentAliasRequest"),
+  BedrockDeleteAgentVersionRequest(BEDROCKAGENTOPERATION, "DeleteAgentVersionRequest"),
+  BedrockGetAgentActionGroupRequest(BEDROCKAGENTOPERATION, "GetAgentActionGroupRequest"),
+  BedrockGetAgentAliasRequest(BEDROCKAGENTOPERATION, "GetAgentAliasRequest"),
+  BedrockGetAgentRequest(BEDROCKAGENTOPERATION, "GetAgentRequest"),
+  BedrockGetAgentVersionRequest(BEDROCKAGENTOPERATION, "GetAgentVersionRequest"),
+  BedrockListAgentActionGroupsRequest(BEDROCKAGENTOPERATION, "ListAgentActionGroupsRequest"),
+  BedrockListAgentAliasesRequest(BEDROCKAGENTOPERATION, "ListAgentAliasesRequest"),
+  BedrockListAgentKnowledgeBasesRequest(BEDROCKAGENTOPERATION, "ListAgentKnowledgeBasesRequest"),
+  BedrocListAgentVersionsRequest(BEDROCKAGENTOPERATION, "ListAgentVersionsRequest"),
+  BedrockPrepareAgentRequest(BEDROCKAGENTOPERATION, "PrepareAgentRequest"),
+  BedrockUpdateAgentActionGroupRequest(BEDROCKAGENTOPERATION, "UpdateAgentActionGroupRequest"),
+  BedrockUpdateAgentAliasRequest(BEDROCKAGENTOPERATION, "UpdateAgentAliasRequest"),
+  BedrockUpdateAgentRequest(BEDROCKAGENTOPERATION, "UpdateAgentRequest"),
+  BedrockBedrockAgentRequest(BEDROCKAGENTOPERATION, "BedrockAgentRequest"),
+  BedrockDeleteDataSourceRequest(BEDROCKDATASOURCEOPERATION, "DeleteDataSourceRequest"),
+  BedrockGetDataSourceRequest(BEDROCKDATASOURCEOPERATION, "GetDataSourceRequest"),
+  BedrockUpdateDataSourceRequest(BEDROCKDATASOURCEOPERATION, "UpdateDataSourceRequest"),
+  BedrocAssociateAgentKnowledgeBaseRequest(
+      BEDROCKKNOWLEDGEBASEOPERATION, "AssociateAgentKnowledgeBaseRequest"),
+  BedrockCreateDataSourceRequest(BEDROCKKNOWLEDGEBASEOPERATION, "CreateDataSourceRequest"),
+  BedrockDeleteKnowledgeBaseRequest(BEDROCKKNOWLEDGEBASEOPERATION, "DeleteKnowledgeBaseRequest"),
+  BedrockDisassociateAgentKnowledgeBaseRequest(
+      BEDROCKKNOWLEDGEBASEOPERATION, "DisassociateAgentKnowledgeBaseRequest"),
+  BedrockGetAgentKnowledgeBaseRequest(
+      BEDROCKKNOWLEDGEBASEOPERATION, "GetAgentKnowledgeBaseRequest"),
+  BedrockGetKnowledgeBaseRequest(BEDROCKKNOWLEDGEBASEOPERATION, "GetKnowledgeBaseRequest"),
+  BedrockListDataSourcesRequest(BEDROCKKNOWLEDGEBASEOPERATION, "ListDataSourcesRequest"),
+  BedrockUpdateAgentKnowledgeBaseRequest(
+      BEDROCKKNOWLEDGEBASEOPERATION, "UpdateAgentKnowledgeBaseRequest"),
   // specific requests
   BatchGetItem(
       DYNAMODB,

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

@@ -5,17 +5,48 @@
 
 package io.opentelemetry.instrumentation.awssdk.v2_2;
 
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.AWS_BEDROCK_AGENT_ID;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.AWS_BEDROCK_DATASOURCE_ID;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.AWS_BEDROCK_GUARDRAIL_ID;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.AWS_BEDROCK_KNOWLEDGEBASE_ID;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.AWS_BUCKET_NAME;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.AWS_QUEUE_NAME;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.AWS_QUEUE_URL;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.AWS_STREAM_NAME;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.AWS_TABLE_NAME;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.GEN_AI_COMPLETION_TOKENS;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.GEN_AI_FINISH_REASON;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.GEN_AI_MAX_TOKENS;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.GEN_AI_MODEL;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.GEN_AI_PROMPT_TOKENS;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.GEN_AI_TEMPERATURE;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.GEN_AI_TOP_P;
 import static io.opentelemetry.instrumentation.awssdk.v2_2.FieldMapping.request;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.FieldMapping.response;
 
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
 enum AwsSdkRequestType {
-  S3(request("aws.bucket.name", "Bucket")),
-  SQS(request("aws.queue.url", "QueueUrl"), request("aws.queue.name", "QueueName")),
-  KINESIS(request("aws.stream.name", "StreamName")),
-  DYNAMODB(request("aws.table.name", "TableName"));
+  S3(request(AWS_BUCKET_NAME.getKey(), "Bucket")),
+  SQS(request(AWS_QUEUE_URL.getKey(), "QueueUrl"), request(AWS_QUEUE_NAME.getKey(), "QueueName")),
+  KINESIS(request(AWS_STREAM_NAME.getKey(), "StreamName")),
+  DYNAMODB(request(AWS_TABLE_NAME.getKey(), "TableName")),
+  BEDROCK(request(AWS_BEDROCK_GUARDRAIL_ID.getKey(), "guardrailIdentifier")),
+  BEDROCKAGENTOPERATION(
+      request(AWS_BEDROCK_AGENT_ID.getKey(), "agentId"),
+      response(AWS_BEDROCK_AGENT_ID.getKey(), "agentId")),
+  BEDROCKDATASOURCEOPERATION(request(AWS_BEDROCK_DATASOURCE_ID.getKey(), "dataSourceId")),
+  BEDROCKKNOWLEDGEBASEOPERATION(request(AWS_BEDROCK_KNOWLEDGEBASE_ID.getKey(), "knowledgeBaseId")),
+  BEDROCKRUNTIME(
+      request(GEN_AI_MODEL.getKey(), "modelId"),
+      request(GEN_AI_TEMPERATURE.getKey(), "body"),
+      request(GEN_AI_TOP_P.getKey(), "body"),
+      request(GEN_AI_MAX_TOKENS.getKey(), "body"),
+      response(GEN_AI_FINISH_REASON.getKey(), "body"),
+      response(GEN_AI_PROMPT_TOKENS.getKey(), "body"),
+      response(GEN_AI_COMPLETION_TOKENS.getKey(), "body"));
 
   // Wrapping in unmodifiableMap
   @SuppressWarnings("ImmutableEnumChecker")

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

@@ -65,8 +65,15 @@ private void mapToAttributes(
     for (int i = 1; i < path.size() && target != null; i++) {
       target = next(target, path.get(i));
     }
+    String value;
     if (target != null) {
-      String value = serializer.serialize(target);
+      if (AwsExperimentalAttributes.isGenAiAttribute(fieldMapping.getAttribute())) {
+        value = serializer.serialize(fieldMapping.getAttribute(), target);
+        span.setAttribute("gen_ai.system", "AWS Bedrock");
+      } else {
+        value = serializer.serialize(target);
+      }
+
       if (!StringUtils.isEmpty(value)) {
         span.setAttribute(fieldMapping.getAttribute(), value);
       }

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

@@ -5,13 +5,17 @@
 
 package io.opentelemetry.instrumentation.awssdk.v2_2;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.Collection;
 import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
+import software.amazon.awssdk.core.SdkBytes;
 import software.amazon.awssdk.core.SdkPojo;
 import software.amazon.awssdk.http.ContentStreamProvider;
 import software.amazon.awssdk.http.SdkHttpFullRequest;
@@ -20,6 +24,7 @@
 import software.amazon.awssdk.utils.StringUtils;
 
 class Serializer {
+  private static final ObjectMapper objectMapper = new ObjectMapper();
 
   @Nullable
   String serialize(Object target) {
@@ -37,6 +42,9 @@ String serialize(Object target) {
     if (target instanceof Map) {
       return serialize(((Map<?, ?>) target).keySet());
     }
+    if (target instanceof SdkBytes) {
+      return serialize((SdkBytes) target);
+    }
     // simple type
     return target.toString();
   }
@@ -65,4 +73,120 @@ private String serialize(Collection<?> collection) {
     String serialized = collection.stream().map(this::serialize).collect(Collectors.joining(","));
     return (StringUtils.isEmpty(serialized) ? null : "[" + serialized + "]");
   }
+
+  @Nullable
+  String serialize(String attributeName, Object target) {
+    try {
+      JsonNode jsonBody;
+      if (target instanceof SdkBytes) {
+        jsonBody = objectMapper.readTree(((SdkBytes) target).asUtf8String());
+      } else {
+        if (target != null) {
+          return target.toString();
+        }
+        return null;
+      }
+      switch (attributeName) {
+        case "gen_ai.response.finish_reason":
+          return getFinishReason(jsonBody);
+        case "gen_ai.usage.completion_tokens":
+          return getOutputTokens(jsonBody);
+        case "gen_ai.usage.prompt_tokens":
+          return getInputTokens(jsonBody);
+        case "gen_ai.request.top_p":
+          return getTopP(jsonBody);
+        case "gen_ai.request.temperature":
+          return getTemperature(jsonBody);
+        case "gen_ai.request.max_tokens":
+          return getMaxTokens(jsonBody);
+        default:
+          return null;
+      }
+    } catch (RuntimeException | JsonProcessingException e) {
+      throw new IllegalStateException("Failed to instantiate operation class", e);
+    }
+  }
+
+  private static String getFinishReason(JsonNode body) {
+    if (body.has("stop_reason")) {
+      return body.get("stop_reason").asText();
+    } else if (body.has("results")) {
+      JsonNode result = body.get("results").get(0);
+      if (result.has("completionReason")) {
+        return result.get("completionReason").asText();
+      }
+    }
+    return null;
+  }
+
+  private static String getInputTokens(JsonNode body) {
+    if (body.has("prompt_token_count")) {
+      return String.valueOf(body.get("prompt_token_count").asInt());
+    } else if (body.has("inputTextTokenCount")) {
+      return String.valueOf(body.get("inputTextTokenCount").asInt());
+    } else if (body.has("usage")) {
+      JsonNode usage = body.get("usage");
+      if (usage.has("input_tokens")) {
+        return String.valueOf(usage.get("input_tokens").asInt());
+      }
+    }
+    return null;
+  }
+
+  private static String getOutputTokens(JsonNode body) {
+    if (body.has("generation_token_count")) {
+      return String.valueOf(body.get("generation_token_count").asInt());
+    } else if (body.has("results")) {
+      JsonNode result = body.get("results").get(0);
+      if (result.has("tokenCount")) {
+        return String.valueOf(result.get("tokenCount").asInt());
+      }
+    } else if (body.has("inputTextTokenCount")) {
+      return String.valueOf(body.get("inputTextTokenCount").asInt());
+    } else if (body.has("usage")) {
+      JsonNode usage = body.get("usage");
+      if (usage.has("output_tokens")) {
+        return String.valueOf(usage.get("output_tokens").asInt());
+      }
+    }
+    return null;
+  }
+
+  private static String getTopP(JsonNode body) {
+    if (body.has("top_p")) {
+      return String.valueOf(body.get("top_p").asDouble());
+    } else if (body.has("textGenerationConfig")) {
+      JsonNode usage = body.get("textGenerationConfig");
+      if (usage.has("topP")) {
+        return String.valueOf(usage.get("topP").asInt());
+      }
+    }
+    return null;
+  }
+
+  private static String getTemperature(JsonNode body) {
+    if (body.has("temperature")) {
+      return String.valueOf(body.get("temperature").asDouble());
+    } else if (body.has("textGenerationConfig")) {
+      JsonNode usage = body.get("textGenerationConfig");
+      if (usage.has("temperature")) {
+        return String.valueOf(usage.get("temperature").asDouble());
+      }
+    }
+    return null;
+  }
+
+  private static String getMaxTokens(JsonNode body) {
+    if (body.has("max_tokens")) {
+      return String.valueOf(body.get("max_tokens").asInt());
+    } else if (body.has("max_gen_len")) {
+      return String.valueOf(body.get("max_gen_len").asInt());
+    } else if (body.has("textGenerationConfig")) {
+      JsonNode usage = body.get("textGenerationConfig");
+      if (usage.has("maxTokenCount")) {
+        return String.valueOf(usage.get("maxTokenCount").asInt());
+      }
+    }
+    return null;
+  }
 }

instrumentation/aws-sdk/aws-sdk-2.2/testing/src/main/groovy/io/opentelemetry/instrumentation/awssdk/v2_2/AbstractAws2ClientTest.groovy

@@ -120,7 +120,25 @@ abstract class AbstractAws2ClientTest extends AbstractAws2ClientCoreTest {
               "$SemanticAttributes.MESSAGING_SYSTEM" "AmazonSQS"
             } else if (service == "Kinesis") {
               "aws.stream.name" "somestream"
+            } else if (service == "Kinesis") {
+              "aws.stream.name" "somestream"
+            } else if (service == "Bedrock" && operation == "GetGuardrail") {
+              "aws.bedrock.guardrail_id" "guardrailId"
+            } else if (service == "BedrockAgent" && operation == "GetAgent") {
+              "aws.bedrock.agent_id" "agentId"
+            } else if (service == "BedrockAgent" && operation == "GetKnowledgeBase") {
+              "aws.bedrock.knowledgebase_id" "knowledgeBaseId"
+            } else if (service == "BedrockAgent" && operation == "GetDataSource") {
+              "aws.bedrock.datasource_id" "datasourceId"
+            } else if (service == "BedrockRuntime" && operation == "InvokeModel") {
+              "gen_ai.request.top_p" "0.9"
+              "gen_ai.request.temperature" "0.7"
+              "gen_ai.request.model" "meta.llama2-13b-chat-v1"
+              "gen_ai.request.max_tokens" "100"
+              "gen_ai.system" "AWS Bedrock"
+              "gen_ai.response.finish_reasons" "length"
             }
+
           }
         }
       }