diff --git a/CHANGELOG.md b/CHANGELOG.md index d39093671..e0f969c13 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,3 +11,5 @@ For any change that affects end users of this package, please add an entry under If your change does not need a CHANGELOG entry, add the "skip changelog" label to your PR. ## Unreleased +- Add botocore instrumentation extension for Bedrock AgentCore services with span attributes + ([#490](https://github.com/aws-observability/aws-otel-python-instrumentation/pull/490)) 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 d349b8890..2ccbf64b9 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,6 +20,8 @@ # TODO:Move to Semantic Conventions when these attributes are added. AWS_AUTH_ACCESS_KEY: str = "aws.auth.account.access_key" AWS_AUTH_REGION: str = "aws.auth.region" +AWS_AUTH_CREDENTIAL_PROVIDER_ARN: str = "aws.auth.credential_provider.arn" +AWS_GATEWAY_TARGET_ID = "aws.gateway.target.id" AWS_SQS_QUEUE_URL: str = "aws.sqs.queue.url" AWS_SQS_QUEUE_NAME: str = "aws.sqs.queue.name" AWS_KINESIS_STREAM_ARN: str = "aws.kinesis.stream.arn" @@ -41,3 +43,10 @@ AWS_REMOTE_RESOURCE_ACCOUNT_ID: str = "aws.remote.resource.account.id" AWS_REMOTE_RESOURCE_REGION: str = "aws.remote.resource.region" AWS_SERVICE_TYPE: str = "aws.service.type" +AWS_BEDROCK_AGENTCORE_RUNTIME_ARN: str = "aws.bedrock.agentcore.runtime.arn" +AWS_BEDROCK_AGENTCORE_RUNTIME_ENDPOINT_ARN: str = "aws.bedrock.agentcore.runtime_endpoint.arn" +AWS_BEDROCK_AGENTCORE_CODE_INTERPRETER_ARN: str = "aws.bedrock.agentcore.code_interpreter.arn" +AWS_BEDROCK_AGENTCORE_BROWSER_ARN: str = "aws.bedrock.agentcore.browser.arn" +AWS_BEDROCK_AGENTCORE_MEMORY_ARN: str = "aws.bedrock.agentcore.memory.arn" +AWS_BEDROCK_AGENTCORE_GATEWAY_ARN: str = "aws.bedrock.agentcore.gateway.arn" +AWS_BEDROCK_AGENTCORE_WORKLOAD_IDENTITY_ARN: str = "aws.bedrock.agentcore.identity.workload_identity.arn" diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_agentcore_patches.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_agentcore_patches.py new file mode 100644 index 000000000..0e227b2f6 --- /dev/null +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_agentcore_patches.py @@ -0,0 +1,92 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +from typing import Any, Dict + +from amazon.opentelemetry.distro._aws_attribute_keys import ( + AWS_AUTH_CREDENTIAL_PROVIDER_ARN, + AWS_BEDROCK_AGENTCORE_BROWSER_ARN, + AWS_BEDROCK_AGENTCORE_CODE_INTERPRETER_ARN, + AWS_BEDROCK_AGENTCORE_GATEWAY_ARN, + AWS_BEDROCK_AGENTCORE_MEMORY_ARN, + AWS_BEDROCK_AGENTCORE_RUNTIME_ARN, + AWS_BEDROCK_AGENTCORE_RUNTIME_ENDPOINT_ARN, + AWS_BEDROCK_AGENTCORE_WORKLOAD_IDENTITY_ARN, + AWS_GATEWAY_TARGET_ID, +) +from opentelemetry.instrumentation.botocore.extensions.types import ( + _AttributeMapT, + _AwsSdkExtension, + _BotocoreInstrumentorContext, + _BotoResultT, +) +from opentelemetry.trace.span import Span + +GEN_AI_RUNTIME_ID = "gen_ai.runtime.id" +GEN_AI_BROWSER_ID = "gen_ai.browser.id" +GEN_AI_CODE_INTERPRETER_ID = "gen_ai.code_interpreter.id" +GEN_AI_MEMORY_ID = "gen_ai.memory.id" +GEN_AI_GATEWAY_ID = "gen_ai.gateway.id" + +# Mapping of flattened JSON paths to attribute keys +_ATTRIBUTE_MAPPING = { + "agentRuntimeArn": AWS_BEDROCK_AGENTCORE_RUNTIME_ARN, + "agentRuntimeEndpointArn": AWS_BEDROCK_AGENTCORE_RUNTIME_ENDPOINT_ARN, + "agentRuntimeId": GEN_AI_RUNTIME_ID, + "browserArn": AWS_BEDROCK_AGENTCORE_BROWSER_ARN, + "browserId": GEN_AI_BROWSER_ID, + "browserIdentifier": GEN_AI_BROWSER_ID, + "codeInterpreterArn": AWS_BEDROCK_AGENTCORE_CODE_INTERPRETER_ARN, + "codeInterpreterId": GEN_AI_CODE_INTERPRETER_ID, + "codeInterpreterIdentifier": GEN_AI_CODE_INTERPRETER_ID, + "gatewayArn": AWS_BEDROCK_AGENTCORE_GATEWAY_ARN, + "gatewayId": GEN_AI_GATEWAY_ID, + "gatewayIdentifier": GEN_AI_GATEWAY_ID, + "targetId": AWS_GATEWAY_TARGET_ID, + "memory.arn": AWS_BEDROCK_AGENTCORE_MEMORY_ARN, + "memory.id": GEN_AI_MEMORY_ID, + "memoryId": GEN_AI_MEMORY_ID, + "credentialProviderArn": AWS_AUTH_CREDENTIAL_PROVIDER_ARN, + "workloadIdentityArn": AWS_BEDROCK_AGENTCORE_WORKLOAD_IDENTITY_ARN, + "workloadIdentityDetails.workloadIdentityArn": AWS_BEDROCK_AGENTCORE_WORKLOAD_IDENTITY_ARN, +} + + +class _BedrockAgentCoreExtension(_AwsSdkExtension): + def extract_attributes(self, attributes: _AttributeMapT): + extracted_attrs = self._extract_attributes(self._call_context.params) + attributes.update(extracted_attrs) + + def on_success( + self, + span: Span, + result: _BotoResultT, + instrumentor_context: _BotocoreInstrumentorContext, + ): + if span is None or not span.is_recording(): + return + + extracted_attrs = self._extract_attributes(result) + for attr_name, attr_value in extracted_attrs.items(): + span.set_attribute(attr_name, attr_value) + + @staticmethod + def _extract_attributes(params: Dict[str, Any]): + """Extracts all Bedrock AgentCore attributes using mapping-based traversal""" + attrs = {} + for path, attr_key in _ATTRIBUTE_MAPPING.items(): + value = _BedrockAgentCoreExtension._get_nested_value(params, path) + if value: + attrs[attr_key] = value + return attrs + + @staticmethod + def _get_nested_value(data: Dict[str, Any], path: str): + """Get value from nested dictionary using dot notation path""" + keys = path.split(".") + value = data + for key in keys: + if isinstance(value, dict) and key in value: + value = value[key] + else: + return None + return value diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py index 10fc77182..900ce688f 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py @@ -23,6 +23,9 @@ AWS_STEPFUNCTIONS_ACTIVITY_ARN, AWS_STEPFUNCTIONS_STATEMACHINE_ARN, ) +from amazon.opentelemetry.distro.patches._bedrock_agentcore_patches import ( # noqa # pylint: disable=unused-import + _BedrockAgentCoreExtension, +) from amazon.opentelemetry.distro.patches._bedrock_patches import ( # noqa # pylint: disable=unused-import _BedrockAgentExtension, _BedrockAgentRuntimeExtension, @@ -247,6 +250,10 @@ def _apply_botocore_bedrock_patch() -> None: # pylint: disable=too-many-stateme _KNOWN_EXTENSIONS["bedrock"] = _lazy_load(".", "_BedrockExtension") _KNOWN_EXTENSIONS["bedrock-agent"] = _lazy_load(".", "_BedrockAgentExtension") _KNOWN_EXTENSIONS["bedrock-agent-runtime"] = _lazy_load(".", "_BedrockAgentRuntimeExtension") + _KNOWN_EXTENSIONS["bedrock-agentcore"] = _lazy_load(".._bedrock_agentcore_patches", "_BedrockAgentCoreExtension") + _KNOWN_EXTENSIONS["bedrock-agentcore-control"] = _lazy_load( + ".._bedrock_agentcore_patches", "_BedrockAgentCoreExtension" + ) # TODO: The following code is to patch bedrock-runtime bugs that are fixed in # opentelemetry-instrumentation-botocore==0.56b0 in these PRs: diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_instrumentation_patch.py similarity index 89% rename from aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py rename to aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_instrumentation_patch.py index 256ee3673..979f79f7b 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/patches/test_instrumentation_patch.py @@ -10,6 +10,24 @@ import opentelemetry.sdk.extension.aws.resource.ec2 as ec2_resource import opentelemetry.sdk.extension.aws.resource.eks as eks_resource +from amazon.opentelemetry.distro._aws_attribute_keys import ( + AWS_AUTH_CREDENTIAL_PROVIDER_ARN, + AWS_BEDROCK_AGENTCORE_BROWSER_ARN, + AWS_BEDROCK_AGENTCORE_CODE_INTERPRETER_ARN, + AWS_BEDROCK_AGENTCORE_GATEWAY_ARN, + AWS_BEDROCK_AGENTCORE_MEMORY_ARN, + AWS_BEDROCK_AGENTCORE_RUNTIME_ARN, + AWS_BEDROCK_AGENTCORE_RUNTIME_ENDPOINT_ARN, + AWS_BEDROCK_AGENTCORE_WORKLOAD_IDENTITY_ARN, + AWS_GATEWAY_TARGET_ID, +) +from amazon.opentelemetry.distro.patches._bedrock_agentcore_patches import ( + GEN_AI_BROWSER_ID, + GEN_AI_CODE_INTERPRETER_ID, + GEN_AI_GATEWAY_ID, + GEN_AI_MEMORY_ID, + GEN_AI_RUNTIME_ID, +) from amazon.opentelemetry.distro.patches._instrumentation_patch import ( AWS_GEVENT_PATCH_MODULES, apply_instrumentation_patches, @@ -38,6 +56,24 @@ _LAMBDA_FUNCTION_NAME: str = "lambdaFunctionName" _LAMBDA_SOURCE_MAPPING_ID: str = "lambdaEventSourceMappingID" _TABLE_ARN: str = "arn:aws:dynamodb:us-west-2:123456789012:table/testTable" +_AGENTCORE_RUNTIME_ARN: str = "arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/test-runtime-123" +_AGENTCORE_RUNTIME_ENDPOINT_ARN: str = "arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime-endpoint/test-endpoint" +_AGENTCORE_RUNTIME_ID: str = "test-runtime-123" +_AGENTCORE_BROWSER_ARN: str = "arn:aws:bedrock-agentcore:us-east-1:123456789012:browser/testBrowser-1234567890" +_AGENTCORE_BROWSER_ID: str = "testBrowser-1234567890" +_AGENTCORE_CODE_INTERPRETER_ARN: str = ( + "arn:aws:bedrock-agentcore:us-east-1:123456789012:code-interpreter/testCodeInt-1234567890" +) +_AGENTCORE_CODE_INTERPRETER_ID: str = "testCodeInt-1234567890" +_AGENTCORE_GATEWAY_ARN: str = "arn:aws:bedrock-agentcore:us-east-1:123456789012:gateway/agentGateway-123456789" +_AGENTCORE_GATEWAY_ID: str = "agentGateway-123456789" +_AGENTCORE_TARGET_ID: str = "target-123456789" +_AGENTCORE_MEMORY_ARN: str = "arn:aws:bedrock-agentcore:us-east-1:123456789012:memory/agentMemory-123456789" +_AGENTCORE_MEMORY_ID: str = "agentMemory-123456789" +_AGENTCORE_CREDENTIAL_PROVIDER_ARN: str = ( + "arn:aws:acps:us-east-1:123456789012:token-vault/test-vault/apikeycredentialprovider/test-provider" +) +_AGENTCORE_WORKLOAD_IDENTITY_ARN: str = "arn:aws:bedrock-agentcore:us-east-1:123456789012:workload-identity/test-wi" # Patch names IMPORTLIB_METADATA_VERSION_PATCH: str = "amazon.opentelemetry.distro._utils.version" @@ -160,6 +196,14 @@ def _test_unpatched_botocore_instrumentation(self): "bedrock-agent-runtime" in _KNOWN_EXTENSIONS, "Upstream has added a Bedrock Agent Runtime extension" ) + # Bedrock AgentCore + self.assertFalse("bedrock-agentcore" in _KNOWN_EXTENSIONS, "Upstream has added a Bedrock AgentCore extension") + + # Bedrock AgentCore Control + self.assertFalse( + "bedrock-agentcore-control" in _KNOWN_EXTENSIONS, "Upstream has added a Bedrock AgentCore Control extension" + ) + # BedrockRuntime self.assertTrue("bedrock-runtime" in _KNOWN_EXTENSIONS, "Upstream has added a bedrock-runtime extension") @@ -237,6 +281,41 @@ def _test_patched_botocore_instrumentation(self): bedrock_agent_runtime_sucess_attributes: Dict[str, str] = _do_on_success_bedrock("bedrock-agent-runtime") self.assertEqual(len(bedrock_agent_runtime_sucess_attributes), 0) + # Bedrock AgentCore + self.assertTrue("bedrock-agentcore" in _KNOWN_EXTENSIONS) + bedrock_agentcore_attributes: Dict[str, str] = _do_extract_bedrock_agentcore_attributes() + # Runtime attributes + self.assertEqual(bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_RUNTIME_ARN], _AGENTCORE_RUNTIME_ARN) + self.assertEqual( + bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_RUNTIME_ENDPOINT_ARN], _AGENTCORE_RUNTIME_ENDPOINT_ARN + ) + self.assertEqual(bedrock_agentcore_attributes[GEN_AI_RUNTIME_ID], _AGENTCORE_RUNTIME_ID) + # Browser attributes + self.assertEqual(bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_BROWSER_ARN], _AGENTCORE_BROWSER_ARN) + self.assertEqual(bedrock_agentcore_attributes[GEN_AI_BROWSER_ID], _AGENTCORE_BROWSER_ID) + # Code interpreter attributes + self.assertEqual( + bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_CODE_INTERPRETER_ARN], _AGENTCORE_CODE_INTERPRETER_ARN + ) + self.assertEqual(bedrock_agentcore_attributes[GEN_AI_CODE_INTERPRETER_ID], _AGENTCORE_CODE_INTERPRETER_ID) + # Gateway attributes + self.assertEqual(bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_GATEWAY_ARN], _AGENTCORE_GATEWAY_ARN) + self.assertEqual(bedrock_agentcore_attributes[GEN_AI_GATEWAY_ID], _AGENTCORE_GATEWAY_ID) + self.assertEqual(bedrock_agentcore_attributes[AWS_GATEWAY_TARGET_ID], _AGENTCORE_TARGET_ID) + # Memory attributes + self.assertEqual(bedrock_agentcore_attributes[GEN_AI_MEMORY_ID], _AGENTCORE_MEMORY_ID) + self.assertEqual(bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_MEMORY_ARN], _AGENTCORE_MEMORY_ARN) + # Auth and identity attributes + self.assertEqual( + bedrock_agentcore_attributes[AWS_AUTH_CREDENTIAL_PROVIDER_ARN], _AGENTCORE_CREDENTIAL_PROVIDER_ARN + ) + self.assertEqual( + bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_WORKLOAD_IDENTITY_ARN], _AGENTCORE_WORKLOAD_IDENTITY_ARN + ) + + # Bedrock AgentCore Control + self.assertTrue("bedrock-agentcore-control" in _KNOWN_EXTENSIONS) + # BedrockRuntime self.assertTrue("bedrock-runtime" in _KNOWN_EXTENSIONS) @@ -854,6 +933,31 @@ def _do_extract_lambda_attributes() -> Dict[str, str]: return _do_extract_attributes(service_name, params) +def _do_extract_bedrock_agentcore_attributes() -> Dict[str, str]: + service_name: str = "bedrock-agentcore" + params: Dict[str, Any] = { + "agentRuntimeArn": _AGENTCORE_RUNTIME_ARN, + "agentRuntimeEndpointArn": _AGENTCORE_RUNTIME_ENDPOINT_ARN, + "agentRuntimeId": _AGENTCORE_RUNTIME_ID, + "browserArn": _AGENTCORE_BROWSER_ARN, + "browserId": _AGENTCORE_BROWSER_ID, + "browserIdentifier": _AGENTCORE_BROWSER_ID, + "codeInterpreterArn": _AGENTCORE_CODE_INTERPRETER_ARN, + "codeInterpreterId": _AGENTCORE_CODE_INTERPRETER_ID, + "codeInterpreterIdentifier": _AGENTCORE_CODE_INTERPRETER_ID, + "gatewayArn": _AGENTCORE_GATEWAY_ARN, + "gatewayId": _AGENTCORE_GATEWAY_ID, + "gatewayIdentifier": _AGENTCORE_GATEWAY_ID, + "targetId": _AGENTCORE_TARGET_ID, + "memoryId": _AGENTCORE_MEMORY_ID, + "credentialProviderArn": _AGENTCORE_CREDENTIAL_PROVIDER_ARN, + "workloadIdentityArn": _AGENTCORE_WORKLOAD_IDENTITY_ARN, + "memory": {"arn": _AGENTCORE_MEMORY_ARN, "id": _AGENTCORE_MEMORY_ID}, + "workloadIdentityDetails": {"workloadIdentityArn": _AGENTCORE_WORKLOAD_IDENTITY_ARN}, + } + 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