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>

CHANGELOG.md

@@ -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
 

instrumentation/opentelemetry-instrumentation-botocore/examples/bedrock-runtime/zero-code/invoke_model.py

@@ -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()

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/bedrock.py

@@ -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]):
+    # pylint: disable=no-self-use
-        if self._call_context.operation not in self._HANDLED_OPERATIONS:
+    def _converse_on_success(self, span: Span, result: dict[str, Any]):
-            return
-
-        if not span.is_recording():
-            return
-
         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

instrumentation/opentelemetry-instrumentation-botocore/tests/bedrock_utils.py

@@ -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(

instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_converse_with_content.yaml

@@ -1,26 +1,8 @@
 interactions:
 - request:
-    body: |-
+    body: '{"messages": [{"role": "user", "content": [{"text": "Say this is a test"}]}],
-      {
+      "inferenceConfig": {"maxTokens": 10, "temperature": 0.8, "topP": 1, "stopSequences":
-        "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
+        Um9vdD0xLTZjNTNiNTMyLTI0MDMzZTUwYzQ0M2JkODY2YTVhODhmMztQYXJlbnQ9MWM4ZDk4NmE2
-        MDc1NGEzZjtTYW1wbGVkPTE=
+        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: |-
+      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}}'
-          "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
-          }
-        }
     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

instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_converse_with_invalid_model.yaml

@@ -1,18 +1,6 @@
 interactions:
 - request:
-    body: |-
+    body: '{"messages": [{"role": "user", "content": [{"text": "Say this is a test"}]}]}'
-      {
-        "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
+        Um9vdD0xLTY4MzBlNjVhLTY4Y2JlMzA5ZTI2ZDA1ZjA4ZDZkY2M1YjtQYXJlbnQ9NjdlMDRlNjRj
-        MzkzMzUxNTtTYW1wbGVkPTE=
+        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: |-
+      string: '{"message":"The provided model identifier is invalid."}'
-        {
-          "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

instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_with_content[amazon.nova].yaml

@@ -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

instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_with_content[amazon.titan].yaml

@@ -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

instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_with_content[anthropic.claude].yaml

@@ -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

instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_with_invalid_model.yaml

@@ -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

instrumentation/opentelemetry-instrumentation-botocore/tests/conftest.py

@@ -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!

instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_bedrock.py

@@ -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