-
Notifications
You must be signed in to change notification settings - Fork 26
feat: Add Resource and CFN Attributes for Bedrock AgentCore spans #495
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 21 commits
be08bab
98d567a
52744ef
17716ac
25f91ac
9e39da9
28c1fae
dc491e9
82bd892
7dbabb5
e0b5fd6
32f27b2
2fb7895
7968b5b
77031e6
f51f997
6a54c95
37c398e
d29964e
2741494
2d7f5ca
03d6654
f2ac7f7
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -8,14 +8,22 @@ | |
|
||
from amazon.opentelemetry.distro._aws_attribute_keys import ( | ||
AWS_AUTH_ACCESS_KEY, | ||
AWS_AUTH_CREDENTIAL_PROVIDER_ARN, | ||
AWS_AUTH_REGION, | ||
AWS_BEDROCK_AGENT_ID, | ||
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_DATA_SOURCE_ID, | ||
AWS_BEDROCK_GUARDRAIL_ARN, | ||
AWS_BEDROCK_GUARDRAIL_ID, | ||
AWS_BEDROCK_KNOWLEDGE_BASE_ID, | ||
AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER, | ||
AWS_DYNAMODB_TABLE_ARN, | ||
AWS_GATEWAY_TARGET_ID, | ||
AWS_KINESIS_STREAM_ARN, | ||
AWS_KINESIS_STREAM_NAME, | ||
AWS_LAMBDA_FUNCTION_ARN, | ||
|
@@ -63,6 +71,13 @@ | |
SERVICE_METRIC, | ||
MetricAttributeGenerator, | ||
) | ||
from amazon.opentelemetry.distro.patches.semconv._incubating.attributes.gen_ai_attributes 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.regional_resource_arn_parser import RegionalResourceArnParser | ||
from amazon.opentelemetry.distro.sqs_url_parser import SqsUrlParser | ||
from opentelemetry.sdk.resources import Resource | ||
|
@@ -105,6 +120,7 @@ | |
_NORMALIZED_SQS_SERVICE_NAME: str = "AWS::SQS" | ||
_NORMALIZED_BEDROCK_SERVICE_NAME: str = "AWS::Bedrock" | ||
_NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME: str = "AWS::BedrockRuntime" | ||
_NORMALIZED_BEDROCK_AGENTCORE_SERVICE_NAME: str = "AWS::BedrockAgentCore" | ||
_NORMALIZED_SECRETSMANAGER_SERVICE_NAME: str = "AWS::SecretsManager" | ||
_NORMALIZED_SNS_SERVICE_NAME: str = "AWS::SNS" | ||
_NORMALIZED_STEPFUNCTIONS_SERVICE_NAME: str = "AWS::StepFunctions" | ||
|
@@ -118,6 +134,19 @@ | |
# Special DEPENDENCY attribute value if GRAPHQL_OPERATION_TYPE attribute key is present. | ||
_GRAPHQL: str = "graphql" | ||
|
||
# AWS SDK service mapping for normalization | ||
_AWS_SDK_SERVICE_MAPPING = { | ||
"Bedrock Agent": _NORMALIZED_BEDROCK_SERVICE_NAME, | ||
"Bedrock Agent Runtime": _NORMALIZED_BEDROCK_SERVICE_NAME, | ||
"Bedrock Runtime": _NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME, | ||
"Bedrock AgentCore Control": _NORMALIZED_BEDROCK_AGENTCORE_SERVICE_NAME, | ||
"Bedrock AgentCore": _NORMALIZED_BEDROCK_AGENTCORE_SERVICE_NAME, | ||
"Secrets Manager": _NORMALIZED_SECRETSMANAGER_SERVICE_NAME, | ||
"SNS": _NORMALIZED_SNS_SERVICE_NAME, | ||
"SFN": _NORMALIZED_STEPFUNCTIONS_SERVICE_NAME, | ||
"Lambda": _NORMALIZED_LAMBDA_SERVICE_NAME, | ||
} | ||
|
||
_logger: Logger = getLogger(__name__) | ||
|
||
|
||
|
@@ -327,16 +356,6 @@ def _normalize_remote_service_name(span: ReadableSpan, service_name: str) -> str | |
as the associated remote resource (Model) is not listed in Cloud Control. | ||
""" | ||
if is_aws_sdk_span(span): | ||
aws_sdk_service_mapping = { | ||
"Bedrock Agent": _NORMALIZED_BEDROCK_SERVICE_NAME, | ||
"Bedrock Agent Runtime": _NORMALIZED_BEDROCK_SERVICE_NAME, | ||
"Bedrock Runtime": _NORMALIZED_BEDROCK_RUNTIME_SERVICE_NAME, | ||
"Secrets Manager": _NORMALIZED_SECRETSMANAGER_SERVICE_NAME, | ||
"SNS": _NORMALIZED_SNS_SERVICE_NAME, | ||
"SFN": _NORMALIZED_STEPFUNCTIONS_SERVICE_NAME, | ||
"Lambda": _NORMALIZED_LAMBDA_SERVICE_NAME, | ||
} | ||
|
||
# Special handling for Lambda invoke operations | ||
if _is_lambda_invoke_operation(span): | ||
lambda_function_name = span.attributes.get(AWS_LAMBDA_FUNCTION_NAME) | ||
|
@@ -345,7 +364,7 @@ def _normalize_remote_service_name(span: ReadableSpan, service_name: str) -> str | |
# is missing rather than falling back to a generic service name | ||
return lambda_function_name if lambda_function_name else UNKNOWN_REMOTE_SERVICE | ||
|
||
return aws_sdk_service_mapping.get(service_name, "AWS::" + service_name) | ||
return _AWS_SDK_SERVICE_MAPPING.get(service_name, "AWS::" + service_name) | ||
return service_name | ||
|
||
|
||
|
@@ -466,6 +485,16 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri | |
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 ( | ||
_AWS_SDK_SERVICE_MAPPING.get(str(span.attributes.get(_RPC_SERVICE))) | ||
== _NORMALIZED_BEDROCK_AGENTCORE_SERVICE_NAME | ||
): | ||
agentcore_type, agentcore_identifier, agentcore_cfn_id = ( | ||
_get_bedrock_agentcore_resource_type_and_identifier(span) | ||
) | ||
remote_resource_type = agentcore_type | ||
remote_resource_identifier = _escape_delimiters(agentcore_identifier) if agentcore_identifier else None | ||
cloudformation_primary_identifier = _escape_delimiters(agentcore_cfn_id) if agentcore_cfn_id else None | ||
elif is_key_present(span, AWS_SECRETSMANAGER_SECRET_ARN): | ||
remote_resource_type = _NORMALIZED_SECRETSMANAGER_SERVICE_NAME + "::Secret" | ||
remote_resource_identifier = _escape_delimiters( | ||
|
@@ -674,6 +703,149 @@ def _set_span_kind_for_dependency(span: ReadableSpan, attributes: BoundedAttribu | |
attributes[AWS_SPAN_KIND] = span_kind | ||
|
||
|
||
def _get_bedrock_agentcore_resource_type_and_identifier( | ||
span: ReadableSpan, | ||
) -> tuple[Optional[str], Optional[str], Optional[str]]: | ||
"""Get BedrockAgentCore resource type, identifier, and CFN primary identifier based on span attributes.""" | ||
attrs = span.attributes | ||
if not attrs: | ||
return None, None, None | ||
|
||
def format_resource_type(resource_type: Optional[str]) -> Optional[str]: | ||
return f"{_NORMALIZED_BEDROCK_AGENTCORE_SERVICE_NAME}::{resource_type}" if resource_type else None | ||
|
||
for handler in [ | ||
_handle_browser_attrs, | ||
_handle_gateway_attrs, | ||
_handle_runtime_attrs, | ||
_handle_code_interpreter_attrs, | ||
_handle_identity_attrs, | ||
_handle_memory_attrs, | ||
]: | ||
resource_type, resource_identifier, cfn_primary_identifier = handler(attrs) | ||
if resource_type: | ||
return format_resource_type(resource_type), resource_identifier, cfn_primary_identifier | ||
|
||
return None, None, None | ||
|
||
|
||
def _handle_browser_attrs(attrs) -> tuple[Optional[str], Optional[str], Optional[str]]: | ||
browser_id = attrs.get(GEN_AI_BROWSER_ID) | ||
browser_arn = attrs.get(AWS_BEDROCK_AGENTCORE_BROWSER_ARN) | ||
if browser_id or browser_arn: | ||
resource_identifier = None | ||
if browser_arn: | ||
resource_identifier = extract_bedrock_agentcore_resource_id_from_arn(str(browser_arn)) | ||
if browser_id: | ||
resource_identifier = str(browser_id) | ||
resource_type = "Browser" if resource_identifier == "aws.browser.v1" else "BrowserCustom" | ||
return resource_type, resource_identifier, resource_identifier | ||
return None, None, None | ||
|
||
|
||
def _handle_gateway_attrs(attrs) -> tuple[Optional[str], Optional[str], Optional[str]]: | ||
gateway_id = attrs.get(GEN_AI_GATEWAY_ID) | ||
gateway_arn = attrs.get(AWS_BEDROCK_AGENTCORE_GATEWAY_ARN) | ||
gateway_target_id = attrs.get(AWS_GATEWAY_TARGET_ID) | ||
|
||
if gateway_target_id: | ||
resource_identifier = str(gateway_target_id) | ||
if gateway_arn: | ||
resource_identifier = extract_bedrock_agentcore_resource_id_from_arn(str(gateway_arn)) | ||
if gateway_id: | ||
resource_identifier = str(gateway_id) | ||
return "GatewayTarget", resource_identifier, resource_identifier | ||
|
||
if gateway_arn or gateway_id: | ||
resource_identifier = None | ||
if gateway_arn: | ||
resource_identifier = extract_bedrock_agentcore_resource_id_from_arn(str(gateway_arn)) | ||
if gateway_id: | ||
resource_identifier = str(gateway_id) | ||
return "Gateway", resource_identifier, resource_identifier | ||
|
||
|
||
return None, None, None | ||
|
||
|
||
def _handle_runtime_attrs(attrs) -> tuple[Optional[str], Optional[str], Optional[str]]: | ||
runtime_id = attrs.get(GEN_AI_RUNTIME_ID) | ||
runtime_arn = attrs.get(AWS_BEDROCK_AGENTCORE_RUNTIME_ARN) | ||
runtime_endpoint_arn = attrs.get(AWS_BEDROCK_AGENTCORE_RUNTIME_ENDPOINT_ARN) | ||
|
||
if runtime_endpoint_arn: | ||
resource_identifier = extract_bedrock_agentcore_resource_id_from_arn(str(runtime_endpoint_arn)) | ||
return "RuntimeEndpoint", resource_identifier, str(runtime_endpoint_arn) | ||
|
||
if runtime_arn or runtime_id: | ||
resource_identifier = None | ||
if runtime_arn: | ||
resource_identifier = extract_bedrock_agentcore_resource_id_from_arn(str(runtime_arn)) | ||
if runtime_id: | ||
resource_identifier = str(runtime_id) | ||
return "Runtime", resource_identifier, resource_identifier | ||
|
||
return None, None, None | ||
|
||
|
||
def _handle_code_interpreter_attrs(attrs) -> tuple[Optional[str], Optional[str], Optional[str]]: | ||
code_interpreter_id = attrs.get(GEN_AI_CODE_INTERPRETER_ID) | ||
code_interpreter_arn = attrs.get(AWS_BEDROCK_AGENTCORE_CODE_INTERPRETER_ARN) | ||
|
||
if code_interpreter_id or code_interpreter_arn: | ||
resource_identifier = None | ||
if code_interpreter_arn: | ||
resource_identifier = extract_bedrock_agentcore_resource_id_from_arn(str(code_interpreter_arn)) | ||
if code_interpreter_id: | ||
resource_identifier = str(code_interpreter_id) | ||
resource_type = ( | ||
"CodeInterpreter" if resource_identifier == "aws.codeinterpreter.v1" else "CodeInterpreterCustom" | ||
) | ||
return resource_type, resource_identifier, resource_identifier | ||
|
||
return None, None, None | ||
|
||
|
||
def _handle_identity_attrs(attrs) -> tuple[Optional[str], Optional[str], Optional[str]]: | ||
credential_arn = attrs.get(AWS_AUTH_CREDENTIAL_PROVIDER_ARN) | ||
|
||
if credential_arn: | ||
credential_arn_str = str(credential_arn) | ||
resource_type = None | ||
if "apikeycredentialprovider" in credential_arn_str: | ||
resource_type = "APIKeyCredentialProvider" | ||
if "oauth2credentialprovider" in credential_arn_str: | ||
resource_type = "OAuth2CredentialProvider" | ||
resource_identifier = extract_bedrock_agentcore_resource_id_from_arn(credential_arn_str) | ||
return resource_type, resource_identifier, resource_identifier | ||
|
||
return None, None, None | ||
|
||
|
||
def _handle_memory_attrs(attrs) -> tuple[Optional[str], Optional[str], Optional[str]]: | ||
memory_id = attrs.get(GEN_AI_MEMORY_ID) | ||
memory_arn = attrs.get(AWS_BEDROCK_AGENTCORE_MEMORY_ARN) | ||
|
||
if memory_id or memory_arn: | ||
resource_identifier = None | ||
cfn_primary_identifier = None | ||
if memory_arn: | ||
resource_identifier = extract_bedrock_agentcore_resource_id_from_arn(str(memory_arn)) | ||
cfn_primary_identifier = str(memory_arn) | ||
if memory_id: | ||
resource_identifier = str(memory_id) | ||
return "Memory", resource_identifier, cfn_primary_identifier | ||
|
||
return None, None, None | ||
|
||
|
||
def extract_bedrock_agentcore_resource_id_from_arn(arn: str) -> Optional[str]: | ||
"""Extract resource ID from ARN resource part.""" | ||
resource_part = RegionalResourceArnParser.extract_resource_name_from_arn(arn) | ||
if not resource_part: | ||
return None | ||
return resource_part.split("/")[-1] if "/" in resource_part else resource_part | ||
|
||
|
||
def _log_unknown_attribute(attribute_key: str, span: ReadableSpan) -> None: | ||
message: str = "No valid %s value found for %s span %s" | ||
if _logger.isEnabledFor(DEBUG): | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,13 @@ | ||
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. | ||
# SPDX-License-Identifier: Apache-2.0 | ||
""" | ||
Utility module holding attribute keys for incubating Gen AI semantic conventions. | ||
|
||
Remove this once we've contributed them to upstream. | ||
""" | ||
|
||
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" |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Do you always need to check on two if
browser_id
value exists?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Or this
resource_identifier = str(browser_id)
should beagentcore_cfn_id = str(browser_id)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
ack. That makes sense.