Skip to content

Commit

Permalink
fix: update bedrock guardrail cfn primary id to be arn (#278)
Browse files Browse the repository at this point in the history
### *Summary of Changes:*

- Updating the CFN Primary Id for Guardrail to be the ARN. (Verified
this in the AWS Cloudformation Console)

![Screenshot 2024-10-09 at 12 44
12 PM](https://github.com/user-attachments/assets/1a7437f1-4269-42f0-b4d9-db9f5f37b3a7)

### *Testing Plan:*

Set up sample apps and manually verified correct span attributes were
generated via auto-instrumentation.


![bedrock-guardrail-span-data-verification-python](https://github.com/user-attachments/assets/802b95c9-85d0-404b-b64b-be53f38f3964)

By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.
  • Loading branch information
yiyuan-he authored Oct 11, 2024
1 parent 75af98a commit d723f96
Show file tree
Hide file tree
Showing 3 changed files with 12 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
AWS_BEDROCK_KNOWLEDGE_BASE_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"
AWS_BEDROCK_GUARDRAIL_ARN: str = "aws.bedrock.guardrail.arn"
AWS_SECRETSMANAGER_SECRET_ARN: str = "aws.secretsmanager.secret.arn"
AWS_SNS_TOPIC_ARN: str = "aws.sns.topic.arn"
AWS_STEPFUNCTIONS_STATEMACHINE_ARN: str = "aws.stepfunctions.state_machine.arn"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
from amazon.opentelemetry.distro._aws_attribute_keys import (
AWS_BEDROCK_AGENT_ID,
AWS_BEDROCK_DATA_SOURCE_ID,
AWS_BEDROCK_GUARDRAIL_ARN,
AWS_BEDROCK_GUARDRAIL_ID,
AWS_BEDROCK_KNOWLEDGE_BASE_ID,
AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER,
Expand Down Expand Up @@ -423,6 +424,7 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
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))
cloudformation_primary_identifier = _escape_delimiters(span.attributes.get(AWS_BEDROCK_GUARDRAIL_ARN))
elif is_key_present(span, AWS_BEDROCK_KNOWLEDGE_BASE_ID):
remote_resource_type = _NORMALIZED_BEDROCK_SERVICE_NAME + "::KnowledgeBase"
remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_BEDROCK_KNOWLEDGE_BASE_ID))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
from amazon.opentelemetry.distro._aws_attribute_keys import (
AWS_BEDROCK_AGENT_ID,
AWS_BEDROCK_DATA_SOURCE_ID,
AWS_BEDROCK_GUARDRAIL_ARN,
AWS_BEDROCK_GUARDRAIL_ID,
AWS_BEDROCK_KNOWLEDGE_BASE_ID,
)
Expand All @@ -23,6 +24,7 @@
_KNOWLEDGE_BASE_ID: str = "knowledgeBaseId"
_DATA_SOURCE_ID: str = "dataSourceId"
_GUARDRAIL_ID: str = "guardrailId"
_GUARDRAIL_ARN: str = "guardrailArn"
_MODEL_ID: str = "modelId"
_AWS_BEDROCK_SYSTEM: str = "aws_bedrock"

Expand Down Expand Up @@ -217,6 +219,13 @@ def on_success(self, span: Span, result: _BotoResultT):
guardrail_id,
)

guardrail_arn = result.get(_GUARDRAIL_ARN)
if guardrail_arn:
span.set_attribute(
AWS_BEDROCK_GUARDRAIL_ARN,
guardrail_arn,
)


class _BedrockRuntimeExtension(_AwsSdkExtension):
"""
Expand Down

0 comments on commit d723f96

Please sign in to comment.