scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 114

feat: make streaming chunk size configurable in bedrock converse and invoke handlers (#17357)

Browse source
This commit is contained in:
Deepak Tammali 2025-12-03 11:44:48 +05:30 • committed by GitHub
parent 43dd9e4a90
commit e289f5e454
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 21 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/llms/bedrock/chat/converse_handler.py
View file

@ -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,

20
litellm/llms/bedrock/chat/invoke_handler.py
View file

@ -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",