botocore: add basic handling for bedrock invoke.model (#3200)

* Add basic handling for invoke.model

* Add changelog a please pylint

* Record converse cassettes against us-east-1

* Avoid double copy of streaming body

---------

Co-authored-by: Adrian Cole <64215+codefromthecrypt@users.noreply.github.com>
This commit is contained in:
Riccardo Magliocchetti 2025-01-23 17:14:49 +01:00 committed by GitHub
parent ec3c51dcd1
commit 2756c1edff
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
12 changed files with 664 additions and 157 deletions

View File

@ -43,6 +43,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
([#3186](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3186))
- `opentelemetry-opentelemetry-botocore` Add basic support for GenAI attributes for AWS Bedrock Converse API
([#3161](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3161))
- `opentelemetry-opentelemetry-botocore` Add basic support for GenAI attributes for AWS Bedrock InvokeModel API
([#3200](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/3200))
### Fixed

View File

@ -0,0 +1,25 @@
import json
import os
import boto3
def main():
client = boto3.client("bedrock-runtime")
response = client.invoke_model(
modelId=os.getenv("CHAT_MODEL", "amazon.titan-text-lite-v1"),
body=json.dumps(
{
"inputText": "Write a short poem on OpenTelemetry.",
"textGenerationConfig": {},
},
),
)
body = response["body"].read()
response_data = json.loads(body.decode("utf-8"))
print(response_data["results"][0]["outputText"])
if __name__ == "__main__":
main()

View File

@ -18,9 +18,13 @@
from __future__ import annotations
import io
import json
import logging
from typing import Any
from botocore.response import StreamingBody
from opentelemetry.instrumentation.botocore.extensions.types import (
_AttributeMapT,
_AwsSdkExtension,
@ -58,7 +62,7 @@ class _BedrockRuntimeExtension(_AwsSdkExtension):
Amazon Bedrock Runtime</a>.
"""
_HANDLED_OPERATIONS = {"Converse"}
_HANDLED_OPERATIONS = {"Converse", "InvokeModel"}
def extract_attributes(self, attributes: _AttributeMapT):
if self._call_context.operation not in self._HANDLED_OPERATIONS:
@ -73,6 +77,7 @@ class _BedrockRuntimeExtension(_AwsSdkExtension):
GenAiOperationNameValues.CHAT.value
)
# Converse
if inference_config := self._call_context.params.get(
"inferenceConfig"
):
@ -97,6 +102,84 @@ class _BedrockRuntimeExtension(_AwsSdkExtension):
inference_config.get("stopSequences"),
)
# InvokeModel
# Get the request body if it exists
body = self._call_context.params.get("body")
if body:
try:
request_body = json.loads(body)
if "amazon.titan" in model_id:
# titan interface is a text completion one
attributes[GEN_AI_OPERATION_NAME] = (
GenAiOperationNameValues.TEXT_COMPLETION.value
)
self._extract_titan_attributes(
attributes, request_body
)
elif "amazon.nova" in model_id:
self._extract_nova_attributes(attributes, request_body)
elif "anthropic.claude" in model_id:
self._extract_claude_attributes(
attributes, request_body
)
except json.JSONDecodeError:
_logger.debug("Error: Unable to parse the body as JSON")
def _extract_titan_attributes(self, attributes, request_body):
config = request_body.get("textGenerationConfig", {})
self._set_if_not_none(
attributes, GEN_AI_REQUEST_TEMPERATURE, config.get("temperature")
)
self._set_if_not_none(
attributes, GEN_AI_REQUEST_TOP_P, config.get("topP")
)
self._set_if_not_none(
attributes, GEN_AI_REQUEST_MAX_TOKENS, config.get("maxTokenCount")
)
self._set_if_not_none(
attributes,
GEN_AI_REQUEST_STOP_SEQUENCES,
config.get("stopSequences"),
)
def _extract_nova_attributes(self, attributes, request_body):
config = request_body.get("inferenceConfig", {})
self._set_if_not_none(
attributes, GEN_AI_REQUEST_TEMPERATURE, config.get("temperature")
)
self._set_if_not_none(
attributes, GEN_AI_REQUEST_TOP_P, config.get("topP")
)
self._set_if_not_none(
attributes, GEN_AI_REQUEST_MAX_TOKENS, config.get("max_new_tokens")
)
self._set_if_not_none(
attributes,
GEN_AI_REQUEST_STOP_SEQUENCES,
config.get("stopSequences"),
)
def _extract_claude_attributes(self, attributes, request_body):
self._set_if_not_none(
attributes,
GEN_AI_REQUEST_MAX_TOKENS,
request_body.get("max_tokens"),
)
self._set_if_not_none(
attributes,
GEN_AI_REQUEST_TEMPERATURE,
request_body.get("temperature"),
)
self._set_if_not_none(
attributes, GEN_AI_REQUEST_TOP_P, request_body.get("top_p")
)
self._set_if_not_none(
attributes,
GEN_AI_REQUEST_STOP_SEQUENCES,
request_body.get("stop_sequences"),
)
@staticmethod
def _set_if_not_none(attributes, key, value):
if value is not None:
@ -115,13 +198,8 @@ class _BedrockRuntimeExtension(_AwsSdkExtension):
if operation_name and request_model:
span.update_name(f"{operation_name} {request_model}")
def on_success(self, span: Span, result: dict[str, Any]):
if self._call_context.operation not in self._HANDLED_OPERATIONS:
return
if not span.is_recording():
return
# pylint: disable=no-self-use
def _converse_on_success(self, span: Span, result: dict[str, Any]):
if usage := result.get("usage"):
if input_tokens := usage.get("inputTokens"):
span.set_attribute(
@ -140,6 +218,109 @@ class _BedrockRuntimeExtension(_AwsSdkExtension):
[stop_reason],
)
def _invoke_model_on_success(
self, span: Span, result: dict[str, Any], model_id: str
):
original_body = None
try:
original_body = result["body"]
body_content = original_body.read()
# Replenish stream for downstream application use
new_stream = io.BytesIO(body_content)
result["body"] = StreamingBody(new_stream, len(body_content))
response_body = json.loads(body_content.decode("utf-8"))
if "amazon.titan" in model_id:
self._handle_amazon_titan_response(span, response_body)
elif "amazon.nova" in model_id:
self._handle_amazon_nova_response(span, response_body)
elif "anthropic.claude" in model_id:
self._handle_anthropic_claude_response(span, response_body)
except json.JSONDecodeError:
_logger.debug("Error: Unable to parse the response body as JSON")
except Exception as exc: # pylint: disable=broad-exception-caught
_logger.debug("Error processing response: %s", exc)
finally:
if original_body is not None:
original_body.close()
def on_success(self, span: Span, result: dict[str, Any]):
if self._call_context.operation not in self._HANDLED_OPERATIONS:
return
if not span.is_recording():
return
# Converse
self._converse_on_success(span, result)
model_id = self._call_context.params.get(_MODEL_ID_KEY)
if not model_id:
return
# InvokeModel
if "body" in result and isinstance(result["body"], StreamingBody):
self._invoke_model_on_success(span, result, model_id)
# pylint: disable=no-self-use
def _handle_amazon_titan_response(
self, span: Span, response_body: dict[str, Any]
):
if "inputTextTokenCount" in response_body:
span.set_attribute(
GEN_AI_USAGE_INPUT_TOKENS, response_body["inputTextTokenCount"]
)
if "results" in response_body and response_body["results"]:
result = response_body["results"][0]
if "tokenCount" in result:
span.set_attribute(
GEN_AI_USAGE_OUTPUT_TOKENS, result["tokenCount"]
)
if "completionReason" in result:
span.set_attribute(
GEN_AI_RESPONSE_FINISH_REASONS,
[result["completionReason"]],
)
# pylint: disable=no-self-use
def _handle_amazon_nova_response(
self, span: Span, response_body: dict[str, Any]
):
if "usage" in response_body:
usage = response_body["usage"]
if "inputTokens" in usage:
span.set_attribute(
GEN_AI_USAGE_INPUT_TOKENS, usage["inputTokens"]
)
if "outputTokens" in usage:
span.set_attribute(
GEN_AI_USAGE_OUTPUT_TOKENS, usage["outputTokens"]
)
if "stopReason" in response_body:
span.set_attribute(
GEN_AI_RESPONSE_FINISH_REASONS, [response_body["stopReason"]]
)
# pylint: disable=no-self-use
def _handle_anthropic_claude_response(
self, span: Span, response_body: dict[str, Any]
):
if usage := response_body.get("usage"):
if "input_tokens" in usage:
span.set_attribute(
GEN_AI_USAGE_INPUT_TOKENS, usage["input_tokens"]
)
if "output_tokens" in usage:
span.set_attribute(
GEN_AI_USAGE_OUTPUT_TOKENS, usage["output_tokens"]
)
if "stop_reason" in response_body:
span.set_attribute(
GEN_AI_RESPONSE_FINISH_REASONS, [response_body["stop_reason"]]
)
def on_error(self, span: Span, exception: _BotoClientErrorT):
if self._call_context.operation not in self._HANDLED_OPERATIONS:
return

View File

@ -14,14 +14,83 @@
from __future__ import annotations
import json
from typing import Any
from botocore.response import StreamingBody
from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.semconv._incubating.attributes import (
gen_ai_attributes as GenAIAttributes,
)
# pylint: disable=too-many-branches, too-many-locals
def assert_completion_attributes_from_streaming_body(
span: ReadableSpan,
request_model: str,
response: StreamingBody | None,
operation_name: str = "chat",
request_top_p: int | None = None,
request_temperature: int | None = None,
request_max_tokens: int | None = None,
request_stop_sequences: list[str] | None = None,
):
input_tokens = None
output_tokens = None
finish_reason = None
if response is not None:
original_body = response["body"]
body_content = original_body.read()
response = json.loads(body_content.decode("utf-8"))
assert response
if "amazon.titan" in request_model:
input_tokens = response.get("inputTextTokenCount")
results = response.get("results")
if results:
first_result = results[0]
output_tokens = first_result.get("tokenCount")
finish_reason = (first_result["completionReason"],)
elif "amazon.nova" in request_model:
if usage := response.get("usage"):
input_tokens = usage["inputTokens"]
output_tokens = usage["outputTokens"]
else:
input_tokens, output_tokens = None, None
if "stopReason" in response:
finish_reason = (response["stopReason"],)
else:
finish_reason = None
elif "anthropic.claude" in request_model:
if usage := response.get("usage"):
input_tokens = usage["input_tokens"]
output_tokens = usage["output_tokens"]
else:
input_tokens, output_tokens = None, None
if "stop_reason" in response:
finish_reason = (response["stop_reason"],)
else:
finish_reason = None
return assert_all_attributes(
span,
request_model,
input_tokens,
output_tokens,
finish_reason,
operation_name,
request_top_p,
request_temperature,
request_max_tokens,
tuple(request_stop_sequences)
if request_stop_sequences is not None
else request_stop_sequences,
)
def assert_completion_attributes(
span: ReadableSpan,
request_model: str,
@ -38,7 +107,7 @@ def assert_completion_attributes(
else:
input_tokens, output_tokens = None, None
if response:
if response and "stopReason" in response:
finish_reason = (response["stopReason"],)
else:
finish_reason = None
@ -60,10 +129,10 @@ def assert_completion_attributes(
def assert_equal_or_not_present(value, attribute_name, span):
if value:
if value is not None:
assert value == span.attributes[attribute_name]
else:
assert attribute_name not in span.attributes
assert attribute_name not in span.attributes, attribute_name
def assert_all_attributes(

View File

@ -1,26 +1,8 @@
interactions:
- request:
body: |-
{
"messages": [
{
"role": "user",
"content": [
{
"text": "Say this is a test"
}
]
}
],
"inferenceConfig": {
"maxTokens": 10,
"temperature": 0.8,
"topP": 1,
"stopSequences": [
"|"
]
}
}
body: '{"messages": [{"role": "user", "content": [{"text": "Say this is a test"}]}],
"inferenceConfig": {"maxTokens": 10, "temperature": 0.8, "topP": 1, "stopSequences":
["|"]}}'
headers:
Content-Length:
- '170'
@ -34,59 +16,39 @@ interactions:
aG9uIGNmZy9yZXRyeS1tb2RlI2xlZ2FjeSBCb3RvY29yZS8xLjM1LjU2
X-Amz-Date:
- !!binary |
MjAyNDEyMzFUMTMyMDQxWg==
MjAyNTAxMjJUMTYwODQwWg==
X-Amz-Security-Token:
- test_aws_security_token
X-Amzn-Trace-Id:
- !!binary |
Um9vdD0xLWY1MWY4NGM1LTNiZjk4YzY0YWMyNmJhNTk1OWJjODgxNjtQYXJlbnQ9YjNmOGZhM2Mz
MDc1NGEzZjtTYW1wbGVkPTE=
Um9vdD0xLTZjNTNiNTMyLTI0MDMzZTUwYzQ0M2JkODY2YTVhODhmMztQYXJlbnQ9MWM4ZDk4NmE2
Zjk1Y2Y0NTtTYW1wbGVkPTE=
amz-sdk-invocation-id:
- !!binary |
OTIyMjczMzItY2I5ZS00NGM1LTliZGUtYjU0NmJmODkxYmEy
MmRkMzAxNjUtYTdmOC00MjAzLWJlOGItZmE1ZWEzYmFjOGUy
amz-sdk-request:
- !!binary |
YXR0ZW1wdD0x
authorization:
- Bearer test_aws_authorization
method: POST
uri: https://bedrock-runtime.eu-central-1.amazonaws.com/model/amazon.titan-text-lite-v1/converse
uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-lite-v1/converse
response:
body:
string: |-
{
"metrics": {
"latencyMs": 811
},
"output": {
"message": {
"content": [
{
"text": "I am happy to assist you today"
}
],
"role": "assistant"
}
},
"stopReason": "max_tokens",
"usage": {
"inputTokens": 8,
"outputTokens": 10,
"totalTokens": 18
}
}
string: '{"metrics":{"latencyMs":742},"output":{"message":{"content":[{"text":"Hey
there! Is there anything else"}],"role":"assistant"}},"stopReason":"max_tokens","usage":{"inputTokens":8,"outputTokens":10,"totalTokens":18}}'
headers:
Connection:
- keep-alive
Content-Length:
- '212'
- '215'
Content-Type:
- application/json
Date:
- Tue, 31 Dec 2024 13:20:42 GMT
- Wed, 22 Jan 2025 16:08:41 GMT
Set-Cookie: test_set_cookie
x-amzn-RequestId:
- 63dfbcb2-3536-4906-b10d-e5b126b3c0ae
- 9fe3b968-40b3-400c-a48d-96fdf682557c
status:
code: 200
message: OK

View File

@ -1,18 +1,6 @@
interactions:
- request:
body: |-
{
"messages": [
{
"role": "user",
"content": [
{
"text": "Say this is a test"
}
]
}
]
}
body: '{"messages": [{"role": "user", "content": [{"text": "Say this is a test"}]}]}'
headers:
Content-Length:
- '77'
@ -26,29 +14,26 @@ interactions:
aG9uIGNmZy9yZXRyeS1tb2RlI2xlZ2FjeSBCb3RvY29yZS8xLjM1LjU2
X-Amz-Date:
- !!binary |
MjAyNTAxMTVUMTEwMTQ3Wg==
MjAyNTAxMjJUMTYwODQxWg==
X-Amz-Security-Token:
- test_aws_security_token
X-Amzn-Trace-Id:
- !!binary |
Um9vdD0xLWIzM2JhNTkxLTdkYmQ0ZDZmYTBmZTdmYzc2MTExOThmNztQYXJlbnQ9NzRmNmQ1NTEz
MzkzMzUxNTtTYW1wbGVkPTE=
Um9vdD0xLTY4MzBlNjVhLTY4Y2JlMzA5ZTI2ZDA1ZjA4ZDZkY2M1YjtQYXJlbnQ9NjdlMDRlNjRj
NGZhOTI3MDtTYW1wbGVkPTE=
amz-sdk-invocation-id:
- !!binary |
NTQ5MmQ0NTktNzhkNi00ZWY4LTlmMDMtZTA5ODhkZGRiZDI5
N2VhMWVmYzktMzlkYS00NDU1LWJiYTctMDNmYTM1ZWUyODU2
amz-sdk-request:
- !!binary |
YXR0ZW1wdD0x
authorization:
- Bearer test_aws_authorization
method: POST
uri: https://bedrock-runtime.eu-central-1.amazonaws.com/model/does-not-exist/converse
uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/does-not-exist/converse
response:
body:
string: |-
{
"message": "The provided model identifier is invalid."
}
string: '{"message":"The provided model identifier is invalid."}'
headers:
Connection:
- keep-alive
@ -57,12 +42,12 @@ interactions:
Content-Type:
- application/json
Date:
- Wed, 15 Jan 2025 11:01:47 GMT
- Wed, 22 Jan 2025 16:08:41 GMT
Set-Cookie: test_set_cookie
x-amzn-ErrorType:
- ValidationException:http://internal.amazon.com/coral/com.amazon.bedrock/
x-amzn-RequestId:
- d425bf99-8a4e-4d83-8d77-a48410dd82b2
- 9ecb3c28-f72f-4350-8746-97c02140ced1
status:
code: 400
message: Bad Request

View File

@ -0,0 +1,58 @@
interactions:
- request:
body: '{"messages": [{"role": "user", "content": [{"text": "Say this is a test"}]}],
"inferenceConfig": {"max_new_tokens": 10, "temperature": 0.8, "topP": 1, "stopSequences":
["|"]}, "schemaVersion": "messages-v1"}'
headers:
Content-Length:
- '207'
User-Agent:
- !!binary |
Qm90bzMvMS4zNS41NiBtZC9Cb3RvY29yZSMxLjM1LjU2IHVhLzIuMCBvcy9saW51eCM2LjEuMC0x
MDM0LW9lbSBtZC9hcmNoI3g4Nl82NCBsYW5nL3B5dGhvbiMzLjEwLjEyIG1kL3B5aW1wbCNDUHl0
aG9uIGNmZy9yZXRyeS1tb2RlI2xlZ2FjeSBCb3RvY29yZS8xLjM1LjU2
X-Amz-Date:
- !!binary |
MjAyNTAxMjJUMTUyNDA0Wg==
X-Amz-Security-Token:
- test_aws_security_token
X-Amzn-Trace-Id:
- !!binary |
Um9vdD0xLTY0ZGIzYWIxLTc2YWUzYmUxYmQ0NzI4Mzg1ZjdmOTEzZTtQYXJlbnQ9ZGRmYTdlZjI4
NWNiYTIxNTtTYW1wbGVkPTE=
amz-sdk-invocation-id:
- !!binary |
ZDZlMGIyOTUtYjM5Yi00NGU3LThiMmItZjgyODM2OTlkZTZk
amz-sdk-request:
- !!binary |
YXR0ZW1wdD0x
authorization:
- Bearer test_aws_authorization
method: POST
uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.nova-micro-v1%3A0/invoke
response:
body:
string: '{"output":{"message":{"content":[{"text":"It sounds like you might
be in the middle of"}],"role":"assistant"}},"stopReason":"max_tokens","usage":{"inputTokens":5,"outputTokens":10,"totalTokens":15}}'
headers:
Connection:
- keep-alive
Content-Length:
- '198'
Content-Type:
- application/json
Date:
- Wed, 22 Jan 2025 15:24:05 GMT
Set-Cookie: test_set_cookie
X-Amzn-Bedrock-Input-Token-Count:
- '5'
X-Amzn-Bedrock-Invocation-Latency:
- '237'
X-Amzn-Bedrock-Output-Token-Count:
- '10'
x-amzn-RequestId:
- 32f3134e-fc64-4db5-94bf-0279159cf79d
status:
code: 200
message: OK
version: 1

View File

@ -0,0 +1,57 @@
interactions:
- request:
body: '{"inputText": "Say this is a test", "textGenerationConfig": {"maxTokenCount":
10, "temperature": 0.8, "topP": 1, "stopSequences": ["|"]}}'
headers:
Content-Length:
- '137'
User-Agent:
- !!binary |
Qm90bzMvMS4zNS41NiBtZC9Cb3RvY29yZSMxLjM1LjU2IHVhLzIuMCBvcy9saW51eCM2LjEuMC0x
MDM0LW9lbSBtZC9hcmNoI3g4Nl82NCBsYW5nL3B5dGhvbiMzLjEwLjEyIG1kL3B5aW1wbCNDUHl0
aG9uIGNmZy9yZXRyeS1tb2RlI2xlZ2FjeSBCb3RvY29yZS8xLjM1LjU2
X-Amz-Date:
- !!binary |
MjAyNTAxMjJUMTUyNDA1Wg==
X-Amz-Security-Token:
- test_aws_security_token
X-Amzn-Trace-Id:
- !!binary |
Um9vdD0xLWZmMzM4ODA0LWMwMzYyNzgzNjczNjAzMWI0ZTZlZTIwNTtQYXJlbnQ9MmJjZmVlZGE5
NWVjZWUyYztTYW1wbGVkPTE=
amz-sdk-invocation-id:
- !!binary |
YmZjOGJiMjEtY2Q2MS00MDNmLWE2NzEtZmQ4YmMzNzBkOTJl
amz-sdk-request:
- !!binary |
YXR0ZW1wdD0x
authorization:
- Bearer test_aws_authorization
method: POST
uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-lite-v1/invoke
response:
body:
string: '{"inputTextTokenCount":5,"results":[{"tokenCount":9,"outputText":"
comment\nHello! How are you?","completionReason":"FINISH"}]}'
headers:
Connection:
- keep-alive
Content-Length:
- '127'
Content-Type:
- application/json
Date:
- Wed, 22 Jan 2025 15:24:06 GMT
Set-Cookie: test_set_cookie
X-Amzn-Bedrock-Input-Token-Count:
- '5'
X-Amzn-Bedrock-Invocation-Latency:
- '1104'
X-Amzn-Bedrock-Output-Token-Count:
- '9'
x-amzn-RequestId:
- ef788ecb-b5ed-404e-ace7-de59741cded5
status:
code: 200
message: OK
version: 1

View File

@ -0,0 +1,58 @@
interactions:
- request:
body: '{"messages": [{"role": "user", "content": [{"text": "Say this is a test",
"type": "text"}]}], "anthropic_version": "bedrock-2023-05-31", "max_tokens":
10, "temperature": 0.8, "top_p": 1, "stop_sequences": ["|"]}'
headers:
Content-Length:
- '211'
User-Agent:
- !!binary |
Qm90bzMvMS4zNS41NiBtZC9Cb3RvY29yZSMxLjM1LjU2IHVhLzIuMCBvcy9saW51eCM2LjEuMC0x
MDM0LW9lbSBtZC9hcmNoI3g4Nl82NCBsYW5nL3B5dGhvbiMzLjEwLjEyIG1kL3B5aW1wbCNDUHl0
aG9uIGNmZy9yZXRyeS1tb2RlI2xlZ2FjeSBCb3RvY29yZS8xLjM1LjU2
X-Amz-Date:
- !!binary |
MjAyNTAxMjJUMTUyNDA2Wg==
X-Amz-Security-Token:
- test_aws_security_token
X-Amzn-Trace-Id:
- !!binary |
Um9vdD0xLWQ2MDZiNDAzLWFhYzE1Y2I3ODBiOTkwMmIxNGU1NWM4ZjtQYXJlbnQ9YjJmMzRlMThk
ZWE4NjdkMztTYW1wbGVkPTE=
amz-sdk-invocation-id:
- !!binary |
YTlhN2I5YzEtNmEyNy00MDFjLTljMWUtM2EyN2YxZGZhMjQ4
amz-sdk-request:
- !!binary |
YXR0ZW1wdD0x
authorization:
- Bearer test_aws_authorization
method: POST
uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/invoke
response:
body:
string: '{"id":"msg_bdrk_01FJozYaVhprPHUzRZ2uVcMg","type":"message","role":"assistant","model":"claude-2.0","content":[{"type":"text","text":"OK,
I heard you say \"Say this is"}],"stop_reason":"max_tokens","stop_sequence":null,"usage":{"input_tokens":14,"output_tokens":10}}'
headers:
Connection:
- keep-alive
Content-Length:
- '265'
Content-Type:
- application/json
Date:
- Wed, 22 Jan 2025 15:24:07 GMT
Set-Cookie: test_set_cookie
X-Amzn-Bedrock-Input-Token-Count:
- '14'
X-Amzn-Bedrock-Invocation-Latency:
- '595'
X-Amzn-Bedrock-Output-Token-Count:
- '10'
x-amzn-RequestId:
- 5057dca6-bd9d-4e1e-9093-2bbbac1a19b4
status:
code: 200
message: OK
version: 1

View File

@ -0,0 +1,51 @@
interactions:
- request:
body: null
headers:
Content-Length:
- '0'
User-Agent:
- !!binary |
Qm90bzMvMS4zNS41NiBtZC9Cb3RvY29yZSMxLjM1LjU2IHVhLzIuMCBvcy9saW51eCM2LjEuMC0x
MDM0LW9lbSBtZC9hcmNoI3g4Nl82NCBsYW5nL3B5dGhvbiMzLjEwLjEyIG1kL3B5aW1wbCNDUHl0
aG9uIGNmZy9yZXRyeS1tb2RlI2xlZ2FjeSBCb3RvY29yZS8xLjM1LjU2
X-Amz-Date:
- !!binary |
MjAyNTAxMjJUMTUyNDA3Wg==
X-Amz-Security-Token:
- test_aws_security_token
X-Amzn-Trace-Id:
- !!binary |
Um9vdD0xLWVmZWZjYTdkLTM0OTI0ZjRmYTVlMDJmOTRhODFiY2M3NjtQYXJlbnQ9YWZiYmEwYjRh
MmU1NTQ0NDtTYW1wbGVkPTE=
amz-sdk-invocation-id:
- !!binary |
ODI0ZDAwZDgtMmE1Yy00Mzk4LWIwYTItOWY5ZmNlYjQ2MGNh
amz-sdk-request:
- !!binary |
YXR0ZW1wdD0x
authorization:
- Bearer test_aws_authorization
method: POST
uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/does-not-exist/invoke
response:
body:
string: '{"message":"The provided model identifier is invalid."}'
headers:
Connection:
- keep-alive
Content-Length:
- '55'
Content-Type:
- application/json
Date:
- Wed, 22 Jan 2025 15:24:08 GMT
Set-Cookie: test_set_cookie
x-amzn-ErrorType:
- ValidationException:http://internal.amazon.com/coral/com.amazon.bedrock/
x-amzn-RequestId:
- 9739ef10-1ae7-4694-ba63-3a39e7ca02c1
status:
code: 400
message: Bad Request
version: 1

View File

@ -1,11 +1,9 @@
"""Unit tests configuration module."""
import json
import os
import boto3
import pytest
import yaml
from opentelemetry.instrumentation.botocore import BotocoreInstrumentor
from opentelemetry.instrumentation.botocore.environment_variables import (
@ -66,7 +64,7 @@ def environment():
if not os.getenv("AWS_SESSION_TOKEN"):
os.environ["AWS_SESSION_TOKEN"] = "test_aws_session_token"
if not os.getenv("AWS_DEFAULT_REGION"):
os.environ["AWS_DEFAULT_REGION"] = "eu-central-1"
os.environ["AWS_DEFAULT_REGION"] = "us-east-1"
@pytest.fixture(scope="module")
@ -115,73 +113,6 @@ def instrument_with_content(tracer_provider, event_logger_provider):
instrumentor.uninstrument()
class LiteralBlockScalar(str):
"""Formats the string as a literal block scalar, preserving whitespace and
without interpreting escape characters"""
def literal_block_scalar_presenter(dumper, data):
"""Represents a scalar string as a literal block, via '|' syntax"""
return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")
yaml.add_representer(LiteralBlockScalar, literal_block_scalar_presenter)
def process_string_value(string_value):
"""Pretty-prints JSON or returns long strings as a LiteralBlockScalar"""
try:
json_data = json.loads(string_value)
return LiteralBlockScalar(json.dumps(json_data, indent=2))
except (ValueError, TypeError):
if len(string_value) > 80:
return LiteralBlockScalar(string_value)
return string_value
def convert_body_to_literal(data):
"""Searches the data for body strings, attempting to pretty-print JSON"""
if isinstance(data, dict):
for key, value in data.items():
# Handle response body case (e.g., response.body.string)
if key == "body" and isinstance(value, dict) and "string" in value:
value["string"] = process_string_value(value["string"])
# Handle request body case (e.g., request.body)
elif key == "body" and isinstance(value, str):
data[key] = process_string_value(value)
else:
convert_body_to_literal(value)
elif isinstance(data, list):
for idx, choice in enumerate(data):
data[idx] = convert_body_to_literal(choice)
return data
class PrettyPrintJSONBody:
"""This makes request and response body recordings more readable."""
@staticmethod
def serialize(cassette_dict):
cassette_dict = convert_body_to_literal(cassette_dict)
return yaml.dump(
cassette_dict, default_flow_style=False, allow_unicode=True
)
@staticmethod
def deserialize(cassette_string):
return yaml.load(cassette_string, Loader=yaml.Loader)
@pytest.fixture(scope="module", autouse=True)
def fixture_vcr(vcr):
vcr.register_serializer("yaml", PrettyPrintJSONBody)
return vcr
def scrub_response_headers(response):
"""
This scrubs sensitive response headers. Note they are case-sensitive!

View File

@ -14,6 +14,8 @@
from __future__ import annotations
import json
import boto3
import pytest
@ -22,7 +24,10 @@ from opentelemetry.semconv._incubating.attributes.error_attributes import (
)
from opentelemetry.trace.status import StatusCode
from .bedrock_utils import assert_completion_attributes
from .bedrock_utils import (
assert_completion_attributes,
assert_completion_attributes_from_streaming_body,
)
BOTO3_VERSION = tuple(int(x) for x in boto3.__version__.split("."))
@ -100,3 +105,126 @@ def test_converse_with_invalid_model(
logs = log_exporter.get_finished_logs()
assert len(logs) == 0
def get_invoke_model_body(
llm_model,
max_tokens=None,
temperature=None,
top_p=None,
stop_sequences=None,
):
def set_if_not_none(config, key, value):
if value is not None:
config[key] = value
prompt = "Say this is a test"
if llm_model == "amazon.nova-micro-v1:0":
config = {}
set_if_not_none(config, "max_new_tokens", max_tokens)
set_if_not_none(config, "temperature", temperature)
set_if_not_none(config, "topP", top_p)
set_if_not_none(config, "stopSequences", stop_sequences)
body = {
"messages": [{"role": "user", "content": [{"text": prompt}]}],
"inferenceConfig": config,
"schemaVersion": "messages-v1",
}
elif llm_model == "amazon.titan-text-lite-v1":
config = {}
set_if_not_none(config, "maxTokenCount", max_tokens)
set_if_not_none(config, "temperature", temperature)
set_if_not_none(config, "topP", top_p)
set_if_not_none(config, "stopSequences", stop_sequences)
body = {"inputText": prompt, "textGenerationConfig": config}
elif llm_model == "anthropic.claude-v2":
body = {
"messages": [
{"role": "user", "content": [{"text": prompt, "type": "text"}]}
],
"anthropic_version": "bedrock-2023-05-31",
}
set_if_not_none(body, "max_tokens", max_tokens)
set_if_not_none(body, "temperature", temperature)
set_if_not_none(body, "top_p", top_p)
set_if_not_none(body, "stop_sequences", stop_sequences)
else:
raise ValueError(f"No config for {llm_model}")
return json.dumps(body)
def get_model_name_from_family(llm_model):
llm_model_name = {
"amazon.titan": "amazon.titan-text-lite-v1",
"amazon.nova": "amazon.nova-micro-v1:0",
"anthropic.claude": "anthropic.claude-v2",
}
return llm_model_name[llm_model]
@pytest.mark.parametrize(
"model_family",
["amazon.nova", "amazon.titan", "anthropic.claude"],
)
@pytest.mark.vcr()
def test_invoke_model_with_content(
span_exporter,
log_exporter,
bedrock_runtime_client,
instrument_with_content,
model_family,
):
llm_model_value = get_model_name_from_family(model_family)
max_tokens, temperature, top_p, stop_sequences = 10, 0.8, 1, ["|"]
body = get_invoke_model_body(
llm_model_value, max_tokens, temperature, top_p, stop_sequences
)
response = bedrock_runtime_client.invoke_model(
body=body,
modelId=llm_model_value,
)
(span,) = span_exporter.get_finished_spans()
assert_completion_attributes_from_streaming_body(
span,
llm_model_value,
response,
"text_completion" if model_family == "amazon.titan" else "chat",
top_p,
temperature,
max_tokens,
stop_sequences,
)
logs = log_exporter.get_finished_logs()
assert len(logs) == 0
@pytest.mark.vcr()
def test_invoke_model_with_invalid_model(
span_exporter,
log_exporter,
bedrock_runtime_client,
instrument_with_content,
):
llm_model_value = "does-not-exist"
with pytest.raises(bedrock_runtime_client.exceptions.ClientError):
bedrock_runtime_client.invoke_model(
body=b"",
modelId=llm_model_value,
)
(span,) = span_exporter.get_finished_spans()
assert_completion_attributes_from_streaming_body(
span,
llm_model_value,
None,
"chat",
)
assert span.status.status_code == StatusCode.ERROR
assert span.attributes[ERROR_TYPE] == "ValidationException"
logs = log_exporter.get_finished_logs()
assert len(logs) == 0