Skip to content
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.

Sign up for GitHub

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

Jump to bottom

Add Bedrock, Bedrock Agent, Bedrock Agent Runtime Support. #209

Merged
merged 31 commits into from
Jul 9, 2024
Merged

Add Bedrock, Bedrock Agent, Bedrock Agent Runtime Support. #209

merged 31 commits into from
Jul 9, 2024

Conversation

zzhlogin
Copy link
Contributor

@zzhlogin zzhlogin commented Jun 13, 2024 •
edited
Loading

Add bedrock support as following:

  1. Instrumentation enhancement:

A. Bedrock:
Extract guardrailId from API response, and add into "aws.bedrock.guardrail.id" span attribute.

B. Bedrock Agent:
Extract agentId from both API request and response, and add into "aws.bedrock.agent.id" span attribute.
Extract knowledgeBaseId from API request, and add into "aws.bedrock.knowledgebase.id" span attribute.
Extract dataSourceId from both API request and response,, and add into "aws.bedrock.datasource.id" span attribute.

The instrumentation is on API operation level, we make sure only one attribute is extracted per API call, there will be no overlap/conflict to identify the resouce.

C. Bedrock Agent Runtime:
Extract agentId from API request, and add into "aws.bedrock.agent.id" span attribute.
Extract knowledgeBaseId from API request, and add into "aws.bedrock.knowledgebase.id" span attribute.

We have checked all APIs and make sure there is no API call having both attributes.

D. Bedrock Runtime:
extract the following attributes and add into span according to Gen AI semantic-conventions:

gen_ai.request.model
gen_ai.system
  1. Populate RemoteResourceType and RemoteResourceIdentifier with:
"RemoteResourceType": "AWS::Bedrock::Agent"
"RemoteResourceIdentifier": "<agent_id>"

"RemoteResourceType": "AWS::Bedrock::DataSource"
"RemoteResourceIdentifier": "<datasource_id>"

"RemoteResourceType": "AWS::Bedrock::Guardrail"
"RemoteResourceIdentifier": "<guardrail_id>"

"RemoteResourceType": "AWS::Bedrock::KnowledgeBase"
"RemoteResourceIdentifier": "<knowledgeBase_id>"

"RemoteResourceType": "AWS::Bedrock::Model"
"RemoteResourceIdentifier": "<modelId>"

Testing: E2E test performed to confirm trace, metrics are generated:
Traces:
Screenshot 2024-07-08 at 3 52 21 PM
Screenshot 2024-07-08 at 3 52 54 PM
Screenshot 2024-07-08 at 3 53 23 PM
Screenshot 2024-07-08 at 4 06 29 PM
Screenshot 2024-07-08 at 4 06 44 PM

metrics:
Screenshot 2024-07-08 at 4 17 27 PM
Screenshot 2024-07-08 at 4 16 47 PM
Screenshot 2024-07-08 at 4 18 24 PM
Screenshot 2024-07-08 at 4 18 50 PM

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

@zzhlogin zzhlogin requested a review from a team as a code owner June 13, 2024 21:56
@zzhlogin zzhlogin changed the title Add Bedrock, Bedrock Agent, Bedrock Agent Runtime Support. (NOT READY) Add Bedrock, Bedrock Agent, Bedrock Agent Runtime Support. Jun 14, 2024
@zzhlogin zzhlogin mentioned this pull request Jul 1, 2024
Closed
mxiamxia
mxiamxia reviewed Jul 2, 2024
View reviewed changes
Copy link
Member

@mxiamxia mxiamxia left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM in high level, since I am not very familiar with Python impl atm, will need to get other reviewer's feedback as well.

