From 60840ea292dc391d3862c9ccc8cf16d0f7d5e5dc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Dinis=20Ferreira?= Date: Thu, 22 Jan 2026 05:57:14 +0100 Subject: [PATCH] fix(bedrock): correct streaming choice index for tool calls (#19506) Bedrock's contentBlockIndex identifies content blocks within a message (text=0, tool_call=1), not OpenAI's choice index (which varies with n>1). This caused OpenAI SDK's ChatCompletionAccumulator to fail when tool call chunks arrived on index 1 while finish_reason arrived on index 0. Bedrock doesn't support n>1 (no such parameter exists): https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InferenceConfiguration.html OpenAI choice index spec: https://platform.openai.com/docs/api-reference/chat/streaming --- litellm/llms/bedrock/chat/invoke_handler.py | 8 +- .../chat/test_streaming_choice_index.py | 114 ++++++++++++++++++ 2 files changed, 118 insertions(+), 4 deletions(-) create mode 100644 tests/test_litellm/llms/bedrock/chat/test_streaming_choice_index.py diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 032283a2d2..c479de1209 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -1502,7 +1502,7 @@ class AWSEventStreamDecoder: ] ] = None - index = int(chunk_data.get("contentBlockIndex", 0)) + content_block_index = int(chunk_data.get("contentBlockIndex", 0)) if "start" in chunk_data: start_obj = ContentBlockStartEvent(**chunk_data["start"]) tool_use, provider_specific_fields, thinking_blocks = ( @@ -1516,11 +1516,11 @@ class AWSEventStreamDecoder: provider_specific_fields, reasoning_content, thinking_blocks, - ) = self._handle_converse_delta_event(delta_obj, index) + ) = self._handle_converse_delta_event(delta_obj, content_block_index) elif ( "contentBlockIndex" in chunk_data ): # stop block, no 'start' or 'delta' object - tool_use = self._handle_converse_stop_event(index) + tool_use = self._handle_converse_stop_event(content_block_index) elif "stopReason" in chunk_data: finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop")) elif "usage" in chunk_data: @@ -1534,7 +1534,7 @@ class AWSEventStreamDecoder: choices=[ StreamingChoices( finish_reason=finish_reason, - index=index, + index=0, # Always 0 - Bedrock never returns multiple choices delta=Delta( content=text, role="assistant", diff --git a/tests/test_litellm/llms/bedrock/chat/test_streaming_choice_index.py b/tests/test_litellm/llms/bedrock/chat/test_streaming_choice_index.py new file mode 100644 index 0000000000..7a28429fda --- /dev/null +++ b/tests/test_litellm/llms/bedrock/chat/test_streaming_choice_index.py @@ -0,0 +1,114 @@ +""" +Test that Bedrock streaming responses always use choice index 0, +regardless of contentBlockIndex value. + +Bedrock's contentBlockIndex identifies content blocks within a message (e.g., +text=0, toolUse=1), NOT parallel completions. Since Bedrock doesn't support +n > 1, all chunks must use choice index 0. + +References: +- Bedrock InferenceConfiguration (no n parameter): + https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InferenceConfiguration.html +- OpenAI choice.index (for n > 1): + https://platform.openai.com/docs/api-reference/chat/object +""" + +from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + + +class TestBedrockStreamingChoiceIndex: + """Test that all streaming chunks use choice index 0.""" + + def test_tool_call_chunk_uses_choice_index_zero(self): + """ + Core regression test: tool call chunks must use choice index 0, + not contentBlockIndex (which is 1 for tool calls). + + This was the bug - contentBlockIndex was incorrectly used as choice.index, + breaking OpenAI SDK's ChatCompletionAccumulator. + """ + handler = AWSEventStreamDecoder(model="anthropic.claude-3-sonnet-20240229-v1:0") + + # First, simulate a tool use start event on contentBlockIndex 1 + start_chunk = { + "start": { + "toolUse": { + "toolUseId": "tooluse_abc123", + "name": "get_weather", + } + }, + "contentBlockIndex": 1, # Tool calls are on index 1 + } + + start_result = handler.converse_chunk_parser(start_chunk) + + # Choice index should be 0, NOT contentBlockIndex (1) + assert start_result.choices[0].index == 0 + assert start_result.choices[0].delta.tool_calls is not None + assert start_result.choices[0].delta.tool_calls[0]["id"] == "tooluse_abc123" + + # Now simulate tool use delta on contentBlockIndex 1 + delta_chunk = { + "delta": { + "toolUse": { + "input": '{"location": "San Francisco"}' + } + }, + "contentBlockIndex": 1, # Tool calls are on index 1 + } + + delta_result = handler.converse_chunk_parser(delta_chunk) + + # Choice index should still be 0, NOT contentBlockIndex (1) + assert delta_result.choices[0].index == 0 + assert delta_result.choices[0].delta.tool_calls is not None + assert delta_result.choices[0].delta.tool_calls[0]["function"]["arguments"] == '{"location": "San Francisco"}' + + def test_mixed_content_blocks_all_use_choice_index_zero(self): + """ + Integration test simulating a realistic streaming session: + text (contentBlockIndex=0) → tool call (contentBlockIndex=1) → finish. + + All chunks must have choice.index=0 for OpenAI SDK compatibility. + """ + handler = AWSEventStreamDecoder(model="anthropic.claude-3-sonnet-20240229-v1:0") + + # Chunk 1: Text on contentBlockIndex 0 + text_chunk = { + "delta": {"text": "Let me check the weather."}, + "contentBlockIndex": 0, + } + result1 = handler.converse_chunk_parser(text_chunk) + assert result1.choices[0].index == 0, "Text chunk should have index=0" + + # Chunk 2: Tool call start on contentBlockIndex 1 + tool_start_chunk = { + "start": { + "toolUse": { + "toolUseId": "tool_xyz", + "name": "get_weather", + } + }, + "contentBlockIndex": 1, + } + result2 = handler.converse_chunk_parser(tool_start_chunk) + assert result2.choices[0].index == 0, "Tool start should have index=0, not contentBlockIndex=1" + + # Chunk 3: Tool call delta on contentBlockIndex 1 + tool_delta_chunk = { + "delta": { + "toolUse": { + "input": '{"city": "NYC"}' + } + }, + "contentBlockIndex": 1, + } + result3 = handler.converse_chunk_parser(tool_delta_chunk) + assert result3.choices[0].index == 0, "Tool delta should have index=0, not contentBlockIndex=1" + + # Chunk 4: Finish reason + finish_chunk = { + "stopReason": "tool_use", + } + result4 = handler.converse_chunk_parser(finish_chunk) + assert result4.choices[0].index == 0, "Finish reason should have index=0"