add support for choice events and messages

Author: ezhang6811

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

@@ -380,11 +380,25 @@ def _get_request_messages(self):
 
                 messages = decoded_body.get("messages", [])
                 if not messages:
-                    # transform old school amazon titan invokeModel api to messages
-                    if input_text := decoded_body.get("inputText"):
-                        messages = [
-                            {"role": "user", "content": [{"text": input_text}]}
-                        ]
+                    model_id = self._call_context.params.get(_MODEL_ID_KEY)
+                    if "amazon.titan" in model_id:
+                        # transform old school amazon titan invokeModel api to messages
+                        if input_text := decoded_body.get("inputText"):
+                            messages = [
+                                {"role": "user", "content": [{"text": input_text}]}
+                            ]
+                    elif "cohere.command-r" in model_id:
+                        # chat_history can be converted to messages; for now, just use message
+                        if input_text := decoded_body.get("message"):
+                            messages = [
+                                {"role": "user", "content": [{"text": input_text}]}
+                            ]
+                    elif "cohere.command" in model_id or "meta.llama" in model_id or "mistral.mistral" in model_id:
+                        # transform old school cohere command api to messages
+                        if input_text := decoded_body.get("prompt"):
+                            messages = [
+                                {"role": "user", "content": [{"text": input_text}]}
+                            ]
 
         return system_messages + messages
 
