Merge pull request #13 from yiyuan-he/update-bedrock

feat: update bedrock resource attributes
open-telemetry · Oct 11, 2024 · 170b4c2 · 170b4c2
2 parents 965f787 + 0a93834
commit 170b4c2
Show file tree

Hide file tree

Showing 6 changed files with 38 additions and 2 deletions.
diff --git a/...rc/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/AwsExperimentalAttributes.java b/...rc/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/AwsExperimentalAttributes.java
@@ -23,6 +23,7 @@ final class AwsExperimentalAttributes {
       stringKey("aws.bedrock.knowledge_base.id");
   static final AttributeKey<String> AWS_DATA_SOURCE_ID = stringKey("aws.bedrock.data_source.id");
   static final AttributeKey<String> AWS_GUARDRAIL_ID = stringKey("aws.bedrock.guardrail.id");
+  static final AttributeKey<String> AWS_GUARDRAIL_ARN = stringKey("aws.bedrock.guardrail.arn");
 
   // TODO: Merge in gen_ai attributes in opentelemetry-semconv-incubating once upgrade to v1.26.0
   static final AttributeKey<String> AWS_BEDROCK_RUNTIME_MODEL_ID =

diff --git a/.../io/opentelemetry/instrumentation/awssdk/v1_11/AwsSdkExperimentalAttributesExtractor.java b/.../io/opentelemetry/instrumentation/awssdk/v1_11/AwsSdkExperimentalAttributesExtractor.java
@@ -11,6 +11,7 @@
 import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_BEDROCK_SYSTEM;
 import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_BUCKET_NAME;
 import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_ENDPOINT;
+import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_GUARDRAIL_ARN;
 import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_GUARDRAIL_ID;
 import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_KNOWLEDGE_BASE_ID;
 import static io.opentelemetry.instrumentation.awssdk.v1_11.AwsExperimentalAttributes.AWS_LAMBDA_NAME;
@@ -154,6 +155,7 @@ private static void bedrockOnEnd(
     switch (serviceName) {
       case BEDROCK_SERVICE:
         setAttribute(attributes, AWS_GUARDRAIL_ID, awsResp, RequestAccess::getGuardrailId);
+        setAttribute(attributes, AWS_GUARDRAIL_ARN, awsResp, RequestAccess::getGuardrailArn);
         break;
       case BEDROCK_AGENT_SERVICE:
         String responseClassName = awsResp.getClass().getSimpleName();

diff --git a/...11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/RequestAccess.java b/...11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/RequestAccess.java
@@ -8,6 +8,7 @@
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.MethodType;
+import java.lang.reflect.Method;
 import javax.annotation.Nullable;
 
 final class RequestAccess {
@@ -110,6 +111,11 @@ static String getGuardrailId(Object request) {
     return invokeOrNull(access.getGuardrailId, request);
   }
 
+  @Nullable
+  static String getGuardrailArn(Object request) {
+    return findNestedAccessorOrNull(request, "getGuardrailArn");
+  }
+
   @Nullable
   static String getModelId(Object request) {
     RequestAccess access = REQUEST_ACCESSORS.get(request.getClass());
@@ -173,4 +179,21 @@ private static MethodHandle findAccessorOrNull(Class<?> clz, String methodName)
       return null;
     }
   }
+
+  @Nullable
+  private static String findNestedAccessorOrNull(Object obj, String... methodNames) {
+    Object current = obj;
+    for (String methodName : methodNames) {
+      if (current == null) {
+        return null;
+      }
+      try {
+        Method method = current.getClass().getMethod(methodName);
+        current = method.invoke(current);
+      } catch (Exception e) {
+        return null;
+      }
+    }
+    return (current instanceof String) ? (String) current : null;
+  }
 }
diff --git a/...src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/AwsExperimentalAttributes.java b/...src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/AwsExperimentalAttributes.java
@@ -16,6 +16,7 @@ final class AwsExperimentalAttributes {
   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_GUARDRAIL_ID = stringKey("aws.bedrock.guardrail.id");
+  static final AttributeKey<String> AWS_GUARDRAIL_ARN = stringKey("aws.bedrock.guardrail.arn");
   static final AttributeKey<String> AWS_AGENT_ID = stringKey("aws.bedrock.agent.id");
   static final AttributeKey<String> AWS_DATA_SOURCE_ID = stringKey("aws.bedrock.data_source.id");
   static final AttributeKey<String> AWS_KNOWLEDGE_BASE_ID =

diff --git a/...2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/AwsSdkRequest.java b/...2.2/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/AwsSdkRequest.java
@@ -7,6 +7,7 @@
 
 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.BEDROCKAGENTRUNTIMEOPERATION;
 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;
@@ -42,7 +43,7 @@ enum AwsSdkRequest {
   SqsRequest(SQS, "SqsRequest"),
   KinesisRequest(KINESIS, "KinesisRequest"),
   BedrockRequest(BEDROCK, "BedrockRequest"),
-  BedrockAgentRuntimeRequest(BEDROCKAGENTOPERATION, "BedrockAgentRuntimeRequest"),
+  BedrockAgentRuntimeRequest(BEDROCKAGENTRUNTIMEOPERATION, "BedrockAgentRuntimeRequest"),
   BedrockRuntimeRequest(BEDROCKRUNTIME, "BedrockRuntimeRequest"),
   // BedrockAgent API based requests. We only support operations that are related to
   // Agent/DataSources/KnowledgeBases

diff --git a/...library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/AwsSdkRequestType.java b/...library/src/main/java/io/opentelemetry/instrumentation/awssdk/v2_2/AwsSdkRequestType.java
@@ -8,6 +8,7 @@
 import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.AWS_AGENT_ID;
 import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.AWS_BUCKET_NAME;
 import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.AWS_DATA_SOURCE_ID;
+import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.AWS_GUARDRAIL_ARN;
 import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.AWS_GUARDRAIL_ID;
 import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.AWS_KNOWLEDGE_BASE_ID;
 import static io.opentelemetry.instrumentation.awssdk.v2_2.AwsExperimentalAttributes.AWS_LAMBDA_ARN;
@@ -34,9 +35,16 @@ enum AwsSdkRequestType {
   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_GUARDRAIL_ID.getKey(), "guardrailIdentifier")),
+  BEDROCK(
+      request(AWS_GUARDRAIL_ID.getKey(), "guardrailIdentifier"),
+      response(AWS_GUARDRAIL_ARN.getKey(), "guardrailArn")),
   BEDROCKAGENTOPERATION(
       request(AWS_AGENT_ID.getKey(), "agentId"), response(AWS_AGENT_ID.getKey(), "agentId")),
+  BEDROCKAGENTRUNTIMEOPERATION(
+      request(AWS_AGENT_ID.getKey(), "agentId"),
+      response(AWS_AGENT_ID.getKey(), "agentId"),
+      request(AWS_KNOWLEDGE_BASE_ID.getKey(), "knowledgeBaseId"),
+      response(AWS_KNOWLEDGE_BASE_ID.getKey(), "knowledgeBaseId")),
   BEDROCKDATASOURCEOPERATION(
       request(AWS_DATA_SOURCE_ID.getKey(), "dataSourceId"),
       response(AWS_DATA_SOURCE_ID.getKey(), "dataSourceId")),