added knowledge base id to span

Author: liustve

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

@@ -61,6 +61,7 @@ const knowledgeBaseOperationAttributeKeyMapping = {
 };
 const dataSourceOperationAttributeKeyMapping = {
   [AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID]: DATA_SOURCE_ID,
+  [AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]: KNOWLEDGE_BASE_ID,
 };
 
 // This map allows us to get all relevant attribute key mappings for a given operation

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

@@ -134,7 +134,7 @@ describe('BedrockAgent', () => {
       expect(getDataSourceSpans.length).toBe(1);
       const getDataSourceSpan = getDataSourceSpans[0];
       expect(getDataSourceSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID]).toBeUndefined();
-      expect(getDataSourceSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]).toBeUndefined();
+      expect(getDataSourceSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]);
       expect(getDataSourceSpan.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID]).toBe(dummyDataSourceId);
       expect(getDataSourceSpan.kind).toBe(SpanKind.CLIENT);
     });

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

@@ -284,18 +284,25 @@ describe('InstrumentationPatchTest', () => {
       ListDataSources: { 'aws.bedrock.knowledge_base.id': _BEDROCK_KNOWLEDGEBASE_ID },
       UpdateAgentKnowledgeBase: { 'aws.bedrock.knowledge_base.id': _BEDROCK_KNOWLEDGEBASE_ID },
       DeleteDataSource: { 'aws.bedrock.data_source.id': _BEDROCK_DATASOURCE_ID },
-      GetDataSource: { 'aws.bedrock.data_source.id': _BEDROCK_DATASOURCE_ID },
+      GetDataSource: {
+        'aws.bedrock.data_source.id': _BEDROCK_DATASOURCE_ID,
+        'aws.bedrock.knowledge_base.id': _BEDROCK_KNOWLEDGEBASE_ID,
+      },
       UpdateDataSource: { 'aws.bedrock.data_source.id': _BEDROCK_DATASOURCE_ID },
     };
 
     for (const [operation, attribute_tuple] of Object.entries(operation_to_expected_attribute)) {
       const bedrockAttributes: Attributes = doExtractBedrockAttributes(services, 'BedrockAgent', operation);
-      const [attribute_key, attribute_value] = Object.entries(attribute_tuple)[0];
-      expect(Object.entries(bedrockAttributes).length).toBe(1);
-      expect(bedrockAttributes[attribute_key]).toEqual(attribute_value);
+
+      for (const [attribute_key, attribute_value] of Object.entries(attribute_tuple)) {
+        expect(bedrockAttributes[attribute_key]).toEqual(attribute_value);
+      }
+
       const bedrockAgentSuccessAttributes: Attributes = doResponseHookBedrock(services, 'BedrockAgent', operation);
-      expect(Object.entries(bedrockAgentSuccessAttributes).length).toBe(1);
-      expect(bedrockAgentSuccessAttributes[attribute_key]).toEqual(attribute_value);
+
+      for (const [attribute_key, attribute_value] of Object.entries(attribute_tuple)) {
+        expect(bedrockAgentSuccessAttributes[attribute_key]).toEqual(attribute_value);
+      }
     }
   });