Add Bedrock, Bedrock Agent, Bedrock Agent Runtime Support.

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py

@@ -17,3 +17,8 @@
 AWS_QUEUE_URL: str = "aws.sqs.queue_url"
 AWS_QUEUE_NAME: str = "aws.sqs.queue_name"
 AWS_STREAM_NAME: str = "aws.kinesis.stream_name"
+AWS_BEDROCK_DATASOURCE_ID: str = "aws.bedrock.data_source.id"
+AWS_BEDROCK_KNOWLEDGEBASE_ID: str = "aws.bedrock.knowledge_base.id"
+AWS_BEDROCK_AGENT_ID: str = "aws.bedrock.agent.id"
+AWS_BEDROCK_GUARDRAIL_ID: str = "aws.bedrock.guardrail.id"
+GEN_AI_REQUEST_MODEL: str = "gen_ai.request.model"

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py

@@ -6,6 +6,10 @@
 from urllib.parse import ParseResult, urlparse
 
 from amazon.opentelemetry.distro._aws_attribute_keys import (
+    AWS_BEDROCK_AGENT_ID,
+    AWS_BEDROCK_DATASOURCE_ID,
+    AWS_BEDROCK_GUARDRAIL_ID,
+    AWS_BEDROCK_KNOWLEDGEBASE_ID,
     AWS_LOCAL_OPERATION,
     AWS_LOCAL_SERVICE,
     AWS_QUEUE_NAME,
@@ -17,6 +21,7 @@
     AWS_REMOTE_SERVICE,
     AWS_SPAN_KIND,
     AWS_STREAM_NAME,
+    GEN_AI_REQUEST_MODEL,
 )
 from amazon.opentelemetry.distro._aws_span_processing_util import (
     LOCAL_ROOT,
@@ -80,6 +85,8 @@
 _NORMALIZED_KINESIS_SERVICE_NAME: str = "AWS::Kinesis"
 _NORMALIZED_S3_SERVICE_NAME: str = "AWS::S3"
 _NORMALIZED_SQS_SERVICE_NAME: str = "AWS::SQS"
+_NORMALIZED_BEDROCK_SERVICE_NAME: str = "AWS::Bedrock"
+_NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME: str = "AWS::BedrockRuntime"
 _DB_CONNECTION_STRING_TYPE: str = "DB::Connection"
 
 # Special DEPENDENCY attribute value if GRAPHQL_OPERATION_TYPE attribute key is present.
@@ -291,9 +298,18 @@ def _normalize_remote_service_name(span: ReadableSpan, service_name: str) -> str
     If the span is an AWS SDK span, normalize the name to align with <a
     href="https://docs.aws.amazon.com/cloudcontrolapi/latest/userguide/supported-resources.html">AWS Cloud Control
     resource format</a> as much as possible. Long term, we would like to normalize service name in the upstream.
+
+    For Bedrock, Bedrock Agent, Bedrock Agent Runtime, we follow the cloudformation template,
+    use RemoteService as AWS::Bedrock. For BedrockRuntime, there is no corresponding reference in Cloudformation found,
+    so we use AWS::BedrockRuntime.
     """
     if is_aws_sdk_span(span):
-        return "AWS::" + service_name
+        aws_sdk_service_mapping = {
+            "Bedrock Agent": _NORMALIZED_BEDROCK_SERVICE_NAME,
+            "Bedrock Agent Runtime": _NORMALIZED_BEDROCK_SERVICE_NAME,
+            "Bedrock Runtime": _NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME,
+        }
+        return aws_sdk_service_mapping.get(service_name, "AWS::" + service_name)
     return service_name
 
 
@@ -342,6 +358,7 @@ def _generate_remote_operation(span: ReadableSpan) -> str:
     return remote_operation
 
 
+# pylint: disable=too-many-branches
 def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttributes) -> None:
     """
     Remote resource attributes {@link AwsAttributeKeys#AWS_REMOTE_RESOURCE_TYPE} and {@link
@@ -375,6 +392,21 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
             remote_resource_identifier = _escape_delimiters(
                 SqsUrlParser.get_queue_name(span.attributes.get(AWS_QUEUE_URL))
             )
+        elif is_key_present(span, AWS_BEDROCK_AGENT_ID):
+            remote_resource_type = _NORMALIZED_BEDROCK_SERVICE_NAME + "::Agent"
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_BEDROCK_AGENT_ID))
+        elif is_key_present(span, AWS_BEDROCK_DATASOURCE_ID):
+            remote_resource_type = _NORMALIZED_BEDROCK_SERVICE_NAME + "::DataSource"
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_BEDROCK_DATASOURCE_ID))
+        elif is_key_present(span, AWS_BEDROCK_GUARDRAIL_ID):
+            remote_resource_type = _NORMALIZED_BEDROCK_SERVICE_NAME + "::Guardrail"
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_BEDROCK_GUARDRAIL_ID))
+        elif is_key_present(span, AWS_BEDROCK_KNOWLEDGEBASE_ID):
+            remote_resource_type = _NORMALIZED_BEDROCK_SERVICE_NAME + "::KnowledgeBase"
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_BEDROCK_KNOWLEDGEBASE_ID))
+        elif is_key_present(span, GEN_AI_REQUEST_MODEL):
+            remote_resource_type = _NORMALIZED_BEDROCK_SERVICE_NAME + "::Model"
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(GEN_AI_REQUEST_MODEL))
     elif is_db_span(span):
         remote_resource_type = _DB_CONNECTION_STRING_TYPE
         remote_resource_identifier = _get_db_connection(span)

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_patches.py

