diff --git a/litellm/llms/bedrock_mantle/chat/streaming_handler.py b/litellm/llms/bedrock_mantle/chat/streaming_handler.py new file mode 100644 index 00000000000..a2cbc4deee8 --- /dev/null +++ b/litellm/llms/bedrock_mantle/chat/streaming_handler.py @@ -0,0 +1,37 @@ +"""Streaming handler for the Amazon Bedrock Mantle OpenAI-compatible surface. + +Mantle emits streaming ``tool_calls[].index`` values that start at 1 for the +first tool call, whereas the OpenAI spec (and every client that aggregates the +deltas, including the OpenAI SDK) expects a 0-based index. This handler shifts +the provider's indices down so the first tool call reported in a stream is +index 0, leaving already-0-based streams untouched. +""" + +from typing import Optional + +from litellm.llms.openai.chat.gpt_transformation import ( + OpenAIChatCompletionStreamingHandler, +) +from litellm.types.utils import ModelResponseStream + + +class BedrockMantleChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + # Offset of the first tool_call index seen in the stream; every index is + # rebased against it so the first tool call becomes index 0. + self._tool_call_index_offset: Optional[int] = None + + def chunk_parser(self, chunk: dict) -> ModelResponseStream: + parsed = super().chunk_parser(chunk) + for choice in parsed.choices: + tool_calls = choice.delta.tool_calls + if not tool_calls: + continue + offset = self._tool_call_index_offset + if offset is None: + offset = min(tool_call.index for tool_call in tool_calls) + self._tool_call_index_offset = offset + for tool_call in tool_calls: + tool_call.index -= offset + return parsed diff --git a/litellm/llms/bedrock_mantle/chat/transformation.py b/litellm/llms/bedrock_mantle/chat/transformation.py index 64d7ef2bed6..9e1dc6705a8 100644 --- a/litellm/llms/bedrock_mantle/chat/transformation.py +++ b/litellm/llms/bedrock_mantle/chat/transformation.py @@ -111,11 +111,11 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig): sync_stream: bool, json_mode: bool | None = False, ) -> Any: - from litellm.llms.openai.chat.gpt_transformation import ( - OpenAIChatCompletionStreamingHandler, + from litellm.llms.bedrock_mantle.chat.streaming_handler import ( + BedrockMantleChatCompletionStreamingHandler, ) - return OpenAIChatCompletionStreamingHandler( + return BedrockMantleChatCompletionStreamingHandler( streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode, diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index cd775abf136..15aa351808b 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -666,6 +666,85 @@ def test_gemma_4_bedrock_mantle_model_metadata( ) +def _mantle_tool_call_chunk(index: int, *, tool_id: str, name: str, args: str) -> str: + return "data: " + json.dumps( + { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 1733529600, + "model": "openai.gpt-5.5", + "choices": [ + { + "index": 0, + "delta": { + "tool_calls": [ + { + "index": index, + "id": tool_id, + "type": "function", + "function": {"name": name, "arguments": args}, + } + ] + }, + "finish_reason": None, + } + ], + } + ) + + +def _collect_tool_call_indices(chunks): + from litellm.llms.bedrock_mantle.chat.streaming_handler import ( + BedrockMantleChatCompletionStreamingHandler, + ) + + from litellm.types.utils import ModelResponseStream + + handler = BedrockMantleChatCompletionStreamingHandler( + streaming_response=iter(chunks), + sync_stream=True, + ) + return [ + tool_call.index + for parsed in handler + if isinstance(parsed, ModelResponseStream) + for choice in parsed.choices + for tool_call in (choice.delta.tool_calls or []) + ] + + +class TestBedrockMantleStreamingToolCallIndex: + """Regression tests for https://github.com/BerriAI/litellm/issues/32759 + + Mantle streams the first tool call with index 1; the OpenAI spec (and any + client aggregating the deltas) expects the first index to be 0. + """ + + def test_first_tool_call_index_rebased_to_zero(self): + chunks = [ + _mantle_tool_call_chunk(1, tool_id="call_a", name="get_weather", args=""), + _mantle_tool_call_chunk(1, tool_id="call_a", name="get_weather", args='{"city": "NYC"}'), + "data: [DONE]", + ] + assert _collect_tool_call_indices(chunks) == [0, 0] + + def test_parallel_tool_call_indices_rebased_contiguously(self): + chunks = [ + _mantle_tool_call_chunk(1, tool_id="call_a", name="get_weather", args="{}"), + _mantle_tool_call_chunk(2, tool_id="call_b", name="get_time", args="{}"), + "data: [DONE]", + ] + assert _collect_tool_call_indices(chunks) == [0, 1] + + def test_already_zero_based_stream_is_untouched(self): + chunks = [ + _mantle_tool_call_chunk(0, tool_id="call_a", name="get_weather", args="{}"), + _mantle_tool_call_chunk(1, tool_id="call_b", name="get_time", args="{}"), + "data: [DONE]", + ] + assert _collect_tool_call_indices(chunks) == [0, 1] + + @pytest.mark.parametrize( "model_id", [