From 77ce7719848c433e196693b60993bde061c326f9 Mon Sep 17 00:00:00 2001
From: Michael He <53622546+yiyuan-he@users.noreply.github.com>
Date: Fri, 11 Oct 2024 17:08:42 -0700
Subject: [PATCH] feat: update bedrock guardrail cfn primary id (#905)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

### *Summary of Changes:*

- Updating the CFN Primary Id for Guardrail to be the ARN. (Verified
this in the AWS Cloudformation Console)

![Screenshot 2024-10-09 at 12 44
12 PM](https://github.com/user-attachments/assets/1a7437f1-4269-42f0-b4d9-db9f5f37b3a7)

- Updated `BEDROCKAGENTOPERATION` enum.
- Our Bedrock Agent Runtime should support both Agent Operation and
KnowledgeBased Operation [like in Java SDK
v1](https://github.com/mxiamxia/opentelemetry-java-instrumentation/blob/965f7878c79c2be5851f23cf272edd39585f42af/instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/AwsSdkExperimentalAttributesExtractor.java#L133).

Related changes in upstream package:
https://github.com/mxiamxia/opentelemetry-java-instrumentation/pull/13

### *Testing Plan:*

Set up sample apps and manually verified correct span attributes were
generated via auto-instrumentation.

`Java SDK v1`

![bedrock-guardrail-span-data-verification-v1](https://github.com/user-attachments/assets/70b2c126-cf73-4a70-b63f-e4cd78101604)

`Java SDK v2`

![bedrock-guardrail-span-data-verification-v2](https://github.com/user-attachments/assets/d4857400-3f78-4e5d-aa1a-bb0d26b6213e)


By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 license.

Co-authored-by: Min Xia <xiami@amazon.com>
---
 .../opentelemetry/javaagent/providers/AwsAttributeKeys.java    | 2 ++
 .../javaagent/providers/AwsMetricAttributeGenerator.java       | 3 +++
 2 files changed, 5 insertions(+)

diff --git a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsAttributeKeys.java b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsAttributeKeys.java
index ef2999a28a..10cf3e9d1e 100644
--- a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsAttributeKeys.java
+++ b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsAttributeKeys.java
@@ -89,4 +89,6 @@ private AwsAttributeKeys() {}
       AttributeKey.stringKey("aws.bedrock.data_source.id");
   static final AttributeKey<String> AWS_GUARDRAIL_ID =
       AttributeKey.stringKey("aws.bedrock.guardrail.id");
+  static final AttributeKey<String> AWS_GUARDRAIL_ARN =
+      AttributeKey.stringKey("aws.bedrock.guardrail.arn");
 }
diff --git a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsMetricAttributeGenerator.java b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsMetricAttributeGenerator.java
index 63c806ea98..1fcfe02395 100644
--- a/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsMetricAttributeGenerator.java
+++ b/awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsMetricAttributeGenerator.java
@@ -44,6 +44,7 @@
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_BUCKET_NAME;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_DATA_SOURCE_ID;
+import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_GUARDRAIL_ARN;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_GUARDRAIL_ID;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_KNOWLEDGE_BASE_ID;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LAMBDA_NAME;
@@ -455,6 +456,8 @@ private static void setRemoteResourceTypeAndIdentifier(SpanData span, Attributes
         remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::Guardrail");
         remoteResourceIdentifier =
             Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_GUARDRAIL_ID)));
+        cloudformationPrimaryIdentifier =
+            Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_GUARDRAIL_ARN)));
       } else if (isKeyPresent(span, GEN_AI_REQUEST_MODEL)) {
         remoteResourceType = Optional.of(NORMALIZED_BEDROCK_SERVICE_NAME + "::Model");
         remoteResourceIdentifier =