Add changelog a please pylint

open-telemetry · Jan 22, 2025 · 53535c7 · 53535c7
1 parent e09456c
commit 53535c7
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 36 deletions.
diff --git a/CHANGELOG.md b/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
 

diff --git a/...instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/bedrock.py b/...instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/bedrock.py
@@ -198,14 +198,8 @@ def before_service_call(self, span: Span):
         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
-
-        # Converse
+    # 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(
@@ -224,40 +218,53 @@ def on_success(self, span: Span, result: dict[str, Any]):
                 [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()
+
+            # 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)
+            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)
+            # 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 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):
-            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)
-                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)
-                # 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 exc:  # pylint: disable=broad-exception-caught
-                _logger.debug("Error processing response: %s", exc)
-            finally:
-                if original_body is not None:
-                    original_body.close()
+            self._invoke_model_on_success(span, result, model_id)
 
     # pylint: disable=no-self-use
     def _handle_amazon_titan_response(

diff --git a/instrumentation/opentelemetry-instrumentation-botocore/tests/bedrock_utils.py b/instrumentation/opentelemetry-instrumentation-botocore/tests/bedrock_utils.py
@@ -26,6 +26,7 @@
 )
 
 
+# pylint: disable=too-many-branches, too-many-locals
 def assert_completion_attributes_from_streaming_body(
     span: ReadableSpan,
     request_model: str,