thpierce
thpierce reviewed Jul 3, 2024
View reviewed changes
aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py Outdated Show resolved Hide resolved
aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py Outdated Show resolved Hide resolved
aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py Outdated Show resolved Hide resolved
aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py Show resolved Hide resolved
aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py Outdated Show resolved Hide resolved
aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py Outdated Show resolved Hide resolved
aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_patches.py Show resolved Hide resolved
aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py Outdated Show resolved Hide resolved
aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py Show resolved Hide resolved
contract-tests/images/applications/botocore/botocore_server.py Outdated Show resolved Hide resolved
thpierce
thpierce reviewed Jul 5, 2024
View reviewed changes
aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py Outdated Show resolved Hide resolved
aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py Outdated Show resolved Hide resolved
aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py Outdated Show resolved Hide resolved
aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_patches.py Outdated Show resolved Hide resolved
aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_patches.py Outdated Show resolved Hide resolved
aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py Outdated Show resolved Hide resolved
aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py Outdated Show resolved Hide resolved
aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py Outdated Show resolved Hide resolved
aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py Show resolved Hide resolved
aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_instrumentation_patch.py Outdated Show resolved Hide resolved
@zzhlogin zzhlogin merged commit 9f4655a into main Jul 9, 2024
11 checks passed
@zzhlogin zzhlogin deleted the enhance_bedrock branch July 9, 2024 15:54
@zzhlogin zzhlogin mentioned this pull request Jul 11, 2024
Merged
@zzhlogin zzhlogin mentioned this pull request Jul 19, 2024
Closed
@mxiamxia mxiamxia mentioned this pull request Jul 26, 2024
Merged
zzhlogin added a commit that referenced this pull request Aug 16, 2024
@zzhlogin
Add bedrock contract tests. (#227)
0d05ca7 
The PR is a follow up on bedrock service support PR:

#209

We add contract tests for following Bedrock services that covers all
resource attributes we newly support:
1. Bedrock API: `GetGuardrail` 
2. BedrockAgent APIs: `GetAgent`, `GetDataSource`, `GetKnowledgeBase`
3. BedrockRuntime API: `InvokeModel`
4. BedrockAgentRuntime API: `InvokeAgent`

Upgrade `botocore` and `boto3` to the latest version `1.34.143` so that
to support Bedrock services API calls.
Upgrade `localstack/localstack` image to the latest version `3.5.0` so
resolve the SQS API call issue using `localstack/localstack:2.0.1` with
new version of `boto3`:
localstack/localstack#9610

**Contract test limitation:**
The contract tests in current repo is using
[LocalStackContainer](https://github.com/aws-observability/aws-otel-python-instrumentation/blob/912dd93ff19b7a4594bb9ed1a7d8cde4907a735d/contract-tests/tests/test/amazon/botocore/botocore_test.py#L68)
to serve AWS SDK service calls. But it doesn’t has bedrock related
service support (This is [the full service
list](https://docs.localstack.cloud/references/coverage/) it support.).
In this case, no matter which bedrock API we call in contract test, the
response will always be 4XX.

As a workaround, we inject `inject_200_success` and `inject_500_error`
directly into the API call to make sure we receive http response with
expected status code and attributes.

**`_assert_semantic_conventions_span_attributes` function change:**
In [_assert_semantic_conventions_span_attributes
function](https://github.com/aws-observability/aws-otel-python-instrumentation/blob/1753bbf2c3cd41778abc358a7f1e3199e48c7fa9/contract-tests/tests/test/amazon/botocore/botocore_test.py#L448)
it is checking the input `service` equals to `rpc.service`, however, we
pass the input service with
["remote_service"](https://github.com/aws-observability/aws-otel-python-instrumentation/blob/1753bbf2c3cd41778abc358a7f1e3199e48c7fa9/contract-tests/tests/test/amazon/botocore/botocore_test.py#L430),
where there is mismatch for example for Bedrock Agent Runtime: we have
`rpc_service="Bedrock Agent Runtime", remote_service="AWS::Bedrock". `
Thus we change to use `rpc_service` if it is provided by:
```
kwargs.get("rpc_service") if "rpc_service" in kwargs else kwargs.get("remote_service").split("::")[-1],
```



By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@mxiamxia mxiamxia mxiamxia left review comments

@thpierce thpierce thpierce approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

None yet
3 participants
@zzhlogin @mxiamxia @thpierce