From ff089115e3e5f810afcd86bc8057bdbc05e201b4 Mon Sep 17 00:00:00 2001 From: zkllll2002 Date: Thu, 23 Jul 2026 09:56:48 +0800 Subject: [PATCH] chore(bedrock): polish stream comments --- litellm/llms/bedrock/chat/invoke_handler.py | 16 +++------------- .../bedrock/chat/test_converse_transformation.py | 2 +- .../llms/bedrock/chat/test_invoke_handler.py | 4 ++-- 3 files changed, 6 insertions(+), 16 deletions(-) diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index f33fee8d908..5c380891936 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -1283,9 +1283,7 @@ class AWSEventStreamDecoder: self.model = model self.parser = EventStreamJSONParser() - # Bedrock can stream multiple content blocks concurrently. Keep the - # state keyed by contentBlockIndex so a later block start cannot change - # how an earlier block's delta is interpreted. + # Bedrock associates streamed deltas with content blocks by this index. self._content_blocks_by_content_block_index: dict[int, list[ContentBlockDeltaEvent]] = {} self._tool_names_by_content_block_index: dict[int, str] = {} self.tool_calls_index: Optional[int] = None @@ -1380,8 +1378,7 @@ class AWSEventStreamDecoder: response_tool_name = get_bedrock_tool_name(response_tool_name=_response_tool_name) self._tool_names_by_content_block_index[content_block_index] = response_tool_name - # When json_mode is True, suppress the internal json_tool_call - # and convert its content to text in delta events instead + # response_format is represented by an internal tool call. if self.json_mode is True and response_tool_name == RESPONSE_FORMAT_TOOL_NAME: return tool_use, provider_specific_fields, thinking_blocks @@ -1426,8 +1423,6 @@ class AWSEventStreamDecoder: if "text" in delta_obj: text = delta_obj["text"] elif "toolUse" in delta_obj: - # When json_mode is True and this is the internal json_tool_call, - # convert tool input to text content instead of tool call arguments if ( self.json_mode is True and self._tool_names_by_content_block_index.get(index) == RESPONSE_FORMAT_TOOL_NAME @@ -1469,8 +1464,6 @@ class AWSEventStreamDecoder: tool_use: Optional[ChatCompletionToolCallChunk] = None tool_name = self._tool_names_by_content_block_index.pop(index, None) - # If the ending block was the internal json_tool_call, skip emitting - # the empty-args tool chunk. if self.json_mode is True and tool_name == RESPONSE_FORMAT_TOOL_NAME: self._content_blocks_by_content_block_index.pop(index, None) return tool_use @@ -1744,10 +1737,7 @@ class MockResponseIterator: # for returning ai21 streaming responses """ tool_use: Optional[ChatCompletionToolCallChunk] = None if self.json_mode is True and tool_calls is not None: - # ``json_mode`` is also set when response_format is emulated with the - # internal ``json_tool_call``. A fake stream can still contain a real - # user tool call, so converting the first tool unconditionally loses - # the protocol signal needed by streaming SDK clients to execute it. + # Bedrock represents response_format as an internal json_tool_call. json_mode_tool_calls = [ tool_call for tool_call in tool_calls if tool_call["function"].get("name") == RESPONSE_FORMAT_TOOL_NAME ] diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 196b56f74a4..a23de452eac 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -4517,7 +4517,7 @@ def test_streaming_filters_json_tool_call_with_real_tools(): def test_streaming_interleaved_json_tool_call_does_not_hide_real_tool_delta(): - """Keep each Bedrock content block's tool state separate while streaming.""" + """Preserve a real tool delta when an internal JSON block starts later.""" from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder decoder = AWSEventStreamDecoder(model="test-model", json_mode=True) 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 893ab88c376..b9de2302484 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -36,7 +36,7 @@ def test_transform_thinking_blocks_with_redacted_content(): def test_fake_stream_json_mode_preserves_regular_tool_use(): - """A fake stream must not turn a user tool call into JSON text.""" + """Preserve a regular tool call in a fake stream with JSON mode enabled.""" response = litellm.ModelResponse( id="chatcmpl-test", created=0, @@ -76,7 +76,7 @@ def test_fake_stream_json_mode_preserves_regular_tool_use(): def test_fake_stream_json_mode_converts_only_response_format_tool(): - """The internal response-format tool remains regular JSON text.""" + """Convert only the response-format helper tool to JSON content.""" response = litellm.ModelResponse( id="chatcmpl-test", created=0,