Avoid double copy of streaming body

Author: xrmx

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

@@ -226,19 +226,17 @@ def _invoke_model_on_success(
             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"))
+            # 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)
-            # 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")

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

@@ -14,7 +14,6 @@
 
 from __future__ import annotations
 
-import io
 import json
 from typing import Any
 
@@ -40,12 +39,11 @@ def assert_completion_attributes_from_streaming_body(
     input_tokens = None
     output_tokens = None
     finish_reason = None
-    if response:
+    if response is not None:
         original_body = response["body"]
         body_content = original_body.read()
-        stream = io.BytesIO(body_content)
-        telemetry_content = stream.read()
-        response = json.loads(telemetry_content.decode("utf-8"))
+        response = json.loads(body_content.decode("utf-8"))
+        assert response
 
         if "amazon.titan" in request_model:
             input_tokens = response.get("inputTextTokenCount")