fix(llmobs): bedrock converse tracing is resistant to modifying the stream (#13659)

Fixes an issue where modifying chunks returned from a bedrock stream
impacted our tracing of those chunks. Instead, our tracing should
reflect the original response returned regardless of whether or not it
was modified.

This is relevant when libraries like langchain delete data from the raw
streamed response (e.g.
[popping](https://github.com/langchain-ai/langchain-aws/blob/40abb584979a349019d89bbf1cba7d8c56d23664/libs/aws/langchain_aws/chat_models/bedrock_converse.py#L995)
`usageMetadata`).

The fix uses an approach where we immediately process stream chunks as
they are iterated over, instead of waiting until the entire stream has
finished.

The data flow is like this:
1. a streamed chunk is read from `TracedBotocoreConverseStream`
2. we send that chunk to `_output_stream_processor`, which reads all the
relevant data from that chunk and builds the final output messages,
token usage, metadata. This should block until we reach the next yield,
at which point we've read all the data we needed for this chunk. The
parsing logic is **unchanged** from the previous helper we used to parse
the stream stream, except this method is now a generator.
4. we yield the chunk back to the user

In terms of manual testing, i've verified that it fixes the langchain x
bedrock converse streaming issue
<img width="908" alt="image"
src="https://github.com/user-attachments/assets/e4d1e4e5-3d1c-4f2d-a822-f7c7cdb169a6"
/>

if this logic looks good, it may be a pattern we will want to implement
for our other integrations as well

## Checklist
- [ ] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [ ] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

---------

Co-authored-by: lievan <evan.li@datadoqhq.com>

Author: 2 people

ddtrace/contrib/internal/botocore/services/bedrock.py

@@ -111,14 +111,15 @@ class TracedBotocoreConverseStream(wrapt.ObjectProxy):
 
     def __init__(self, wrapped, ctx: core.ExecutionContext):
         super().__init__(wrapped)
-        self._stream_chunks = []
         self._execution_ctx = ctx
 
     def __iter__(self):
+        stream_processor = self._execution_ctx["bedrock_integration"]._converse_output_stream_processor()
         exception_raised = False
         try:
+            next(stream_processor)
             for chunk in self.__wrapped__:
-                self._stream_chunks.append(chunk)
+                stream_processor.send(chunk)
                 yield chunk
         except Exception:
             core.dispatch("botocore.patched_bedrock_api_call.exception", [self._execution_ctx, sys.exc_info()])
@@ -127,7 +128,7 @@ def __iter__(self):
         finally:
             if exception_raised:
                 return
-            core.dispatch("botocore.bedrock.process_response_converse", [self._execution_ctx, self._stream_chunks])
+            core.dispatch("botocore.bedrock.process_response_converse", [self._execution_ctx, stream_processor])
 
 
 def safe_token_count(token_count) -> Optional[int]:

ddtrace/llmobs/_integrations/bedrock.py

@@ -1,5 +1,6 @@
 from typing import Any
 from typing import Dict
+from typing import Generator
 from typing import List
 from typing import Optional
 from typing import Tuple
@@ -94,11 +95,18 @@ def _llmobs_set_tags(
             if ctx["resource"] == "Converse":
                 output_messages = self._extract_output_message_for_converse(response)
             elif ctx["resource"] == "ConverseStream":
-                (
-                    output_messages,
-                    additional_metadata,
-                    streamed_usage_metrics,
-                ) = self._extract_output_message_for_converse_stream(response)
+                """
+                At this point, we signal to `_converse_output_stream_processor` that we're done with the stream
+                and ready to get the final results. This causes `_converse_output_stream_processor` to break out of the
+                while loop, do some final processing, and return the final results.
+                """
+                try:
+                    response.send(None)
+                except StopIteration as e:
+                    output_messages, additional_metadata, streamed_usage_metrics = e.value
+                finally:
+                    response.close()
+
                 metadata.update(additional_metadata)
                 usage_metrics.update(streamed_usage_metrics)
             else:
@@ -205,11 +213,16 @@ def _extract_output_message_for_converse(response: Dict[str, Any]):
         return get_messages_from_converse_content(role, content)
 
     @staticmethod
-    def _extract_output_message_for_converse_stream(
-        streamed_body: List[Dict[str, Any]]
-    ) -> Tuple[List[Dict[str, Any]], Dict[str, str], Dict[str, int]]:
+    def _converse_output_stream_processor() -> (
+        Generator[
+            None,
+            Dict[str, Any],
+            Tuple[List[Dict[str, Any]], Dict[str, str], Dict[str, int]],
+        ]
+    ):
         """
-        Extract output messages from streamed converse responses.
+        Listens for output chunks from a converse streamed response and builds a
+        list of output messages, usage metrics, and metadata.
 
         Converse stream response comes in chunks. The chunks we care about are:
         - a message start/stop event, or
@@ -229,7 +242,9 @@ def _extract_output_message_for_converse_stream(
 
         current_message: Optional[Dict[str, Any]] = None
 
-        for chunk in streamed_body:
+        chunk = yield
+
+        while chunk is not None:
             if "metadata" in chunk and "usage" in chunk["metadata"]:
                 usage = chunk["metadata"]["usage"]
                 for token_type in ("input", "output", "total"):
@@ -276,6 +291,8 @@ def _extract_output_message_for_converse_stream(
                 )
                 current_message = None
 
+            chunk = yield
+
         # Handle the case where we didn't receive an explicit message stop event
         if current_message is not None and current_message.get("content_block_indicies"):
             messages.append(

releasenotes/notes/fix-bedrock-stream-0bd52763872967a5.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    LLM Observability: This fix resolves an issue where modifying bedrock converse streamed chunks caused traced spans to show modified content.

tests/contrib/botocore/test_bedrock_llmobs.py

@@ -349,3 +349,55 @@ def test_llmobs_converse_stream(cls, bedrock_client, request_vcr, mock_tracer, l
             },
             tags={"service": "aws.bedrock-runtime", "ml_app": "<ml-app-name>"},
         )
+
+    @pytest.mark.skipif(BOTO_VERSION < (1, 34, 131), reason="Converse API not available until botocore 1.34.131")
+    def test_llmobs_converse_modified_stream(cls, bedrock_client, request_vcr, mock_tracer, llmobs_events):
+        """
+        Verify that LLM Obs tracing works even if stream chunks are modified mid-stream.
+        """
+        output_msg = ""
+        request_params = create_bedrock_converse_request(**bedrock_converse_args_with_system_and_tool)
+        with request_vcr.use_cassette("bedrock_converse_stream.yaml"):
+            response = bedrock_client.converse_stream(**request_params)
+            for chunk in response["stream"]:
+                if "contentBlockDelta" in chunk and "delta" in chunk["contentBlockDelta"]:
+                    if "text" in chunk["contentBlockDelta"]["delta"]:
+                        output_msg += chunk["contentBlockDelta"]["delta"]["text"]
+                # delete keys from streamed chunk
+                [chunk.pop(key) for key in list(chunk.keys())]
+
+        span = mock_tracer.pop_traces()[0][0]
+        assert len(llmobs_events) == 1
+
+        assert llmobs_events[0] == _expected_llmobs_llm_span_event(
+            span,
+            model_name="claude-3-sonnet-20240229-v1:0",
+            model_provider="anthropic",
+            input_messages=[
+                {"role": "system", "content": request_params.get("system")[0]["text"]},
+                {"role": "user", "content": request_params.get("messages")[0].get("content")[0].get("text")},
+            ],
+            output_messages=[
+                {
+                    "role": "assistant",
+                    "content": output_msg,
+                    "tool_calls": [
+                        {
+                            "arguments": {"concept": "distributed tracing"},
+                            "name": "fetch_concept",
+                            "tool_id": mock.ANY,
+                        }
+                    ],
+                }
+            ],
+            metadata={
+                "temperature": request_params.get("inferenceConfig", {}).get("temperature"),
+                "max_tokens": request_params.get("inferenceConfig", {}).get("maxTokens"),
+            },
+            token_metrics={
+                "input_tokens": 259,
+                "output_tokens": 64,
+                "total_tokens": 323,
+            },
+            tags={"service": "aws.bedrock-runtime", "ml_app": "<ml-app-name>"},
+        )