diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py index 9b6328fa3..3c3801dc0 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py @@ -20,5 +20,4 @@ AWS_BEDROCK_KNOWLEDGEBASE_ID: str = "aws.bedrock.knowledgebase_id" AWS_BEDROCK_AGENT_ID: str = "aws.bedrock.agent_id" AWS_BEDROCK_MODEL_ID: str = "aws.bedrock.model_id" -AWS_BEDROCK_GAURDRAIL_ID: str = "aws.bedrock.guardrail_id" - +AWS_BEDROCK_GUARDRAIL_ID: str = "aws.bedrock.guardrail_id" diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py index 0b76d9bab..027e63881 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py @@ -8,9 +8,8 @@ from amazon.opentelemetry.distro._aws_attribute_keys import ( AWS_BEDROCK_AGENT_ID, AWS_BEDROCK_DATASOURCE_ID, - AWS_BEDROCK_GAURDRAIL_ID, + AWS_BEDROCK_GUARDRAIL_ID, AWS_BEDROCK_KNOWLEDGEBASE_ID, - AWS_BEDROCK_MODEL_ID, AWS_LOCAL_OPERATION, AWS_LOCAL_SERVICE, AWS_QUEUE_NAME, @@ -297,6 +296,8 @@ def _normalize_remote_service_name(span: ReadableSpan, service_name: str) -> str """ if is_aws_sdk_span(span): aws_sdk_service_mapping = { + "Bedrock Agent": _NORMALIZED_BEDROCK_SERVICE_NAME, + "Bedrock Agent Runtime": _NORMALIZED_BEDROCK_SERVICE_NAME, } return aws_sdk_service_mapping.get(service_name, "AWS::" + service_name) return service_name @@ -380,18 +381,15 @@ 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_MODEL_ID): - remote_resource_type = _NORMALIZED_BEDROCK_SERVICE_NAME + "::Model" - remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_BEDROCK_MODEL_ID)) 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_GAURDRAIL_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_GAURDRAIL_ID)) + 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)) diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_patches.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_patches.py index a0c782711..26fbce61a 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_patches.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_patches.py @@ -3,7 +3,7 @@ # Modifications Copyright The OpenTelemetry Authors. Licensed under the Apache License 2.0 License. import abc import inspect -from typing import Any, Dict, Optional +from typing import Dict, Optional from opentelemetry.instrumentation.botocore.extensions.types import ( _AttributeMapT, @@ -31,14 +31,13 @@ class _AgentOperation(_BedrockAgentOperation): # UpdateAgentKnowledgeBase -> KnowledgeBaseId # GetAgentKnowledgeBase -> KnowledgeBaseId start_attributes = { - "aws.bedrock.agent_id": "AgentId", + "aws.bedrock.agent_id": "agentId", } response_attributes = { - "aws.bedrock.agent_id": "AgentId", + "aws.bedrock.agent_id": "agentId", } @classmethod - @abc.abstractmethod def operation_names(cls): return [ "CreateAgentActionGroup", @@ -71,7 +70,7 @@ class _KnowledgeBaseOperation(_BedrockAgentOperation): # ListIngestionJobs -> not support # StartIngestionJob -> not support start_attributes = { - "aws.bedrock.knowledgebase_id": "KnowledgeBaseId", + "aws.bedrock.knowledgebase_id": "knowledgeBaseId", } response_attributes = {} @@ -95,10 +94,10 @@ class _DataSourceOperation(_BedrockAgentOperation): # ListIngestionJobs -> not support # StartIngestionJob -> not support start_attributes = { - "aws.bedrock.datasource_id": "DataSourceId", + "aws.bedrock.datasource_id": "dataSourceId", } response_attributes = { - "aws.bedrock.datasource_id": "DataSourceId", + "aws.bedrock.datasource_id": "dataSourceId", } @classmethod @@ -108,7 +107,7 @@ def operation_names(cls): _OPERATION_MAPPING = { op_name: op_class - for op_class in [_KnowledgeBaseOperation, _DataSourceOperation] + 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) } @@ -136,34 +135,28 @@ def on_success(self, span: Span, result: _BotoResultT): if response_value: span.set_attribute( key, - value, + response_value, ) class _BedrockAgentRuntimeExtension(_AwsSdkExtension): # -> AgentId, KnowledgebaseId -> no overlap def extract_attributes(self, attributes: _AttributeMapT): # AgentId, KnowledgebaseId - agent_id = self._call_context.params.get("AgentId") + agent_id = self._call_context.params.get("agentId") if agent_id: attributes["aws.bedrock.agent_id"] = agent_id - knowledgebase_id = self._call_context.params.get("KnowledgeBaseId") + knowledgebase_id = self._call_context.params.get("knowledgeBaseId") if knowledgebase_id: attributes["aws.bedrock.knowledgebase_id"] = knowledgebase_id class _BedrockExtension(_AwsSdkExtension): # -> ModelId, GaurdrailId -> no overlap - def extract_attributes(self, attributes: _AttributeMapT): - # ModelId - model_id = self._call_context.params.get("ModelId") - if model_id: - attributes["aws.bedrock.model_id"] = model_id - def on_success(self, span: Span, result: _BotoResultT): # GuardrailId - gaurdrail_id = result.get("guardrailId") - if gaurdrail_id: + guardrail_id = result.get("guardrailId") + if guardrail_id: span.set_attribute( "aws.bedrock.guardrail_id", - gaurdrail_id, + guardrail_id, ) diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py index 072e6eeb0..a2f29ec95 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py @@ -8,6 +8,10 @@ from unittest.mock import MagicMock 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_CONSUMER_PARENT_SPAN_KIND, AWS_LOCAL_OPERATION, AWS_LOCAL_SERVICE, @@ -821,6 +825,9 @@ def test_normalize_remote_service_name_aws_sdk(self): self.validate_aws_sdk_service_normalization("Kinesis", "AWS::Kinesis") self.validate_aws_sdk_service_normalization("S3", "AWS::S3") self.validate_aws_sdk_service_normalization("SQS", "AWS::SQS") + self.validate_aws_sdk_service_normalization("Bedrock", "AWS::Bedrock") + self.validate_aws_sdk_service_normalization("Bedrock Agent", "AWS::Bedrock") + self.validate_aws_sdk_service_normalization("Bedrock Agent Runtime", "AWS::Bedrock") def validate_aws_sdk_service_normalization(self, service_name: str, expected_remote_service: str): self._mock_attribute([SpanAttributes.RPC_SYSTEM, SpanAttributes.RPC_SERVICE], ["aws-api", service_name]) @@ -977,6 +984,26 @@ def test_sdk_client_span_with_remote_resource_attributes(self): self._validate_remote_resource_attributes("AWS::DynamoDB::Table", "aws_table^^name") self._mock_attribute([SpanAttributes.AWS_DYNAMODB_TABLE_NAMES], [None]) + # Validate behaviour of AWS_BEDROCK_AGENT_ID attribute with special chars(^), then remove it. + self._mock_attribute([AWS_BEDROCK_AGENT_ID], ["test_agent_id"], keys, values) + self._validate_remote_resource_attributes("AWS::Bedrock::Agent", "test_agent_id") + self._mock_attribute([AWS_BEDROCK_AGENT_ID], [None]) + + # Validate behaviour of AWS_BEDROCK_DATASOURCE_ID attribute with special chars(^), then remove it. + self._mock_attribute([AWS_BEDROCK_DATASOURCE_ID], ["test_datasource_id"], keys, values) + self._validate_remote_resource_attributes("AWS::Bedrock::DataSource", "test_datasource_id") + self._mock_attribute([AWS_BEDROCK_DATASOURCE_ID], [None]) + + # Validate behaviour of AWS_BEDROCK_GUARDRAIL_ID attribute with special chars(^), then remove it. + self._mock_attribute([AWS_BEDROCK_GUARDRAIL_ID], ["test_guardrail_id"], keys, values) + self._validate_remote_resource_attributes("AWS::Bedrock::Guardrail", "test_guardrail_id") + self._mock_attribute([AWS_BEDROCK_GUARDRAIL_ID], [None]) + + # Validate behaviour of AWS_BEDROCK_KNOWLEDGEBASE_ID attribute with special chars(^), then remove it. + self._mock_attribute([AWS_BEDROCK_KNOWLEDGEBASE_ID], ["test_knowledgeBase_id"], keys, values) + self._validate_remote_resource_attributes("AWS::Bedrock::KnowledgeBase", "test_knowledgeBase_id") + self._mock_attribute([AWS_BEDROCK_KNOWLEDGEBASE_ID], [None]) + self._mock_attribute([SpanAttributes.RPC_SYSTEM], [None]) def test_client_db_span_with_remote_resource_attributes(self): diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py index bc6e851a9..18ca1112c 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py @@ -1,6 +1,6 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -from typing import Dict +from typing import Any, Dict from unittest import TestCase from unittest.mock import MagicMock, patch @@ -9,11 +9,17 @@ from amazon.opentelemetry.distro.patches._instrumentation_patch import apply_instrumentation_patches from opentelemetry.instrumentation.botocore.extensions import _KNOWN_EXTENSIONS from opentelemetry.semconv.trace import SpanAttributes +from opentelemetry.trace.span import Span _STREAM_NAME: str = "streamName" _BUCKET_NAME: str = "bucketName" _QUEUE_NAME: str = "queueName" _QUEUE_URL: str = "queueUrl" +_BEDROCK_MODEL_ID: str = "modelId" +_BEDROCK_AGENT_ID: str = "agentId" +_BEDROCK_DATASOURCE_ID: str = "DataSourceId" +_BEDROCK_GUARDRAIL_ID: str = "GuardrailId" +_BEDROCK_KNOWLEDGEBASE_ID: str = "KnowledgeBaseId" class TestInstrumentationPatch(TestCase): @@ -74,6 +80,17 @@ def _validate_unpatched_botocore_instrumentation(self): self.assertFalse("aws.sqs.queue_url" in attributes) self.assertFalse("aws.sqs.queue_name" in attributes) + # Bedrock + self.assertFalse("bedrock" in _KNOWN_EXTENSIONS, "Upstream has added a Bedrock extension") + + # Bedrock Agent + self.assertFalse("bedrock-agent" in _KNOWN_EXTENSIONS, "Upstream has added a Bedrock Agent extension") + + # Bedrock Agent Runtime + self.assertFalse( + "bedrock-agent-runtime" in _KNOWN_EXTENSIONS, "Upstream has added a Bedrock Agent Runtime extension" + ) + def _validate_patched_botocore_instrumentation(self): # Kinesis self.assertTrue("kinesis" in _KNOWN_EXTENSIONS) @@ -96,6 +113,43 @@ def _validate_patched_botocore_instrumentation(self): self.assertTrue("aws.sqs.queue_name" in sqs_attributes) self.assertEqual(sqs_attributes["aws.sqs.queue_name"], _QUEUE_NAME) + # Bedrock + bedrock_sucess_attributes: Dict[str, str] = _do_bedrock_on_success() + self.assertTrue("aws.bedrock.guardrail_id" in bedrock_sucess_attributes) + self.assertEqual(bedrock_sucess_attributes["aws.bedrock.guardrail_id"], _BEDROCK_GUARDRAIL_ID) + + # Bedrock Agent Operation + self.assertTrue("bedrock-agent" in _KNOWN_EXTENSIONS) + bedrock_agent_op_attributes: Dict[str, str] = _do_extract_bedrock_agent_op_attributes() + self.assertTrue("aws.bedrock.agent_id" in bedrock_agent_op_attributes) + self.assertEqual(bedrock_agent_op_attributes["aws.bedrock.agent_id"], _BEDROCK_AGENT_ID) + bedrock_agent_op_sucess_attributes: Dict[str, str] = _do_bedrock_agent_op_on_success() + self.assertTrue("aws.bedrock.agent_id" in bedrock_agent_op_sucess_attributes) + self.assertEqual(bedrock_agent_op_sucess_attributes["aws.bedrock.agent_id"], _BEDROCK_AGENT_ID) + + # Bedrock DataSource Operation + self.assertTrue("bedrock-agent" in _KNOWN_EXTENSIONS) + bedrock_datasource_op_attributes: Dict[str, str] = _do_extract_bedrock_datasource_op_attributes() + self.assertTrue("aws.bedrock.datasource_id" in bedrock_datasource_op_attributes) + self.assertEqual(bedrock_datasource_op_attributes["aws.bedrock.datasource_id"], _BEDROCK_DATASOURCE_ID) + bedrock_datasource_op_sucess_attributes: Dict[str, str] = _do_bedrock_datasource_op_on_success() + self.assertTrue("aws.bedrock.datasource_id" in bedrock_datasource_op_sucess_attributes) + self.assertEqual(bedrock_datasource_op_sucess_attributes["aws.bedrock.datasource_id"], _BEDROCK_DATASOURCE_ID) + + # Bedrock KnowledgeBase Operation + self.assertTrue("bedrock-agent" in _KNOWN_EXTENSIONS) + bedrock_knowledgebase_op_attributes: Dict[str, str] = _do_extract_bedrock_knowledgebase_op_attributes() + self.assertTrue("aws.bedrock.knowledgebase_id" in bedrock_knowledgebase_op_attributes) + self.assertEqual(bedrock_knowledgebase_op_attributes["aws.bedrock.knowledgebase_id"], _BEDROCK_KNOWLEDGEBASE_ID) + + # Bedrock Agent Runtime + self.assertTrue("bedrock-agent-runtime" in _KNOWN_EXTENSIONS) + bedrock_agent_runtime_attributes: Dict[str, str] = _do_extract_bedrock_agent_runtime_attributes() + self.assertTrue("aws.bedrock.agent_id" in bedrock_agent_runtime_attributes) + self.assertEqual(bedrock_agent_runtime_attributes["aws.bedrock.agent_id"], _BEDROCK_AGENT_ID) + self.assertTrue("aws.bedrock.knowledgebase_id" in bedrock_agent_runtime_attributes) + self.assertEqual(bedrock_agent_runtime_attributes["aws.bedrock.knowledgebase_id"], _BEDROCK_KNOWLEDGEBASE_ID) + def _do_extract_kinesis_attributes() -> Dict[str, str]: service_name: str = "kinesis" @@ -115,10 +169,83 @@ def _do_extract_sqs_attributes() -> Dict[str, str]: return _do_extract_attributes(service_name, params) -def _do_extract_attributes(service_name: str, params: Dict[str, str]) -> Dict[str, str]: +def _do_bedrock_on_success() -> Dict[str, str]: + service_name: str = "bedrock" + result: Dict[str, Any] = {"guardrailId": _BEDROCK_GUARDRAIL_ID} + return _do_on_success(service_name, result) + + +def _do_extract_bedrock_agent_op_attributes() -> Dict[str, str]: + service_name: str = "bedrock-agent" + params: Dict[str, str] = {"agentId": _BEDROCK_AGENT_ID} + operation: str = "CreateAgentAlias" + return _do_extract_attributes(service_name, params, operation) + + +def _do_bedrock_agent_op_on_success() -> Dict[str, str]: + service_name: str = "bedrock-agent" + result: Dict[str, Any] = {"agentId": _BEDROCK_AGENT_ID} + operation: str = "CreateAgentAlias" + return _do_on_success(service_name, result, operation) + + +def _do_extract_bedrock_datasource_op_attributes() -> Dict[str, str]: + service_name: str = "bedrock-agent" + params: Dict[str, str] = {"dataSourceId": _BEDROCK_DATASOURCE_ID} + operation: str = "UpdateDataSource" + return _do_extract_attributes(service_name, params, operation) + + +def _do_bedrock_datasource_op_on_success() -> Dict[str, str]: + service_name: str = "bedrock-agent" + result: Dict[str, Any] = {"dataSourceId": _BEDROCK_DATASOURCE_ID} + operation: str = "UpdateDataSource" + return _do_on_success(service_name, result, operation) + + +def _do_extract_bedrock_knowledgebase_op_attributes() -> Dict[str, str]: + service_name: str = "bedrock-agent" + params: Dict[str, str] = {"knowledgeBaseId": _BEDROCK_KNOWLEDGEBASE_ID} + operation: str = "GetKnowledgeBase" + return _do_extract_attributes(service_name, params, operation) + + +def _do_extract_bedrock_agent_runtime_attributes() -> Dict[str, str]: + service_name: str = "bedrock-agent-runtime" + params: Dict[str, str] = {"agentId": _BEDROCK_AGENT_ID, "knowledgeBaseId": _BEDROCK_KNOWLEDGEBASE_ID} + return _do_extract_attributes(service_name, params) + + +def _do_extract_attributes(service_name: str, params: Dict[str, Any], operation: str = None) -> Dict[str, str]: mock_call_context: MagicMock = MagicMock() mock_call_context.params = params + if operation: + mock_call_context.operation = operation attributes: Dict[str, str] = {} sqs_extension = _KNOWN_EXTENSIONS[service_name]()(mock_call_context) sqs_extension.extract_attributes(attributes) return attributes + + +def _do_on_success( + service_name: str, result: Dict[str, Any], operation: str = None, params: Dict[str, Any] = None +) -> Dict[str, str]: + span_mock: Span = MagicMock() + mock_call_context = MagicMock() + span_attributes: Dict[str, str] = {} + + def set_side_effect(set_key, set_value): + span_attributes[set_key] = set_value + + span_mock.set_attribute.side_effect = set_side_effect + + if operation: + mock_call_context.operation = operation + + if params: + mock_call_context.params = params + + extension = _KNOWN_EXTENSIONS[service_name]()(mock_call_context) + extension.on_success(span_mock, result) + + return span_attributes