fix(bedrock): Ensure consistent chunk IDs in Bedrock streaming responses (#16596)

* ensure consistent chunk IDs in streaming responses

* use native conversationId for consistent stream chunk IDs

litellm/llms/bedrock/chat/invoke_handler.py

@@ -493,9 +493,9 @@ class BedrockLLM(BaseAWSLLM):
                             content=None,
                         )
                         model_response.choices[0].message = _message  # type: ignore
-                        model_response._hidden_params["original_response"] = (
-                            outputText  # allow user to access raw anthropic tool calling response
-                        )
+                        model_response._hidden_params[
+                            "original_response"
+                        ] = outputText  # allow user to access raw anthropic tool calling response
                     if (
                         _is_function_call is True
                         and stream is not None
@@ -793,9 +793,9 @@ class BedrockLLM(BaseAWSLLM):
                     ):  # completion(top_k=3) > anthropic_config(top_k=3) <- allows for dynamic variables to be passed in
                         inference_params[k] = v
                 if stream is True:
-                    inference_params["stream"] = (
-                        True  # cohere requires stream = True in inference params
-                    )
+                    inference_params[
+                        "stream"
+                    ] = True  # cohere requires stream = True in inference params
                 data = json.dumps({"prompt": prompt, **inference_params})
         elif provider == "anthropic":
             if model.startswith("anthropic.claude-3"):
@@ -1184,6 +1184,7 @@ class AWSEventStreamDecoder:
         self.parser = EventStreamJSONParser()
         self.content_blocks: List[ContentBlockDeltaEvent] = []
         self.tool_calls_index: Optional[int] = None
+        self.response_id: Optional[str] = None
 
     def check_empty_tool_call_args(self) -> bool:
         """
@@ -1247,6 +1248,17 @@ class AWSEventStreamDecoder:
 
     def converse_chunk_parser(self, chunk_data: dict) -> ModelResponseStream:
         try:
+            # Capture the conversationId from the first messageStart event
+            # and use it as the consistent ID for all subsequent chunks.
+            if self.response_id is None:
+                if "messageStart" in chunk_data:
+                    conversation_id = chunk_data["messageStart"].get("conversationId")
+                    if conversation_id:
+                        self.response_id = f"chatcmpl-{conversation_id}"
+                else:
+                    # Fallback to generating a UUID if the first chunk is not messageStart
+                    self.response_id = f"chatcmpl-{uuid.uuid4()}"
+
             verbose_logger.debug("\n\nRaw Chunk: {}\n\n".format(chunk_data))
             text = ""
             tool_use: Optional[ChatCompletionToolCallChunk] = None
@@ -1378,6 +1390,7 @@ class AWSEventStreamDecoder:
                         ),
                     )
                 ],
+                id=self.response_id,
                 usage=usage,
                 provider_specific_fields=model_response_provider_specific_fields,
             )

tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py

@@ -1,14 +1,10 @@
-import json
 import os
 import sys
 
-import pytest
-from fastapi.testclient import TestClient
 
 sys.path.insert(
     0, os.path.abspath("../../../../..")
 )  # Adds the parent directory to the system path
-from unittest.mock import MagicMock, patch
 
 from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder
 
@@ -167,3 +163,40 @@ def test_transform_tool_calls_index_with_optional_arg_func():
         tool_call_hunk_dict = tool_call_hunk.model_dump()
         for tool_call in tool_call_hunk_dict["choices"][0]["delta"]["tool_calls"]:
             assert tool_call["index"] == 0
+
+
+def test_bedrock_converse_streaming_consistent_id():
+    """
+    Tests that all chunks in a Bedrock Converse stream response share the same ID,
+    capturing the ID from the initial 'messageStart' event.
+    """
+    # Simulate a realistic Bedrock Converse stream
+    native_conversation_id = "a1b2c3d4-e5f6-7890-1234-56789abcdef0"
+    mock_stream_chunks = [
+        {
+            "messageStart": {
+                "conversationId": native_conversation_id,
+                "role": "assistant",
+            }
+        },
+        {"delta": {"text": "Hello"}, "contentBlockIndex": 0},
+        {"delta": {"text": " world!"}, "contentBlockIndex": 0},
+        {"stopReason": "stop"},
+    ]
+
+    decoder = AWSEventStreamDecoder(model="bedrock/anthropic.claude-3-sonnet-v1:0")
+
+    # Process each chunk and collect the parsed responses
+    parsed_responses = []
+    for chunk in mock_stream_chunks:
+        parsed_responses.append(decoder.converse_chunk_parser(chunk))
+
+    # Verify that all parsed responses have the same, non-null ID derived from the native ID
+    assert len(parsed_responses) > 1, "Should have processed multiple chunks"
+
+    expected_id = f"chatcmpl-{native_conversation_id}"
+
+    for response in parsed_responses:
+        assert (
+            response.id == expected_id
+        ), "All chunk IDs must match the one captured from the messageStart event"