added guardrail arn to bedrock

Author: liustve

aws-distro-opentelemetry-node-autoinstrumentation/src/aws-attribute-keys.ts

@@ -36,6 +36,7 @@ export const AWS_ATTRIBUTE_KEYS: { [key: string]: string } = {
   AWS_BEDROCK_KNOWLEDGE_BASE_ID: 'aws.bedrock.knowledge_base.id',
   AWS_BEDROCK_AGENT_ID: 'aws.bedrock.agent.id',
   AWS_BEDROCK_GUARDRAIL_ID: 'aws.bedrock.guardrail.id',
+  AWS_BEDROCK_GUARDRAIL_ARN: 'aws.bedrock.guardrail.arn',
   AWS_SNS_TOPIC_ARN: 'aws.sns.topic.arn',
   AWS_SECRETSMANAGER_SECRET_ARN: 'aws.secretsmanager.secret.arn',
   AWS_STEPFUNCTIONS_STATEMACHINE_ARN: 'aws.stepfunctions.state_machine.arn',

aws-distro-opentelemetry-node-autoinstrumentation/src/aws-metric-attribute-generator.ts

@@ -437,6 +437,9 @@ export class AwsMetricAttributeGenerator implements MetricAttributeGenerator {
         remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
           span.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ID]
         );
+        cloudFormationIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
+          span.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ARN]
+        );
       } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID)) {
         remoteResourceType = NORMALIZED_BEDROCK_SERVICE_NAME + '::KnowledgeBase';
         remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(

aws-distro-opentelemetry-node-autoinstrumentation/src/patches/aws/services/bedrock.ts

@@ -15,6 +15,7 @@ const AGENT_ID: string = 'agentId';
 const KNOWLEDGE_BASE_ID: string = 'knowledgeBaseId';
 const DATA_SOURCE_ID: string = 'dataSourceId';
 const GUARDRAIL_ID: string = 'guardrailId';
+const GUARDRAIL_ARN: string = 'guardrailArn';
 const MODEL_ID: string = 'modelId';
 const AWS_BEDROCK_SYSTEM: string = 'aws_bedrock';
 
@@ -182,10 +183,15 @@ export class BedrockServiceExtension implements ServiceExtension {
     };
   }
   responseHook(response: NormalizedResponse, span: Span, tracer: Tracer, config: AwsSdkInstrumentationConfig): void {
-    const guardrail_id = response.data[GUARDRAIL_ID];
+    const guardrailId = response.data[GUARDRAIL_ID];
+    const guardrailArn = response.data[GUARDRAIL_ARN];
 
-    if (guardrail_id) {
-      span.setAttribute(AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ID, guardrail_id);
+    if (guardrailId) {
+      span.setAttribute(AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ID, guardrailId);
+    }
+
+    if (guardrailArn) {
+      span.setAttribute(AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ARN, guardrailArn);
     }
   }
 }

aws-distro-opentelemetry-node-autoinstrumentation/test/aws-metric-attribute-generator.test.ts

@@ -849,10 +849,12 @@ describe('AwsMetricAttributeGeneratorTest', () => {
     validateRemoteResourceAttributes('AWS::Bedrock::DataSource', 'test_datasource_^^id');
     mockAttribute(AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID, undefined);
 
-    // Validate behaviour of AWS_BEDROCK_GUARDRAIL_ID attribute, then remove it.
+    // Validate behaviour of AWS_BEDROCK_GUARDRAIL_ID and AWS_BEDROCK_GUARDRAIL_ARN attribute, then remove it.
     mockAttribute(AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ID, 'test_guardrail_id');
+    mockAttribute(AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ARN, 'test_guardrail_arn');
     validateRemoteResourceAttributes('AWS::Bedrock::Guardrail', 'test_guardrail_id');
     mockAttribute(AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ID, undefined);
+    mockAttribute(AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ARN, undefined);
 
     // Validate behaviour of AWS_BEDROCK_GUARDRAIL_ID attribute with special chars(^), then remove it.
     mockAttribute(AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ID, 'test_guardrail_^id');

aws-distro-opentelemetry-node-autoinstrumentation/test/patches/aws/services/bedrock.test.ts

@@ -236,9 +236,11 @@ describe('Bedrock', () => {
   describe('GetGuardrail', () => {
     it('adds guardrailId to span', async () => {
       const dummyGuardrailIdentifier: string = 'ABCDEFGH';
+      const guardrailArn: string = 'arn:aws:bedrock:us-east-1:123456789012:guardrail/abc123';
 
       nock(`https://bedrock.${region}.amazonaws.com`).get(`/guardrails/${dummyGuardrailIdentifier}`).reply(200, {
         guardrailId: dummyGuardrailIdentifier,
+        guardrailArn: guardrailArn,
       });
 
       await bedrock
@@ -258,6 +260,7 @@ describe('Bedrock', () => {
       expect(getGuardrailSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]).toBeUndefined();
       expect(getGuardrailSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID]).toBeUndefined();
       expect(getGuardrailSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ID]).toBe(dummyGuardrailIdentifier);
+      expect(getGuardrailSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ARN]).toBe(guardrailArn);
       expect(getGuardrailSpan.kind).toBe(SpanKind.CLIENT);
     });
   });

aws-distro-opentelemetry-node-autoinstrumentation/test/patches/instrumentation-patch.test.ts

@@ -38,6 +38,7 @@ const _QUEUE_URL: string = 'https://sqs.us-east-1.amazonaws.com/123412341234/que
 const _BEDROCK_AGENT_ID: string = 'agentId';
 const _BEDROCK_DATASOURCE_ID: string = 'DataSourceId';
 const _BEDROCK_GUARDRAIL_ID: string = 'GuardrailId';
+const _BEDROCK_GUARDRAIL_ARN: string = 'arn:aws:bedrock:us-east-1:123456789012:guardrail/abc123';
 const _BEDROCK_KNOWLEDGEBASE_ID: string = 'KnowledgeBaseId';
 const _GEN_AI_SYSTEM: string = 'aws_bedrock';
 const _GEN_AI_REQUEST_MODEL: string = 'genAiReuqestModelId';
@@ -250,8 +251,11 @@ describe('InstrumentationPatchTest', () => {
     // Expect no-op from attribute extraction in Bedrock
     expect(Object.entries(bedrockAttributes).length).toEqual(0);
     const bedrockAttributesAfterResponse: Attributes = doResponseHookBedrock(services, 'Bedrock');
-    expect(Object.entries(bedrockAttributesAfterResponse).length).toBe(1);
+    expect(Object.entries(bedrockAttributesAfterResponse).length).toBe(2);
     expect(bedrockAttributesAfterResponse[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ID]).toEqual(_BEDROCK_GUARDRAIL_ID);
+    expect(bedrockAttributesAfterResponse[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ARN]).toEqual(
+      _BEDROCK_GUARDRAIL_ARN
+    );
   });
 
   it('Bedrock Agent with patching', () => {
@@ -519,6 +523,7 @@ describe('InstrumentationPatchTest', () => {
         dataSourceId: _BEDROCK_DATASOURCE_ID,
         knowledgeBaseId: _BEDROCK_KNOWLEDGEBASE_ID,
         guardrailId: _BEDROCK_GUARDRAIL_ID,
+        guardrailArn: _BEDROCK_GUARDRAIL_ARN,
         modelId: _GEN_AI_REQUEST_MODEL,
       },
       request: {