@@ -0,0 +1,202 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+# Modifications Copyright The OpenTelemetry Authors. Licensed under the Apache License 2.0 License.
+import abc
+import inspect
+from typing import Dict, Optional
+
+from opentelemetry.instrumentation.botocore.extensions.types import (
+    _AttributeMapT,
+    _AwsSdkCallContext,
+    _AwsSdkExtension,
+    _BotoResultT,
+)
+from opentelemetry.trace.span import Span
+
+
+class _BedrockAgentOperation(abc.ABC):
+    request_attributes: Optional[Dict[str, str]] = None
+    response_attributes: Optional[Dict[str, str]] = None
+
+    @classmethod
+    @abc.abstractmethod
+    def operation_names(cls):
+        pass
+
+
+class _AgentOperation(_BedrockAgentOperation):
+    """
+    This class primarily supports BedrockAgent Agent related operations.
+    """
+
+    request_attributes = {
+        "aws.bedrock.agent.id": "agentId",
+    }
+    response_attributes = {
+        "aws.bedrock.agent.id": "agentId",
+    }
+
+    @classmethod
+    def operation_names(cls):
+        return [
+            "CreateAgentActionGroup",
+            "CreateAgentAlias",
+            "DeleteAgentActionGroup",
+            "DeleteAgentAlias",
+            "DeleteAgent",
+            "DeleteAgentVersion",
+            "GetAgentActionGroup",
+            "GetAgentAlias",
+            "GetAgent",
+            "GetAgentVersion",
+            "ListAgentActionGroups",
+            "ListAgentAliases",
+            "ListAgentKnowledgeBases",
+            "ListAgentVersions",
+            "PrepareAgent",
+            "UpdateAgentActionGroup",
+            "UpdateAgentAlias",
+            "UpdateAgent",
+        ]
+
+
+class _KnowledgeBaseOperation(_BedrockAgentOperation):
+    """
+    This class primarily supports BedrockAgent KnowledgeBase related operations.
+
+    Note: The 'CreateDataSource' operation does not have a 'dataSourceId' in the context,
+    but it always comes with a 'knowledgeBaseId'. Therefore, we categorize it under 'knowledgeBaseId' operations.
+    """
+
+    request_attributes = {
+        "aws.bedrock.knowledge_base.id": "knowledgeBaseId",
+    }
+    response_attributes = {
+        "aws.bedrock.knowledge_base.id": "knowledgeBaseId",
+    }
+
+    @classmethod
+    def operation_names(cls):
+        return [
+            "AssociateAgentKnowledgeBase",
+            "CreateDataSource",
+            "DeleteKnowledgeBase",
+            "DisassociateAgentKnowledgeBase",
+            "GetAgentKnowledgeBase",
+            "GetKnowledgeBase",
+            "ListDataSources",
+            "UpdateAgentKnowledgeBase",
+        ]
+
+
+class _DataSourceOperation(_BedrockAgentOperation):
+    """
+    This class primarily supports BedrockAgent DataSource related operations.
+    """
+
+    request_attributes = {
+        "aws.bedrock.data_source.id": "dataSourceId",
+    }
+    response_attributes = {
+        "aws.bedrock.data_source.id": "dataSourceId",
+    }
+
+    @classmethod
+    def operation_names(cls):
+        return ["DeleteDataSource", "GetDataSource", "UpdateDataSource"]
+
+
+# _OPERATION_NAME_TO_ClASS_MAPPING maps operation names to their corresponding classes
+# by iterating over all subclasses of _BedrockAgentOperation and extract operation
+# by call operation_names() function.
+_OPERATION_NAME_TO_ClASS_MAPPING = {
+    op_name: op_class
+    for op_class in [_KnowledgeBaseOperation, _DataSourceOperation, _AgentOperation]
+    for op_name in op_class.operation_names()
+    if inspect.isclass(op_class) and issubclass(op_class, _BedrockAgentOperation) and not inspect.isabstract(op_class)
+}
+
+
+class _BedrockAgentExtension(_AwsSdkExtension):
+    """
+    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: _AgentOperation, _KnowledgeBaseOperation,
+    and _DataSourceOperation. We only support operations that are related to the resource
+    and where the context contains the resource ID.
+    """
+
+    def __init__(self, call_context: _AwsSdkCallContext):
+        super().__init__(call_context)
+        self._operation_class = _OPERATION_NAME_TO_ClASS_MAPPING.get(call_context.operation)
+
+    def extract_attributes(self, attributes: _AttributeMapT):
+        if self._operation_class is None:
+            return
+        for attribute_key, request_param_key in self._operation_class.request_attributes.items():
+            request_param_value = self._call_context.params.get(request_param_key)
+            if request_param_value:
+                attributes[attribute_key] = request_param_value
+
+    def on_success(self, span: Span, result: _BotoResultT):
+        if self._operation_class is None:
+            return
+
+        for attribute_key, response_param_key in self._operation_class.response_attributes.items():
+            response_param_value = result.get(response_param_key)
+            if response_param_value:
+                span.set_attribute(
+                    attribute_key,
+                    response_param_value,
+                )
+
+
+class _BedrockAgentRuntimeExtension(_AwsSdkExtension):
+    """
+    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>.
+    """
+
+    def extract_attributes(self, attributes: _AttributeMapT):
+        agent_id = self._call_context.params.get("agentId")
+        if agent_id:
+            attributes["aws.bedrock.agent.id"] = agent_id
+
+        knowledge_base_id = self._call_context.params.get("knowledgeBaseId")
+        if knowledge_base_id:
+            attributes["aws.bedrock.knowledge_base.id"] = knowledge_base_id
+
+
+class _BedrockExtension(_AwsSdkExtension):
+    """
+    This class is an extension for <a
+    href="https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Amazon_Bedrock.html">Bedrock</a>.
+    """
+
+    # pylint: disable=no-self-use
+    def on_success(self, span: Span, result: _BotoResultT):
+        # guardrailId can only be retrieved from the response, not from the request
+        guardrail_id = result.get("guardrailId")
+        if guardrail_id:
+            span.set_attribute(
+                "aws.bedrock.guardrail.id",
+                guardrail_id,
+            )
+
+
+class _BedrockRuntimeExtension(_AwsSdkExtension):
+    """
+    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>.
+    """
+
+    def extract_attributes(self, attributes: _AttributeMapT):
+        attributes["gen_ai.system"] = "aws_bedrock"
+
+        model_id = self._call_context.params.get("modelId")
+        if model_id:
+            attributes["gen_ai.request.model"] = model_id

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py

