|
10 | 10 |
|
11 | 11 | import opentelemetry.sdk.extension.aws.resource.ec2 as ec2_resource |
12 | 12 | import opentelemetry.sdk.extension.aws.resource.eks as eks_resource |
| 13 | +from amazon.opentelemetry.distro._aws_attribute_keys import ( |
| 14 | + AWS_AUTH_CREDENTIAL_PROVIDER_ARN, |
| 15 | + AWS_BEDROCK_AGENTCORE_BROWSER_ARN, |
| 16 | + AWS_BEDROCK_AGENTCORE_CODE_INTERPRETER_ARN, |
| 17 | + AWS_BEDROCK_AGENTCORE_GATEWAY_ARN, |
| 18 | + AWS_BEDROCK_AGENTCORE_MEMORY_ARN, |
| 19 | + AWS_BEDROCK_AGENTCORE_RUNTIME_ARN, |
| 20 | + AWS_BEDROCK_AGENTCORE_RUNTIME_ENDPOINT_ARN, |
| 21 | + AWS_BEDROCK_AGENTCORE_WORKLOAD_IDENTITY_ARN, |
| 22 | + AWS_GATEWAY_TARGET_ID, |
| 23 | +) |
| 24 | +from amazon.opentelemetry.distro.patches._bedrock_agentcore_patches import ( |
| 25 | + GEN_AI_BROWSER_ID, |
| 26 | + GEN_AI_CODE_INTERPRETER_ID, |
| 27 | + GEN_AI_GATEWAY_ID, |
| 28 | + GEN_AI_MEMORY_ID, |
| 29 | + GEN_AI_RUNTIME_ID, |
| 30 | +) |
13 | 31 | from amazon.opentelemetry.distro.patches._instrumentation_patch import ( |
14 | 32 | AWS_GEVENT_PATCH_MODULES, |
15 | 33 | apply_instrumentation_patches, |
|
38 | 56 | _LAMBDA_FUNCTION_NAME: str = "lambdaFunctionName" |
39 | 57 | _LAMBDA_SOURCE_MAPPING_ID: str = "lambdaEventSourceMappingID" |
40 | 58 | _TABLE_ARN: str = "arn:aws:dynamodb:us-west-2:123456789012:table/testTable" |
| 59 | +_AGENTCORE_RUNTIME_ARN: str = "arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/test-runtime-123" |
| 60 | +_AGENTCORE_RUNTIME_ENDPOINT_ARN: str = "arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime-endpoint/test-endpoint" |
| 61 | +_AGENTCORE_RUNTIME_ID: str = "test-runtime-123" |
| 62 | +_AGENTCORE_BROWSER_ARN: str = "arn:aws:bedrock-agentcore:us-east-1:123456789012:browser/testBrowser-1234567890" |
| 63 | +_AGENTCORE_BROWSER_ID: str = "testBrowser-1234567890" |
| 64 | +_AGENTCORE_CODE_INTERPRETER_ARN: str = ( |
| 65 | + "arn:aws:bedrock-agentcore:us-east-1:123456789012:code-interpreter/testCodeInt-1234567890" |
| 66 | +) |
| 67 | +_AGENTCORE_CODE_INTERPRETER_ID: str = "testCodeInt-1234567890" |
| 68 | +_AGENTCORE_GATEWAY_ARN: str = "arn:aws:bedrock-agentcore:us-east-1:123456789012:gateway/agentGateway-123456789" |
| 69 | +_AGENTCORE_GATEWAY_ID: str = "agentGateway-123456789" |
| 70 | +_AGENTCORE_TARGET_ID: str = "target-123456789" |
| 71 | +_AGENTCORE_MEMORY_ARN: str = "arn:aws:bedrock-agentcore:us-east-1:123456789012:memory/agentMemory-123456789" |
| 72 | +_AGENTCORE_MEMORY_ID: str = "agentMemory-123456789" |
| 73 | +_AGENTCORE_CREDENTIAL_PROVIDER_ARN: str = ( |
| 74 | + "arn:aws:acps:us-east-1:123456789012:token-vault/test-vault/apikeycredentialprovider/test-provider" |
| 75 | +) |
| 76 | +_AGENTCORE_WORKLOAD_IDENTITY_ARN: str = "arn:aws:bedrock-agentcore:us-east-1:123456789012:workload-identity/test-wi" |
41 | 77 |
|
42 | 78 | # Patch names |
43 | 79 | IMPORTLIB_METADATA_VERSION_PATCH: str = "amazon.opentelemetry.distro._utils.version" |
@@ -160,6 +196,14 @@ def _test_unpatched_botocore_instrumentation(self): |
160 | 196 | "bedrock-agent-runtime" in _KNOWN_EXTENSIONS, "Upstream has added a Bedrock Agent Runtime extension" |
161 | 197 | ) |
162 | 198 |
|
| 199 | + # Bedrock AgentCore |
| 200 | + self.assertFalse("bedrock-agentcore" in _KNOWN_EXTENSIONS, "Upstream has added a Bedrock AgentCore extension") |
| 201 | + |
| 202 | + # Bedrock AgentCore Control |
| 203 | + self.assertFalse( |
| 204 | + "bedrock-agentcore-control" in _KNOWN_EXTENSIONS, "Upstream has added a Bedrock AgentCore Control extension" |
| 205 | + ) |
| 206 | + |
163 | 207 | # BedrockRuntime |
164 | 208 | self.assertTrue("bedrock-runtime" in _KNOWN_EXTENSIONS, "Upstream has added a bedrock-runtime extension") |
165 | 209 |
|
@@ -237,6 +281,41 @@ def _test_patched_botocore_instrumentation(self): |
237 | 281 | bedrock_agent_runtime_sucess_attributes: Dict[str, str] = _do_on_success_bedrock("bedrock-agent-runtime") |
238 | 282 | self.assertEqual(len(bedrock_agent_runtime_sucess_attributes), 0) |
239 | 283 |
|
| 284 | + # Bedrock AgentCore |
| 285 | + self.assertTrue("bedrock-agentcore" in _KNOWN_EXTENSIONS) |
| 286 | + bedrock_agentcore_attributes: Dict[str, str] = _do_extract_bedrock_agentcore_attributes() |
| 287 | + # Runtime attributes |
| 288 | + self.assertEqual(bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_RUNTIME_ARN], _AGENTCORE_RUNTIME_ARN) |
| 289 | + self.assertEqual( |
| 290 | + bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_RUNTIME_ENDPOINT_ARN], _AGENTCORE_RUNTIME_ENDPOINT_ARN |
| 291 | + ) |
| 292 | + self.assertEqual(bedrock_agentcore_attributes[GEN_AI_RUNTIME_ID], _AGENTCORE_RUNTIME_ID) |
| 293 | + # Browser attributes |
| 294 | + self.assertEqual(bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_BROWSER_ARN], _AGENTCORE_BROWSER_ARN) |
| 295 | + self.assertEqual(bedrock_agentcore_attributes[GEN_AI_BROWSER_ID], _AGENTCORE_BROWSER_ID) |
| 296 | + # Code interpreter attributes |
| 297 | + self.assertEqual( |
| 298 | + bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_CODE_INTERPRETER_ARN], _AGENTCORE_CODE_INTERPRETER_ARN |
| 299 | + ) |
| 300 | + self.assertEqual(bedrock_agentcore_attributes[GEN_AI_CODE_INTERPRETER_ID], _AGENTCORE_CODE_INTERPRETER_ID) |
| 301 | + # Gateway attributes |
| 302 | + self.assertEqual(bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_GATEWAY_ARN], _AGENTCORE_GATEWAY_ARN) |
| 303 | + self.assertEqual(bedrock_agentcore_attributes[GEN_AI_GATEWAY_ID], _AGENTCORE_GATEWAY_ID) |
| 304 | + self.assertEqual(bedrock_agentcore_attributes[AWS_GATEWAY_TARGET_ID], _AGENTCORE_TARGET_ID) |
| 305 | + # Memory attributes |
| 306 | + self.assertEqual(bedrock_agentcore_attributes[GEN_AI_MEMORY_ID], _AGENTCORE_MEMORY_ID) |
| 307 | + self.assertEqual(bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_MEMORY_ARN], _AGENTCORE_MEMORY_ARN) |
| 308 | + # Auth and identity attributes |
| 309 | + self.assertEqual( |
| 310 | + bedrock_agentcore_attributes[AWS_AUTH_CREDENTIAL_PROVIDER_ARN], _AGENTCORE_CREDENTIAL_PROVIDER_ARN |
| 311 | + ) |
| 312 | + self.assertEqual( |
| 313 | + bedrock_agentcore_attributes[AWS_BEDROCK_AGENTCORE_WORKLOAD_IDENTITY_ARN], _AGENTCORE_WORKLOAD_IDENTITY_ARN |
| 314 | + ) |
| 315 | + |
| 316 | + # Bedrock AgentCore Control |
| 317 | + self.assertTrue("bedrock-agentcore-control" in _KNOWN_EXTENSIONS) |
| 318 | + |
240 | 319 | # BedrockRuntime |
241 | 320 | self.assertTrue("bedrock-runtime" in _KNOWN_EXTENSIONS) |
242 | 321 |
|
@@ -854,6 +933,31 @@ def _do_extract_lambda_attributes() -> Dict[str, str]: |
854 | 933 | return _do_extract_attributes(service_name, params) |
855 | 934 |
|
856 | 935 |
|
| 936 | +def _do_extract_bedrock_agentcore_attributes() -> Dict[str, str]: |
| 937 | + service_name: str = "bedrock-agentcore" |
| 938 | + params: Dict[str, Any] = { |
| 939 | + "agentRuntimeArn": _AGENTCORE_RUNTIME_ARN, |
| 940 | + "agentRuntimeEndpointArn": _AGENTCORE_RUNTIME_ENDPOINT_ARN, |
| 941 | + "agentRuntimeId": _AGENTCORE_RUNTIME_ID, |
| 942 | + "browserArn": _AGENTCORE_BROWSER_ARN, |
| 943 | + "browserId": _AGENTCORE_BROWSER_ID, |
| 944 | + "browserIdentifier": _AGENTCORE_BROWSER_ID, |
| 945 | + "codeInterpreterArn": _AGENTCORE_CODE_INTERPRETER_ARN, |
| 946 | + "codeInterpreterId": _AGENTCORE_CODE_INTERPRETER_ID, |
| 947 | + "codeInterpreterIdentifier": _AGENTCORE_CODE_INTERPRETER_ID, |
| 948 | + "gatewayArn": _AGENTCORE_GATEWAY_ARN, |
| 949 | + "gatewayId": _AGENTCORE_GATEWAY_ID, |
| 950 | + "gatewayIdentifier": _AGENTCORE_GATEWAY_ID, |
| 951 | + "targetId": _AGENTCORE_TARGET_ID, |
| 952 | + "memoryId": _AGENTCORE_MEMORY_ID, |
| 953 | + "credentialProviderArn": _AGENTCORE_CREDENTIAL_PROVIDER_ARN, |
| 954 | + "workloadIdentityArn": _AGENTCORE_WORKLOAD_IDENTITY_ARN, |
| 955 | + "memory": {"arn": _AGENTCORE_MEMORY_ARN, "id": _AGENTCORE_MEMORY_ID}, |
| 956 | + "workloadIdentityDetails": {"workloadIdentityArn": _AGENTCORE_WORKLOAD_IDENTITY_ARN}, |
| 957 | + } |
| 958 | + return _do_extract_attributes(service_name, params) |
| 959 | + |
| 960 | + |
857 | 961 | def _do_extract_attributes(service_name: str, params: Dict[str, Any], operation: str = None) -> Dict[str, str]: |
858 | 962 | mock_call_context: MagicMock = MagicMock() |
859 | 963 | mock_call_context.params = params |
|
0 commit comments