diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 53cbafcbe6..71429e4191 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/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, ) diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index 429b1a4389..a415d55021 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/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"