Add Bedrock, Bedrock Agent, Bedrock Runtime, and Bedrock Agent Runtime support

aws-distro-opentelemetry-node-autoinstrumentation/package.json

@@ -38,6 +38,10 @@
     "url": "https://github.com/aws-observability/aws-otel-js-instrumentation/issues"
   },
   "devDependencies": {
+    "@aws-sdk/client-bedrock": "3.632.0",
+    "@aws-sdk/client-bedrock-agent": "3.632.0",
+    "@aws-sdk/client-bedrock-agent-runtime": "3.632.0",
+    "@aws-sdk/client-bedrock-runtime": "3.632.0",
     "@aws-sdk/client-kinesis": "3.632.0",
     "@aws-sdk/client-s3": "3.632.0",
     "@aws-sdk/client-sqs": "3.632.0",

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

@@ -31,4 +31,8 @@ export const AWS_ATTRIBUTE_KEYS: { [key: string]: string } = {
   AWS_SQS_QUEUE_NAME: 'aws.sqs.queue.name',
   AWS_KINESIS_STREAM_NAME: 'aws.kinesis.stream.name',
   AWS_DYNAMODB_TABLE_NAMES: SEMATTRS_AWS_DYNAMODB_TABLE_NAMES,
+  AWS_BEDROCK_DATA_SOURCE_ID: 'aws.bedrock.data_source.id',
+  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-distro-opentelemetry-node-autoinstrumentation/src/aws-metric-attribute-generator.ts

@@ -55,6 +55,8 @@ const NORMALIZED_DYNAMO_DB_SERVICE_NAME: string = 'AWS::DynamoDB';
 const NORMALIZED_KINESIS_SERVICE_NAME: string = 'AWS::Kinesis';
 const NORMALIZED_S3_SERVICE_NAME: string = 'AWS::S3';
 const NORMALIZED_SQS_SERVICE_NAME: string = 'AWS::SQS';
+const NORMALIZED_BEDROCK_SERVICE_NAME: string = 'AWS::Bedrock';
+const NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME: string = 'AWS::BedrockRuntime';
 
 const DB_CONNECTION_RESOURCE_TYPE: string = 'DB::Connection';
 // As per https://opentelemetry.io/docs/specs/semconv/resource/#service, if service name is not specified, SDK defaults
@@ -317,10 +319,19 @@ export class AwsMetricAttributeGenerator implements MetricAttributeGenerator {
    * Cloud Control resource format</a> as much as possible, with special attention to services we
    * can detect remote resource information for. Long term, we would like to normalize service name
    * in the upstream.
+   *
+   * For Bedrock, Bedrock Agent, and Bedrock Agent Runtime, we can align with AWS Cloud Control and use
+   * AWS::Bedrock for RemoteService. For BedrockRuntime, we are using AWS::BedrockRuntime
+   * as the associated remote resource (Model) is not listed in Cloud Control.
    */
   private static normalizeRemoteServiceName(span: ReadableSpan, serviceName: string): string {
     if (AwsSpanProcessingUtil.isAwsSDKSpan(span)) {
-      return 'AWS::' + serviceName;
+      const awsSdkServiceMapping: { [key: string]: string } = {
+        BedrockAgent: NORMALIZED_BEDROCK_SERVICE_NAME,
+        BedrockAgentRuntime: NORMALIZED_BEDROCK_SERVICE_NAME,
+        BedrockRuntime: NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME,
+      };
+      return awsSdkServiceMapping[serviceName] || 'AWS::' + serviceName;
     }
     return serviceName;
   }
@@ -369,6 +380,31 @@ export class AwsMetricAttributeGenerator implements MetricAttributeGenerator {
         remoteResourceIdentifier = SqsUrlParser.getQueueName(
           AwsMetricAttributeGenerator.escapeDelimiters(span.attributes[AWS_ATTRIBUTE_KEYS.AWS_SQS_QUEUE_URL])
         );
+      } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID)) {
+        remoteResourceType = NORMALIZED_BEDROCK_SERVICE_NAME + '::Agent';
+        remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
+          span.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID]
+        );
+      } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID)) {
+        remoteResourceType = NORMALIZED_BEDROCK_SERVICE_NAME + '::DataSource';
+        remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
+          span.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID]
+        );
+      } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ID)) {
+        remoteResourceType = NORMALIZED_BEDROCK_SERVICE_NAME + '::Guardrail';
+        remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
+          span.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ID]
+        );
+      } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID)) {
+        remoteResourceType = NORMALIZED_BEDROCK_SERVICE_NAME + '::KnowledgeBase';
+        remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
+          span.attributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]
+        );
+      } else if (AwsSpanProcessingUtil.isKeyPresent(span, AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL)) {
+        remoteResourceType = NORMALIZED_BEDROCK_SERVICE_NAME + '::Model';
+        remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
+          span.attributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL]
+        );
       }
     } else if (AwsSpanProcessingUtil.isDBSpan(span)) {
       remoteResourceType = DB_CONNECTION_RESOURCE_TYPE;

aws-distro-opentelemetry-node-autoinstrumentation/src/aws-span-processing-util.ts

@@ -38,6 +38,10 @@ export class AwsSpanProcessingUtil {
   static MAX_KEYWORD_LENGTH: number = 27;
   static SQL_DIALECT_PATTERN: string = '^(?:' + AwsSpanProcessingUtil.getDialectKeywords().join('|') + ')\\b';
 
+  // TODO: Use Semantic Conventions once upgraded
+  static GEN_AI_REQUEST_MODEL: string = 'gen_ai.request.model';
+  static GEN_AI_SYSTEM: string = 'gen_ai.system';
+
   static getDialectKeywords(): string[] {
     return SQL_DIALECT_KEYWORDS_JSON.keywords;
   }

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

@@ -0,0 +1,221 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { Attributes, DiagLogger, Span, SpanKind, Tracer } from '@opentelemetry/api';
+import {
+  AwsSdkInstrumentationConfig,
+  NormalizedRequest,
+  NormalizedResponse,
+} from '@opentelemetry/instrumentation-aws-sdk';
+import { AWS_ATTRIBUTE_KEYS } from '../../../aws-attribute-keys';
+import { RequestMetadata, ServiceExtension } from '../../../third-party/otel/aws/services/ServiceExtension';
+import { AwsSpanProcessingUtil } from '../../../aws-span-processing-util';
+
+const AGENT_ID: string = 'agentId';
+const KNOWLEDGE_BASE_ID: string = 'knowledgeBaseId';
+const DATA_SOURCE_ID: string = 'dataSourceId';
+const GUARDRAIL_ID: string = 'guardrailId';
+const MODEL_ID: string = 'modelId';
+const AWS_BEDROCK_SYSTEM: string = 'aws_bedrock';
+
+const AGENT_OPERATIONS = [
+  'CreateAgentActionGroup',
+  'CreateAgentAlias',
+  'DeleteAgentActionGroup',
+  'DeleteAgentAlias',
+  'DeleteAgent',
+  'DeleteAgentVersion',
+  'GetAgentActionGroup',
+  'GetAgentAlias',
+  'GetAgent',
+  'GetAgentVersion',
+  'ListAgentActionGroups',
+  'ListAgentAliases',
+  'ListAgentKnowledgeBases',
+  'ListAgentVersions',
+  'PrepareAgent',
+  'UpdateAgentActionGroup',
+  'UpdateAgentAlias',
+  'UpdateAgent',
+];
+
+const KNOWLEDGE_BASE_OPERATIONS = [
+  'AssociateAgentKnowledgeBase',
+  'CreateDataSource',
+  'DeleteKnowledgeBase',
+  'DisassociateAgentKnowledgeBase',
+  'GetAgentKnowledgeBase',
+  'GetKnowledgeBase',
+  'ListDataSources',
+  'UpdateAgentKnowledgeBase',
+];
+
+const DATA_SOURCE_OPERATIONS = ['DeleteDataSource', 'GetDataSource', 'UpdateDataSource'];
+
+// The following constants map the way we present the data in telemetry to how they appear in request/responses
+// e.g. we put `aws.bedrock.knowledge_base.id` into trace data by finding `knowledgeBaseId`
+const agentOperationAttributeKeyMapping = { [AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID]: AGENT_ID };
+const knowledgeBaseOperationAttributeKeyMapping = {
+  [AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID]: KNOWLEDGE_BASE_ID,
+};
+const dataSourceOperationAttributeKeyMapping = {
+  [AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_DATA_SOURCE_ID]: DATA_SOURCE_ID,
+};
+
+// This map allows us to get all relevant attribute key mappings for a given operation
+const operationToBedrockAgentAttributesMap: { [key: string]: { [key: string]: string } } = {};
+for (const operation of AGENT_OPERATIONS) {
+  operationToBedrockAgentAttributesMap[operation] = agentOperationAttributeKeyMapping;
+}
+for (const operation of KNOWLEDGE_BASE_OPERATIONS) {
+  operationToBedrockAgentAttributesMap[operation] = knowledgeBaseOperationAttributeKeyMapping;
+}
+for (const operation of DATA_SOURCE_OPERATIONS) {
+  operationToBedrockAgentAttributesMap[operation] = dataSourceOperationAttributeKeyMapping;
+}
+
+// This class is an extension for <a
+// href="https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Agents_for_Amazon_Bedrock.html">
+// Agents for Amazon Bedrock</a>.
+// This class primarily identify three types of resource based operations: AGENT_OPERATIONS,
+// KNOWLEDGE_BASE_OPERATIONS, and DATA_SOURCE_OPERATIONS. We only support operations that are related to
+// the resource and where the context contains the resource ID.
+export class BedrockAgentServiceExtension implements ServiceExtension {
+  requestPreSpanHook(
+    request: NormalizedRequest,
+    config: AwsSdkInstrumentationConfig,
+    diag: DiagLogger
+  ): RequestMetadata {
+    const spanAttributes: Attributes = {};
+    const isIncoming = false;
+    const spanKind: SpanKind = SpanKind.CLIENT;
+    let spanName: string | undefined;
+
+    const operation: string = request.commandName;
+    if (operation && operationToBedrockAgentAttributesMap[operation]) {
+      const bedrockAgentServiceInfo = operationToBedrockAgentAttributesMap[operation];
+      for (const serviceInfo of Object.entries(bedrockAgentServiceInfo)) {
+        const [attributeKey, requestParamKey] = serviceInfo;
+        const requestParamValue = request.commandInput?.[requestParamKey];
+
+        if (requestParamValue) {
+          spanAttributes[attributeKey] = requestParamValue;
+        }
+      }
+    }
+
+    return {
+      isIncoming,
+      spanAttributes,
+      spanKind,
+      spanName,
+    };
+  }
+
+  responseHook(response: NormalizedResponse, span: Span, tracer: Tracer, config: AwsSdkInstrumentationConfig): void {
+    const operation: string = response.request.commandName;
+    if (operation && operationToBedrockAgentAttributesMap[operation]) {
+      const bedrockAgentServiceInfo = operationToBedrockAgentAttributesMap[operation];
+      for (const serviceInfo of Object.entries(bedrockAgentServiceInfo)) {
+        const [attributeKey, responseParamKey] = serviceInfo;
+        const responseParamValue = response.data[responseParamKey];
+
+        if (responseParamValue) {
+          span.setAttribute(attributeKey, responseParamValue);
+        }
+      }
+    }
+  }
+}
+
+// This class is an extension for <a
+// href="https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Agents_for_Amazon_Bedrock_Runtime.html">
+// Agents for Amazon Bedrock Runtime</a>.
+export class BedrockAgentRuntimeServiceExtension implements ServiceExtension {
+  requestPreSpanHook(
+    request: NormalizedRequest,
+    config: AwsSdkInstrumentationConfig,
+    diag: DiagLogger
+  ): RequestMetadata {
+    const spanAttributes: Attributes = {};
+    const isIncoming = false;
+    const spanKind: SpanKind = SpanKind.CLIENT;
+    let spanName: string | undefined;
+
+    const agentId = request.commandInput?.[AGENT_ID];
+    const knowledgeBaseId = request.commandInput?.[KNOWLEDGE_BASE_ID];
+
+    if (agentId) {
+      spanAttributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_AGENT_ID] = agentId;
+    }
+    if (knowledgeBaseId) {
+      spanAttributes[AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_KNOWLEDGE_BASE_ID] = knowledgeBaseId;
+    }
+
+    return {
+      isIncoming,
+      spanAttributes,
+      spanKind,
+      spanName,
+    };
+  }
+}
+
+// This class is an extension for <a
+// href="https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Amazon_Bedrock.html">Bedrock</a>.
+export class BedrockServiceExtension implements ServiceExtension {
+  // Must be implemented, returning empty metadata
+  requestPreSpanHook(
+    request: NormalizedRequest,
+    config: AwsSdkInstrumentationConfig,
+    diag: DiagLogger
+  ): RequestMetadata {
+    const spanAttributes: Attributes = {};
+    const isIncoming = false;
+    const spanKind: SpanKind = SpanKind.CLIENT;
+    let spanName: string | undefined;
+    return {
+      isIncoming,
+      spanAttributes,
+      spanKind,
+      spanName,
+    };
+  }
+  responseHook(response: NormalizedResponse, span: Span, tracer: Tracer, config: AwsSdkInstrumentationConfig): void {
+    const guardrail_id = response.data[GUARDRAIL_ID];
+
+    if (guardrail_id) {
+      span.setAttribute(AWS_ATTRIBUTE_KEYS.AWS_BEDROCK_GUARDRAIL_ID, guardrail_id);
+    }
+  }
+}
+
+// This class is an extension for <a
+// href="https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Amazon_Bedrock_Runtime.html">
+// Amazon Bedrock Runtime</a>.
+export class BedrockRuntimeServiceExtension implements ServiceExtension {
+  requestPreSpanHook(
+    request: NormalizedRequest,
+    config: AwsSdkInstrumentationConfig,
+    diag: DiagLogger
+  ): RequestMetadata {
+    const spanAttributes: Attributes = {};
+    const isIncoming = false;
+    const spanKind: SpanKind = SpanKind.CLIENT;
+    let spanName: string | undefined;
+
+    const modelId = request.commandInput?.[MODEL_ID];
+
+    spanAttributes[AwsSpanProcessingUtil.GEN_AI_SYSTEM] = AWS_BEDROCK_SYSTEM;
+    if (modelId) {
+      spanAttributes[AwsSpanProcessingUtil.GEN_AI_REQUEST_MODEL] = modelId;
+    }
+
+    return {
+      isIncoming,
+      spanAttributes,
+      spanKind,
+      spanName,
+    };
+  }
+}

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

@@ -18,6 +18,12 @@ import { AWSXRAY_TRACE_ID_HEADER, AWSXRayPropagator } from '@opentelemetry/propa
 import { APIGatewayProxyEventHeaders, Context } from 'aws-lambda';
 import { AWS_ATTRIBUTE_KEYS } from '../aws-attribute-keys';
 import { RequestMetadata } from '../third-party/otel/aws/services/ServiceExtension';
+import {
+  BedrockAgentRuntimeServiceExtension,
+  BedrockAgentServiceExtension,
+  BedrockRuntimeServiceExtension,
+  BedrockServiceExtension,
+} from './aws/services/bedrock';
 import { KinesisServiceExtension } from './aws/services/kinesis';
 import { S3ServiceExtension } from './aws/services/s3';
 
@@ -51,6 +57,10 @@ export function applyInstrumentationPatches(instrumentations: Instrumentation[])
       if (services) {
         services.set('S3', new S3ServiceExtension());
         services.set('Kinesis', new KinesisServiceExtension());
+        services.set('Bedrock', new BedrockServiceExtension());
+        services.set('BedrockAgent', new BedrockAgentServiceExtension());
+        services.set('BedrockAgentRuntime', new BedrockAgentRuntimeServiceExtension());
+        services.set('BedrockRuntime', new BedrockRuntimeServiceExtension());
         patchSqsServiceExtension(services.get('SQS'));
       }
     } else if (instrumentation.instrumentationName === '@opentelemetry/instrumentation-aws-lambda') {