Add tests and handle only non streaming responses

Author: xrmx

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/environment_variables.py

@@ -0,0 +1,3 @@
+OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT = (
+    "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT"
+)

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

@@ -18,14 +18,11 @@
 
 from __future__ import annotations
 
-import io
 import json
 import logging
 import math
 from typing import Any
 
-from botocore.response import StreamingBody
-
 from opentelemetry.instrumentation.botocore.extensions.types import (
     _AttributeMapT,
     _AwsSdkExtension,
@@ -58,13 +55,19 @@ class _BedrockRuntimeExtension(_AwsSdkExtension):
     """
 
     def extract_attributes(self, attributes: _AttributeMapT):
-        attributes[GEN_AI_SYSTEM] = GenAiSystemValues.AWS_BEDROCK
-        attributes[GEN_AI_OPERATION_NAME] = GenAiOperationNameValues.CHAT
+        attributes[GEN_AI_SYSTEM] = GenAiSystemValues.AWS_BEDROCK.value
 
         model_id = self._call_context.params.get(_MODEL_ID_KEY)
         if model_id:
             attributes[GEN_AI_REQUEST_MODEL] = model_id
 
+            # FIXME: add other model patterns
+            text_model_patterns = ["amazon.titan-text"]
+            if any(pattern in model_id for pattern in text_model_patterns):
+                attributes[GEN_AI_OPERATION_NAME] = (
+                    GenAiOperationNameValues.CHAT.value
+                )
+
             # Get the request body if it exists
             body = self._call_context.params.get("body")
             if body:
@@ -209,176 +212,38 @@ def _set_if_not_none(attributes, key, value):
         if value is not None:
             attributes[key] = value
 
+    def before_service_call(self, span: Span):
+        if not span.is_recording():
+            return
+
+        operation_name = span.attributes.get(GEN_AI_OPERATION_NAME, "")
+        request_model = span.attributes.get(GEN_AI_REQUEST_MODEL, "")
+        # avoid setting to an empty string if are not available
+        if operation_name and request_model:
+            span.update_name(f"{operation_name} {request_model}")
+
     # pylint: disable=too-many-branches
     def on_success(self, span: Span, result: dict[str, Any]):
         model_id = self._call_context.params.get(_MODEL_ID_KEY)
 
         if not model_id:
             return
 
-        if "body" in result and isinstance(result["body"], StreamingBody):
-            original_body = None
-            try:
-                original_body = result["body"]
-                body_content = original_body.read()
-
-                # Use one stream for telemetry
-                stream = io.BytesIO(body_content)
-                telemetry_content = stream.read()
-                response_body = json.loads(telemetry_content.decode("utf-8"))
-                if "amazon.titan" in model_id:
-                    self._handle_amazon_titan_response(span, response_body)
-                if "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)
-                elif "meta.llama" in model_id:
-                    self._handle_meta_llama_response(span, response_body)
-                elif "cohere.command" in model_id:
-                    self._handle_cohere_command_response(span, response_body)
-                elif "ai21.jamba" in model_id:
-                    self._handle_ai21_jamba_response(span, response_body)
-                elif "mistral" in model_id:
-                    self._handle_mistral_mistral_response(span, response_body)
-                # Replenish stream for downstream application use
-                new_stream = io.BytesIO(body_content)
-                result["body"] = StreamingBody(new_stream, len(body_content))
-
-            except json.JSONDecodeError:
-                _logger.debug(
-                    "Error: Unable to parse the response body as JSON"
-                )
-            except Exception as e:  # pylint: disable=broad-exception-caught, invalid-name
-                _logger.debug("Error processing response: %s", e)
-            finally:
-                if original_body is not None:
-                    original_body.close()
-
-    # 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:
+        # FIXME: this is tested only with titan
+        if usage := result.get("usage"):
+            if input_tokens := usage.get("inputTokens"):
                 span.set_attribute(
-                    GEN_AI_USAGE_OUTPUT_TOKENS, usage["outputTokens"]
+                    GEN_AI_USAGE_INPUT_TOKENS,
+                    input_tokens,
                 )
-        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" in response_body:
-            usage = response_body["usage"]
-            if "input_tokens" in usage:
+            if output_tokens := usage.get("outputTokens"):
                 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"]
+                    GEN_AI_USAGE_OUTPUT_TOKENS,
+                    output_tokens,
                 )
-        if "stop_reason" in response_body:
-            span.set_attribute(
-                GEN_AI_RESPONSE_FINISH_REASONS, [response_body["stop_reason"]]
-            )
 
-    # pylint: disable=no-self-use
-    def _handle_cohere_command_response(
-        self, span: Span, response_body: dict[str, Any]
-    ):
-        # Output tokens: Approximate from the response text
-        if "text" in response_body:
-            span.set_attribute(
-                GEN_AI_USAGE_OUTPUT_TOKENS,
-                math.ceil(len(response_body["text"]) / 6),
-            )
-        if "finish_reason" in response_body:
+        if stop_reason := result.get("stopReason"):
             span.set_attribute(
                 GEN_AI_RESPONSE_FINISH_REASONS,
-                [response_body["finish_reason"]],
-            )
-
-    # pylint: disable=no-self-use
-    def _handle_ai21_jamba_response(
-        self, span: Span, response_body: dict[str, Any]
-    ):
-        if "usage" in response_body:
-            usage = response_body["usage"]
-            if "prompt_tokens" in usage:
-                span.set_attribute(
-                    GEN_AI_USAGE_INPUT_TOKENS, usage["prompt_tokens"]
-                )
-            if "completion_tokens" in usage:
-                span.set_attribute(
-                    GEN_AI_USAGE_OUTPUT_TOKENS, usage["completion_tokens"]
-                )
-        if "choices" in response_body:
-            choices = response_body["choices"][0]
-            if "finish_reason" in choices:
-                span.set_attribute(
-                    GEN_AI_RESPONSE_FINISH_REASONS, [choices["finish_reason"]]
-                )
-
-    # pylint: disable=no-self-use
-    def _handle_meta_llama_response(
-        self, span: Span, response_body: dict[str, Any]
-    ):
-        if "prompt_token_count" in response_body:
-            span.set_attribute(
-                GEN_AI_USAGE_INPUT_TOKENS, response_body["prompt_token_count"]
-            )
-        if "generation_token_count" in response_body:
-            span.set_attribute(
-                GEN_AI_USAGE_OUTPUT_TOKENS,
-                response_body["generation_token_count"],
-            )
-        if "stop_reason" in response_body:
-            span.set_attribute(
-                GEN_AI_RESPONSE_FINISH_REASONS, [response_body["stop_reason"]]
-            )
-
-    # pylint: disable=no-self-use
-    def _handle_mistral_mistral_response(
-        self, span: Span, response_body: dict[str, Any]
-    ):
-        if "outputs" in response_body:
-            outputs = response_body["outputs"][0]
-            if "text" in outputs:
-                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"]]
+                [stop_reason],
             )

instrumentation/opentelemetry-instrumentation-botocore/test-requirements.txt renamed to instrumentation/opentelemetry-instrumentation-botocore/test-requirements-0.txt

@@ -19,6 +19,7 @@ pluggy==1.5.0
 py-cpuinfo==9.0.0
 pycparser==2.21
 pytest==7.4.4
+pytest-vcr==1.0.2
 python-dateutil==2.8.2
 pytz==2024.1
 PyYAML==6.0.1

instrumentation/opentelemetry-instrumentation-botocore/test-requirements-1.txt

@@ -0,0 +1,39 @@
+asgiref==3.8.1
+aws-xray-sdk==2.12.1
+boto3==1.35.56
+botocore==1.35.56
+certifi==2024.7.4
+cffi==1.17.0
+charset-normalizer==3.3.2
+cryptography==43.0.1
+Deprecated==1.2.14
+docker==7.0.0
+idna==3.7
+iniconfig==2.0.0
+Jinja2==3.1.4
+jmespath==1.0.1
+MarkupSafe==2.1.5
+moto==5.0.9
+packaging==24.0
+pluggy==1.5.0
+py-cpuinfo==9.0.0
+pycparser==2.21
+pytest==7.4.4
+pytest-vcr==1.0.2
+python-dateutil==2.8.2
+pytz==2024.1
+PyYAML==6.0.1
+requests==2.32.3
+responses==0.25.0
+s3transfer==0.10.0
+six==1.16.0
+tomli==2.0.1
+typing_extensions==4.12.2
+urllib3==1.26.19
+Werkzeug==3.0.6
+wrapt==1.16.0
+xmltodict==0.13.0
+zipp==3.19.2
+-e opentelemetry-instrumentation
+-e propagator/opentelemetry-propagator-aws-xray
+-e instrumentation/opentelemetry-instrumentation-botocore

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

@@ -0,0 +1,88 @@
+interactions:
+- request:
+    body: |-
+      {
+        "messages": [
+          {
+            "role": "user",
+            "content": [
+              {
+                "text": "Say this is a test"
+              }
+            ]
+          }
+        ],
+        "inferenceConfig": {
+          "maxTokens": 100
+        }
+      }
+    headers:
+      Content-Length:
+      - '116'
+      Content-Type:
+      - !!binary |
+        YXBwbGljYXRpb24vanNvbg==
+      User-Agent:
+      - !!binary |
+        Qm90bzMvMS4zNS41NiBtZC9Cb3RvY29yZSMxLjM1LjU2IHVhLzIuMCBvcy9saW51eCM2LjEuMC0x
+        MDM0LW9lbSBtZC9hcmNoI3g4Nl82NCBsYW5nL3B5dGhvbiMzLjEwLjEyIG1kL3B5aW1wbCNDUHl0
+        aG9uIGNmZy9yZXRyeS1tb2RlI2xlZ2FjeSBCb3RvY29yZS8xLjM1LjU2
+      X-Amz-Date:
+      - !!binary |
+        MjAyNDEyMzBUMTY0NjU1Wg==
+      X-Amz-Security-Token:
+      - test_aws_security_token
+      X-Amzn-Trace-Id:
+      - !!binary |
+        Um9vdD0xLWRjYjVjZTc2LWZjMjkyNWZmMTRhZGEzYTgzNDBiNjc3ZTtQYXJlbnQ9ZjU1MTk0MWRm
+        OGM5MGFjMTtTYW1wbGVkPTE=
+      amz-sdk-invocation-id:
+      - !!binary |
+        YzdiMmUzNGMtYThmYi00YzczLTlhNzMtZDIxNzFkMGY3NzAw
+      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
+  response:
+    body:
+      string: |-
+        {
+          "metrics": {
+            "latencyMs": 585
+          },
+          "output": {
+            "message": {
+              "content": [
+                {
+                  "text": "This is a test."
+                }
+              ],
+              "role": "assistant"
+            }
+          },
+          "stopReason": "end_turn",
+          "usage": {
+            "inputTokens": 8,
+            "outputTokens": 9,
+            "totalTokens": 17
+          }
+        }
+    headers:
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '194'
+      Content-Type:
+      - application/json
+      Date:
+      - Mon, 30 Dec 2024 16:46:56 GMT
+      Set-Cookie: test_set_cookie
+      x-amzn-RequestId:
+      - 4d340a17-68b2-4727-b0df-a98dcd3a5af1
+    status:
+      code: 200
+      message: OK
+version: 1