@@ -3,6 +3,12 @@
 # Modifications Copyright The OpenTelemetry Authors. Licensed under the Apache License 2.0 License.
 import importlib
 
+from amazon.opentelemetry.distro.patches._bedrock_patches import (  # noqa # pylint: disable=unused-import
+    _BedrockAgentExtension,
+    _BedrockAgentRuntimeExtension,
+    _BedrockExtension,
+    _BedrockRuntimeExtension,
+)
 from opentelemetry.instrumentation.botocore.extensions import _KNOWN_EXTENSIONS
 from opentelemetry.instrumentation.botocore.extensions.sqs import _SqsExtension
 from opentelemetry.instrumentation.botocore.extensions.types import _AttributeMapT, _AwsSdkExtension
@@ -17,6 +23,7 @@ def _apply_botocore_instrumentation_patches() -> None:
     _apply_botocore_kinesis_patch()
     _apply_botocore_s3_patch()
     _apply_botocore_sqs_patch()
+    _apply_botocore_bedrock_patch()
 
 
 def _apply_botocore_kinesis_patch() -> None:
@@ -65,6 +72,19 @@ def patch_extract_attributes(self, attributes: _AttributeMapT):
     _SqsExtension.extract_attributes = patch_extract_attributes
 
 
+def _apply_botocore_bedrock_patch() -> None:
+    """Botocore instrumentation patch for Bedrock, Bedrock Agent, Bedrock Runtime and Bedrock Agent Runtime
+
+    This patch adds an extension to the upstream's list of known extension for Bedrock.
+    Extensions allow for custom logic for adding service-specific information to spans, such as attributes.
+    Specifically, we are adding logic to add the AWS_BEDROCK attributes referenced in _aws_attribute_keys.
+    """
+    _KNOWN_EXTENSIONS["bedrock"] = _lazy_load(".", "_BedrockExtension")
+    _KNOWN_EXTENSIONS["bedrock-agent"] = _lazy_load(".", "_BedrockAgentExtension")
+    _KNOWN_EXTENSIONS["bedrock-agent-runtime"] = _lazy_load(".", "_BedrockAgentRuntimeExtension")
+    _KNOWN_EXTENSIONS["bedrock-runtime"] = _lazy_load(".", "_BedrockRuntimeExtension")
+
+
 # The OpenTelemetry Authors code
 def _lazy_load(module, cls):
     """Clone of upstream opentelemetry.instrumentation.botocore.extensions.lazy_load