From e289f5e454140042eda28ca27b8797078e2fe9f6 Mon Sep 17 00:00:00 2001 From: Deepak Tammali <45919384+deepaktammali@users.noreply.github.com> Date: Wed, 3 Dec 2025 11:44:48 +0530 Subject: [PATCH] feat: make streaming chunk size configurable in bedrock converse and invoke handlers (#17357) --- litellm/llms/bedrock/chat/converse_handler.py | 10 ++++++++-- litellm/llms/bedrock/chat/invoke_handler.py | 20 ++++++++++++------- 2 files changed, 21 insertions(+), 9 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index fd1f6f0c893..d5bd054118d 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -29,6 +29,7 @@ def make_sync_call( logging_obj: LiteLLMLoggingObject, json_mode: Optional[bool] = False, fake_stream: bool = False, + stream_chunk_size: int = 1024, ): if client is None: client = _get_httpx_client() # Create a new client if none provided @@ -66,7 +67,7 @@ def make_sync_call( ) else: decoder = AWSEventStreamDecoder(model=model) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) + completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) # LOGGING logging_obj.post_call( @@ -102,6 +103,7 @@ class BedrockConverseLLM(BaseAWSLLM): fake_stream: bool = False, json_mode: Optional[bool] = False, api_key: Optional[str] = None, + stream_chunk_size: int = 1024, ) -> CustomStreamWrapper: request_data = await litellm.AmazonConverseConfig()._async_transform_request( model=model, @@ -143,6 +145,7 @@ class BedrockConverseLLM(BaseAWSLLM): logging_obj=logging_obj, fake_stream=fake_stream, json_mode=json_mode, + stream_chunk_size=stream_chunk_size, ) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, @@ -260,6 +263,7 @@ class BedrockConverseLLM(BaseAWSLLM): ): ## SETUP ## stream = optional_params.pop("stream", None) + stream_chunk_size = optional_params.pop("stream_chunk_size", 1024) unencoded_model_id = optional_params.pop("model_id", None) fake_stream = optional_params.pop("fake_stream", False) json_mode = optional_params.get("json_mode", False) @@ -356,7 +360,8 @@ class BedrockConverseLLM(BaseAWSLLM): json_mode=json_mode, fake_stream=fake_stream, credentials=credentials, - api_key=api_key + api_key=api_key, + stream_chunk_size=stream_chunk_size, ) # type: ignore ### ASYNC COMPLETION return self.async_completion( @@ -433,6 +438,7 @@ class BedrockConverseLLM(BaseAWSLLM): logging_obj=logging_obj, json_mode=json_mode, fake_stream=fake_stream, + stream_chunk_size=stream_chunk_size, ) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 7a960fd45d2..5e33a266449 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -192,6 +192,7 @@ async def make_call( fake_stream: bool = False, json_mode: Optional[bool] = False, bedrock_invoke_provider: Optional[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL] = None, + stream_chunk_size: int = 1024, ): try: if client is None: @@ -235,7 +236,7 @@ async def make_call( json_mode=json_mode, ) completion_stream = decoder.aiter_bytes( - response.aiter_bytes(chunk_size=1024) + response.aiter_bytes(chunk_size=stream_chunk_size) ) elif bedrock_invoke_provider == "deepseek_r1": decoder = AmazonDeepSeekR1StreamDecoder( @@ -243,12 +244,12 @@ async def make_call( sync_stream=False, ) completion_stream = decoder.aiter_bytes( - response.aiter_bytes(chunk_size=1024) + response.aiter_bytes(chunk_size=stream_chunk_size) ) else: decoder = AWSEventStreamDecoder(model=model) completion_stream = decoder.aiter_bytes( - response.aiter_bytes(chunk_size=1024) + response.aiter_bytes(chunk_size=stream_chunk_size) ) # LOGGING @@ -281,6 +282,7 @@ def make_sync_call( fake_stream: bool = False, json_mode: Optional[bool] = False, bedrock_invoke_provider: Optional[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL] = None, + stream_chunk_size: int = 1024, ): try: if client is None: @@ -321,16 +323,16 @@ def make_sync_call( sync_stream=True, json_mode=json_mode, ) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) + completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) elif bedrock_invoke_provider == "deepseek_r1": decoder = AmazonDeepSeekR1StreamDecoder( model=model, sync_stream=True, ) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) + completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) else: decoder = AWSEventStreamDecoder(model=model) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) + completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) # LOGGING logging_obj.post_call( @@ -729,6 +731,7 @@ class BedrockLLM(BaseAWSLLM): ## SETUP ## stream = optional_params.pop("stream", None) + stream_chunk_size = optional_params.pop("stream_chunk_size", 1024) provider = self.get_bedrock_invoke_provider(model) modelId = self.get_bedrock_model_id( @@ -1003,6 +1006,7 @@ class BedrockLLM(BaseAWSLLM): headers=prepped.headers, timeout=timeout, client=client, + stream_chunk_size=stream_chunk_size, ) # type: ignore ### ASYNC COMPLETION return self.async_completion( @@ -1048,7 +1052,7 @@ class BedrockLLM(BaseAWSLLM): decoder = AWSEventStreamDecoder(model=model) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) + completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, model=model, @@ -1168,6 +1172,7 @@ class BedrockLLM(BaseAWSLLM): logger_fn=None, headers={}, client: Optional[AsyncHTTPHandler] = None, + stream_chunk_size: int = 1024, ) -> CustomStreamWrapper: # The call is not made here; instead, we prepare the necessary objects for the stream. @@ -1183,6 +1188,7 @@ class BedrockLLM(BaseAWSLLM): messages=messages, logging_obj=logging_obj, fake_stream=True if "ai21" in api_base else False, + stream_chunk_size=stream_chunk_size, ), model=model, custom_llm_provider="bedrock",