add unit tests for invoke_model attributes

Author: ezhang6811

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

@@ -21,6 +21,7 @@
 import io
 import json
 import logging
+import math
 from timeit import default_timer
 from typing import Any
 
@@ -514,11 +515,11 @@ def _invoke_model_on_success(
                     span, response_body, instrumentor_context, capture_content
                 )
             elif "cohere.command-r" in model_id:
-                self._handle_command_r_response(
+                self._handle_cohere_command_r_response(
                     span, response_body, instrumentor_context, capture_content
                 )
             elif "cohere.command" in model_id:
-                self._handle_command_response(
+                self._handle_cohere_command_response(
                     span, response_body, instrumentor_context, capture_content
                 )
             elif "meta.llama" in model_id:

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

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import json
+import math
 from typing import Any
 
 from botocore.response import StreamingBody
@@ -54,6 +55,7 @@ def assert_completion_attributes_from_streaming_body(
     input_tokens = None
     output_tokens = None
     finish_reason = None
+    request_prompt = "Say this is a test"
     if response is not None:
         original_body = response["body"]
         body_content = original_body.read()
@@ -89,6 +91,33 @@ def assert_completion_attributes_from_streaming_body(
                 finish_reason = (response["stop_reason"],)
             else:
                 finish_reason = None
+        elif "cohere.command-r" in request_model:
+            input_tokens = math.ceil(len(request_prompt) / 6)
+            text = response.get("text")
+            if text:
+                output_tokens = math.ceil(len(text) / 6)
+            finish_reason = (response["finish_reason"],)
+        elif "cohere.command" in request_model:
+            input_tokens = math.ceil(len(request_prompt) / 6)
+            generations = response.get("generations")
+            if generations:
+                first_generation = generations[0]
+                output_tokens = math.ceil(len(first_generation["text"]) / 6)
+                finish_reason = (first_generation["finish_reason"],)
+        elif "meta.llama" in request_model:
+            if "prompt_token_count" in response:
+                input_tokens = response.get("prompt_token_count")
+            if "generation_token_count" in response:
+                output_tokens = response.get("generation_token_count")
+            if "stop_reason" in response:
+                finish_reason = (response["stop_reason"],)
+        elif "mistral.mistral" in request_model:
+            input_tokens = math.ceil(len(request_prompt) / 6)
+            outputs = response.get("outputs")
+            if outputs:
+                first_output = outputs[0]
+                output_tokens = math.ceil(len(first_output["text"]) / 6)
+                finish_reason = (first_output["stop_reason"],)
 
     return assert_all_attributes(
         span,

instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_no_content[cohere.command-r].yaml

@@ -0,0 +1,71 @@
+interactions:
+- request:
+    body: |-
+      {
+        "message": "Say this is a test",
+        "max_tokens": 10,
+        "temperature": 0.8,
+        "p": 0.9,
+        "stop_sequences": [
+          "|"
+        ]
+      }
+    headers:
+      Content-Length:
+      - '106'
+      User-Agent:
+      - Boto3/1.35.56 md/Botocore#1.35.56 ua/2.0 os/macos#24.3.0 md/arch#arm64 lang/python#3.10.16
+        md/pyimpl#CPython cfg/retry-mode#legacy Botocore/1.35.56
+      X-Amz-Date:
+      - 20250410T165615Z
+      X-Amz-Security-Token:
+      - test_aws_security_token
+      X-Amzn-Trace-Id:
+      - Root=1-d11162cc-5e3ef3bd07c5f9eb3ad25214;Parent=0a40817ca75eead7;Sampled=1
+      amz-sdk-invocation-id:
+      - 54acc108-181d-4437-a0a2-8293496f892d
+      amz-sdk-request:
+      - attempt=1
+      authorization:
+      - Bearer test_aws_authorization
+    method: POST
+    uri: https://bedrock-runtime.us-west-2.amazonaws.com/model/cohere.command-r-v1%3A0/invoke
+  response:
+    body:
+      string: |-
+        {
+          "response_id": "379ed018/64bcb4c2-902c-4646-9e54-3e5e90c3b11e",
+          "text": "This is a test.<",
+          "generation_id": "dcf7f3e3-a611-4fb5-85fe-ae09e6b2eaea",
+          "chat_history": [
+            {
+              "role": "USER",
+              "message": "Say this is a test"
+            },
+            {
+              "role": "CHATBOT",
+              "message": "This is a test.<"
+            }
+          ],
+          "finish_reason": "STOP_SEQUENCE"
+        }
+    headers:
+      Connection:
+      - keep-alive
+      Content-Type:
+      - application/json
+      Date:
+      - Thu, 10 Apr 2025 16:56:15 GMT
+      Set-Cookie: test_set_cookie
+      X-Amzn-Bedrock-Input-Token-Count:
+      - '5'
+      X-Amzn-Bedrock-Invocation-Latency:
+      - '141'
+      X-Amzn-Bedrock-Output-Token-Count:
+      - '5'
+      x-amzn-RequestId:
+      - 4d1ce10c-03d3-4579-8bab-a0e4eb14a8d0
+    status:
+      code: 200
+      message: OK
+version: 1

instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_no_content[cohere.command].yaml

@@ -0,0 +1,66 @@
+interactions:
+- request:
+    body: |-
+      {
+        "prompt": "Say this is a test",
+        "max_tokens": 10,
+        "temperature": 0.8,
+        "p": 1,
+        "stop_sequences": [
+          "|"
+        ]
+      }
+    headers:
+      Content-Length:
+      - '103'
+      User-Agent:
+      - Boto3/1.35.56 md/Botocore#1.35.56 ua/2.0 os/macos#24.3.0 md/arch#arm64 lang/python#3.10.16
+        md/pyimpl#CPython cfg/retry-mode#legacy Botocore/1.35.56
+      X-Amz-Date:
+      - 20250410T183021Z
+      X-Amz-Security-Token:
+      - test_aws_security_token
+      X-Amzn-Trace-Id:
+      - Root=1-668b3eaf-1927d32dd30fc4f0fedf3a02;Parent=6b49f910f73172c6;Sampled=1
+      amz-sdk-invocation-id:
+      - 98cb0162-a3bb-4d4e-bcbf-290664700ecf
+      amz-sdk-request:
+      - attempt=1
+      authorization:
+      - Bearer test_aws_authorization
+    method: POST
+    uri: https://bedrock-runtime.us-west-2.amazonaws.com/model/cohere.command-light-text-v14/invoke
+  response:
+    body:
+      string: |-
+        {
+          "id": "b68fba0b-67f6-46dd-a659-dab9a99ee354",
+          "generations": [
+            {
+              "id": "195692c2-6457-4ae0-9e5e-f0472a429cdc",
+              "text": " I would be more than happy to assist you with",
+              "finish_reason": "MAX_TOKENS"
+            }
+          ],
+          "prompt": "Say this is a test"
+        }
+    headers:
+      Connection:
+      - keep-alive
+      Content-Type:
+      - application/json
+      Date:
+      - Thu, 10 Apr 2025 18:30:21 GMT
+      Set-Cookie: test_set_cookie
+      X-Amzn-Bedrock-Input-Token-Count:
+      - '5'
+      X-Amzn-Bedrock-Invocation-Latency:
+      - '309'
+      X-Amzn-Bedrock-Output-Token-Count:
+      - '10'
+      x-amzn-RequestId:
+      - b68fba0b-67f6-46dd-a659-dab9a99ee354
+    status:
+      code: 200
+      message: OK
+version: 1

instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_no_content[meta.llama].yaml

@@ -0,0 +1,58 @@
+interactions:
+- request:
+    body: |-
+      {
+        "prompt": "Say this is a test",
+        "max_gen_len": 10,
+        "temperature": 0.8,
+        "top_p": 1
+      }
+    headers:
+      Content-Length:
+      - '83'
+      User-Agent:
+      - Boto3/1.35.56 md/Botocore#1.35.56 ua/2.0 os/macos#24.3.0 md/arch#arm64 lang/python#3.10.16
+        md/pyimpl#CPython cfg/retry-mode#legacy Botocore/1.35.56
+      X-Amz-Date:
+      - 20250410T192408Z
+      X-Amz-Security-Token:
+      - test_aws_security_token
+      X-Amzn-Trace-Id:
+      - Root=1-db2cb3ce-1a9a7362f6f0e5eb7dc0ce56;Parent=a301dca505171c7f;Sampled=1
+      amz-sdk-invocation-id:
+      - 87acada6-b95b-4280-b002-50ef98edf99d
+      amz-sdk-request:
+      - attempt=1
+      authorization:
+      - Bearer test_aws_authorization
+    method: POST
+    uri: https://bedrock-runtime.us-west-2.amazonaws.com/model/meta.llama3-1-70b-instruct-v1%3A0/invoke
+  response:
+    body:
+      string: |-
+        {
+          "generation": " and you are a test subject. You are not",
+          "prompt_token_count": 5,
+          "generation_token_count": 10,
+          "stop_reason": "length"
+        }
+    headers:
+      Connection:
+      - keep-alive
+      Content-Type:
+      - application/json
+      Date:
+      - Thu, 10 Apr 2025 19:24:09 GMT
+      Set-Cookie: test_set_cookie
+      X-Amzn-Bedrock-Input-Token-Count:
+      - '5'
+      X-Amzn-Bedrock-Invocation-Latency:
+      - '617'
+      X-Amzn-Bedrock-Output-Token-Count:
+      - '10'
+      x-amzn-RequestId:
+      - 51b421ef-fd84-4c69-b933-8c928b735d17
+    status:
+      code: 200
+      message: OK
+version: 1

instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_no_content[mistral.mistral].yaml

@@ -0,0 +1,63 @@
+interactions:
+- request:
+    body: |-
+      {
+        "prompt": "Say this is a test",
+        "max_tokens": 10,
+        "temperature": 0.8,
+        "top_p": 1,
+        "stop": [
+          "|"
+        ]
+      }
+    headers:
+      Content-Length:
+      - '97'
+      User-Agent:
+      - Boto3/1.35.56 md/Botocore#1.35.56 ua/2.0 os/macos#24.3.0 md/arch#arm64 lang/python#3.10.16
+        md/pyimpl#CPython cfg/retry-mode#legacy Botocore/1.35.56
+      X-Amz-Date:
+      - 20250410T183022Z
+      X-Amz-Security-Token:
+      - test_aws_security_token
+      X-Amzn-Trace-Id:
+      - Root=1-e325095b-1fab9945412a3dcac8474715;Parent=ae78dbba36fcf5b0;Sampled=1
+      amz-sdk-invocation-id:
+      - cbfbbad5-adc8-4a71-82f0-0777f6c78f83
+      amz-sdk-request:
+      - attempt=1
+      authorization:
+      - Bearer test_aws_authorization
+    method: POST
+    uri: https://bedrock-runtime.us-west-2.amazonaws.com/model/mistral.mistral-7b-instruct-v0%3A2/invoke
+  response:
+    body:
+      string: |-
+        {
+          "outputs": [
+            {
+              "text": " to find out whether you think an animal is cute",
+              "stop_reason": "length"
+            }
+          ]
+        }
+    headers:
+      Connection:
+      - keep-alive
+      Content-Type:
+      - application/json
+      Date:
+      - Thu, 10 Apr 2025 18:30:22 GMT
+      Set-Cookie: test_set_cookie
+      X-Amzn-Bedrock-Input-Token-Count:
+      - '6'
+      X-Amzn-Bedrock-Invocation-Latency:
+      - '219'
+      X-Amzn-Bedrock-Output-Token-Count:
+      - '10'
+      x-amzn-RequestId:
+      - 97084131-1cef-4927-9bab-53a36390153e
+    status:
+      code: 200
+      message: OK
+version: 1