@@ -831,6 +845,12 @@ def _handle_cohere_command_r_response(
             span.set_attribute(
                 GEN_AI_RESPONSE_FINISH_REASONS, [response_body["finish_reason"]]
             )
+        
+        event_logger = instrumentor_context.event_logger
+        choice = _Choice.from_invoke_cohere_command_r(
+            response_body, capture_content
+        )
+        event_logger.emit(choice.to_choice_event())
 
     def _handle_cohere_command_response(
         self,
@@ -849,6 +869,12 @@ def _handle_cohere_command_response(
                 span.set_attribute(
                     GEN_AI_RESPONSE_FINISH_REASONS, [generations["finish_reason"]]
                 )
+        
+        event_logger = instrumentor_context.event_logger
+        choice = _Choice.from_invoke_cohere_command(
+            response_body, capture_content
+        )
+        event_logger.emit(choice.to_choice_event())
 
     def _handle_meta_llama_response(
         self,
@@ -870,6 +896,12 @@ def _handle_meta_llama_response(
                 GEN_AI_RESPONSE_FINISH_REASONS, [response_body["stop_reason"]]
             )
 
+        event_logger = instrumentor_context.event_logger
+        choice = _Choice.from_invoke_meta_llama(
+            response_body, capture_content
+        )
+        event_logger.emit(choice.to_choice_event())
+
     def _handle_mistral_ai_response(
         self,
         span: Span,
@@ -883,6 +915,12 @@ def _handle_mistral_ai_response(
                 span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, math.ceil(len(outputs["text"]) / 6))
             if "stop_reason" in outputs:
                 span.set_attribute(GEN_AI_RESPONSE_FINISH_REASONS, [outputs["stop_reason"]])
+        
+        event_logger = instrumentor_context.event_logger
+        choice = _Choice.from_invoke_mistral_mistral(
+            response_body, capture_content
+        )
+        event_logger.emit(choice.to_choice_event())
 
     def on_error(
         self,

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

@@ -519,6 +519,48 @@ def from_invoke_anthropic_claude(
             message["content"] = response["content"]
         return cls(message, response["stop_reason"], index=0)
 
+    @classmethod
+    def from_invoke_cohere_command_r(
+        cls, response: dict[str, Any], capture_content: bool
+    ) -> _Choice:
+        if capture_content:
+            message = {"content": response["text"]}
+        else:
+            message = {}
+        return cls(message, response["finish_reason"], index=0)
+    
+    @classmethod
+    def from_invoke_cohere_command(
+        cls, response: dict[str, Any], capture_content: bool
+    ) -> _Choice:
+        result = response["generations"][0]
+        if capture_content:
+            message = {"content": result["text"]}
+        else:
+            message = {}
+        return cls(message, result["finish_reason"], index=0)
+
+    @classmethod
+    def from_invoke_meta_llama(
+        cls, response: dict[str, Any], capture_content: bool
+    ) -> _Choice:
+        if capture_content:
+            message = {"content": response["generation"]}
+        else:
+            message = {}
+        return cls(message, response["stop_reason"], index=0)
+    
+    @classmethod
+    def from_invoke_mistral_mistral(
+        cls, response: dict[str, Any], capture_content: bool
+    ) -> _Choice:
+        result = response["outputs"][0]
+        if capture_content:
+            message = {"content": result["text"]}
+        else:
+            message = {}
+        return cls(message, result["stop_reason"], index=0)
+
     def _to_body_dict(self) -> dict[str, Any]:
         return {
             "finish_reason": self.finish_reason,

instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_with_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.99,
+        "stop_sequences": [
+          "|"
+        ]
+      }
+    headers:
+      Content-Length:
+      - '107'
+      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:
+      - 20250410T224018Z
+      X-Amz-Security-Token:
+      - test_aws_security_token
+      X-Amzn-Trace-Id:
+      - Root=1-51673f20-3dd018601b078c785c032a50;Parent=c05b84f54e2c35ee;Sampled=1
+      amz-sdk-invocation-id:
+      - c81a224e-d8b3-4416-aaca-858b478b7db4
+      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/aa2df2bf-edc8-483f-8cd0-d22d04ba34ba",
+          "text": "This is a test. How are you doing today",
+          "generation_id": "2d74c447-266d-4286-8425-a92ad7fc0cbc",
+          "chat_history": [
+            {
+              "role": "USER",
+              "message": "Say this is a test"
+            },
+            {
+              "role": "CHATBOT",
+              "message": "This is a test. How are you doing today"
+            }
+          ],
+          "finish_reason": "MAX_TOKENS"
+        }
+    headers:
+      Connection:
+      - keep-alive
+      Content-Type:
+      - application/json
+      Date:
+      - Thu, 10 Apr 2025 22:40:18 GMT
+      Set-Cookie: test_set_cookie
+      X-Amzn-Bedrock-Input-Token-Count:
+      - '5'
+      X-Amzn-Bedrock-Invocation-Latency:
+      - '196'
+      X-Amzn-Bedrock-Output-Token-Count:
+      - '10'
+      x-amzn-RequestId:
+      - c1fd38df-b669-4464-9f4b-4cf32c32fde8
+    status:
+      code: 200
+      message: OK
+version: 1

instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_with_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:
+      - 20250410T223753Z
+      X-Amz-Security-Token:
+      - test_aws_security_token
+      X-Amzn-Trace-Id:
+      - Root=1-461bf523-e7a0388221ce8a1d7a31fc5b;Parent=a55a9398f0f3d364;Sampled=1
+      amz-sdk-invocation-id:
+      - 17fbe175-373f-4c10-8047-8bba2fd8e22c
+      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": "a09c1c60-6608-482a-b98d-764e4d87fcd1",
+          "generations": [
+            {
+              "id": "8b38ce59-5f77-4b79-82c3-58eb72d4b1a2",
+              "text": " Let it be a test of knowledge, skills,",
+              "finish_reason": "MAX_TOKENS"
+            }
+          ],
+          "prompt": "Say this is a test"
+        }
+    headers:
+      Connection:
+      - keep-alive
+      Content-Type:
+      - application/json
+      Date:
+      - Thu, 10 Apr 2025 22:37:53 GMT
+      Set-Cookie: test_set_cookie
+      X-Amzn-Bedrock-Input-Token-Count:
+      - '5'
+      X-Amzn-Bedrock-Invocation-Latency:
+      - '248'
+      X-Amzn-Bedrock-Output-Token-Count:
+      - '10'
+      x-amzn-RequestId:
+      - a09c1c60-6608-482a-b98d-764e4d87fcd1
+    status:
+      code: 200
+      message: OK
+version: 1

instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_with_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:
+      - 20250410T224059Z
+      X-Amz-Security-Token:
+      - test_aws_security_token
+      X-Amzn-Trace-Id:
+      - Root=1-9930be15-db431eb0bf3c21be1ed23dde;Parent=0c4e46fcf0474877;Sampled=1
+      amz-sdk-invocation-id:
+      - c3afc5f6-9912-4892-b48e-dff457911e1c
+      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": " post. This is a test post. This is",
+          "prompt_token_count": 5,
+          "generation_token_count": 10,
+          "stop_reason": "length"
+        }
+    headers:
+      Connection:
+      - keep-alive
+      Content-Type:
+      - application/json
+      Date:
+      - Thu, 10 Apr 2025 22:40:59 GMT
+      Set-Cookie: test_set_cookie
+      X-Amzn-Bedrock-Input-Token-Count:
+      - '5'
+      X-Amzn-Bedrock-Invocation-Latency:
+      - '604'
+      X-Amzn-Bedrock-Output-Token-Count:
+      - '10'
+      x-amzn-RequestId:
+      - 4b8da626-00b1-4535-b75b-a7e985e9877e
+    status:
+      code: 200
+      message: OK
+version: 1

instrumentation/opentelemetry-instrumentation-botocore/tests/cassettes/test_invoke_model_with_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:
+      - 20250410T224059Z
+      X-Amz-Security-Token:
+      - test_aws_security_token
+      X-Amzn-Trace-Id:
+      - Root=1-332feae9-215a9553da54b4789bf12414;Parent=f3899b50f888d5d7;Sampled=1
+      amz-sdk-invocation-id:
+      - 2d410cd2-ef68-48af-9cad-b7ddf3b1d844
+      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": "\n\nA man stands before a crowd of people",
+              "stop_reason": "length"
+            }
+          ]
+        }
+    headers:
+      Connection:
+      - keep-alive
+      Content-Type:
+      - application/json
+      Date:
+      - Thu, 10 Apr 2025 22:41:00 GMT
+      Set-Cookie: test_set_cookie
+      X-Amzn-Bedrock-Input-Token-Count:
+      - '6'
+      X-Amzn-Bedrock-Invocation-Latency:
+      - '174'
+      X-Amzn-Bedrock-Output-Token-Count:
+      - '10'
+      x-amzn-RequestId:
+      - 68547039-f30c-4aff-b368-31bb6d9b9d07
+    status:
+      code: 200
+      message: OK
+version: 1