diff --git a/litellm/llms/bedrock_mantle/chat/handler.py b/litellm/llms/bedrock_mantle/chat/handler.py new file mode 100644 index 00000000000..390f5e61696 --- /dev/null +++ b/litellm/llms/bedrock_mantle/chat/handler.py @@ -0,0 +1,32 @@ +""" +Streaming handler for Amazon Bedrock Mantle (OpenAI-compatible API). + +Bedrock Mantle emits a fresh ``id`` on every SSE chunk, which violates the +OpenAI streaming contract that all chunks in one response share a single +``id``. Clients that validate this (e.g. the openai-go SDK's +``ChatCompletionAccumulator``) drop every chunk after the first, losing +content and tool-call arguments. We pin the id from the first chunk that +carries one and reuse it for the rest of the stream. +""" + +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, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False): + super().__init__(streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode) + self._response_id: Optional[str] = None + + def chunk_parser(self, chunk: dict) -> ModelResponseStream: + parsed = super().chunk_parser(chunk) + raw_id = chunk.get("id") + if self._response_id is None and raw_id: + self._response_id = raw_id + if self._response_id is not None: + parsed.id = self._response_id + return parsed diff --git a/litellm/llms/bedrock_mantle/chat/transformation.py b/litellm/llms/bedrock_mantle/chat/transformation.py index 64d7ef2bed6..96d8ca68f7a 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.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..7d350bf6464 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,56 @@ def test_gemma_4_bedrock_mantle_model_metadata( ) +def test_streaming_handler_pins_first_chunk_id(): + """Bedrock Mantle emits a fresh id on every SSE chunk; the OpenAI contract + requires all chunks in a response to share one id, so the handler must pin + the first chunk's id onto every subsequent chunk.""" + from litellm.llms.bedrock_mantle.chat.handler import ( + BedrockMantleChatCompletionStreamingHandler, + ) + + handler = BedrockMantleChatCompletionStreamingHandler( + streaming_response=iter([]), sync_stream=True + ) + + raw_chunks = [ + { + "id": f"chatcmpl-{i}", + "created": 1, + "model": "openai.gpt-5.5", + "choices": [{"index": 0, "delta": {"content": token}}], + } + for i, token in enumerate(["Hel", "lo", "!"]) + ] + + parsed_ids = [handler.chunk_parser(chunk).id for chunk in raw_chunks] + + assert parsed_ids == ["chatcmpl-0", "chatcmpl-0", "chatcmpl-0"] + + +def test_streaming_handler_pins_id_when_first_chunk_lacks_one(): + """Some backends omit id on an initial role-only chunk; the handler should + pin the first id it actually sees and keep it stable afterwards.""" + from litellm.llms.bedrock_mantle.chat.handler import ( + BedrockMantleChatCompletionStreamingHandler, + ) + + handler = BedrockMantleChatCompletionStreamingHandler( + streaming_response=iter([]), sync_stream=True + ) + + raw_chunks = [ + {"created": 1, "model": "openai.gpt-5.5", "choices": [{"index": 0, "delta": {"role": "assistant"}}]}, + {"id": "chatcmpl-real", "created": 1, "model": "openai.gpt-5.5", "choices": [{"index": 0, "delta": {"content": "Hi"}}]}, + {"id": "chatcmpl-other", "created": 1, "model": "openai.gpt-5.5", "choices": [{"index": 0, "delta": {"content": "!"}}]}, + ] + + parsed_ids = [handler.chunk_parser(chunk).id for chunk in raw_chunks] + + assert parsed_ids[1] == "chatcmpl-real" + assert parsed_ids[2] == "chatcmpl-real" + + @pytest.mark.parametrize( "model_id", [