From e73f949fbb735b344f75e5f8c395344bf4cfdb54 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 10:27:47 -0700 Subject: [PATCH] fix(params): stop stream_chunk_size reaching provider request bodies (#42664) * fix(params): carry stream_chunk_size through litellm_params instead of provider params * test(integration): fence stream_chunk_size out of every provider request body * test(bedrock): type parametrized stream chunk test params * test(integration): drop the contracts manifest resurrected by the main merge * test(bedrock): type the stream_chunk_size test helpers * test(params): finish AGENTS.md typing pass on stream_chunk_size tests * test(integration): drop the covers marker from the stream_chunk_size wire test --------- Co-authored-by: shrey kharbanda --- .../litellm_core_utils/get_litellm_params.py | 2 + litellm/llms/bedrock/chat/converse_handler.py | 2 +- litellm/main.py | 1 + litellm/types/utils.py | 1 + tests/_support/__init__.py | 0 tests/_support/stream_chunk_size.py | 31 ++ .../providers/test_stream_chunk_size_wire.py | 316 ++++++++++++++++++ .../test_get_litellm_params.py | 8 +- tests/unit/llms/chat/test_converse_handler.py | 162 ++++++++- 9 files changed, 513 insertions(+), 10 deletions(-) create mode 100644 tests/_support/__init__.py create mode 100644 tests/_support/stream_chunk_size.py create mode 100644 tests/integration/providers/test_stream_chunk_size_wire.py diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 112d038039e..f7aaef3a51f 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -130,6 +130,7 @@ def get_litellm_params( api_version: str | None = None, max_retries: int | None = None, litellm_request_debug: bool | None = None, + stream_chunk_size: int | None = None, **kwargs, ) -> dict: _litellm_metadata_dict: Final = litellm_metadata if isinstance(litellm_metadata, dict) else None @@ -192,6 +193,7 @@ def get_litellm_params( "max_retries": max_retries, "use_litellm_proxy": use_litellm_proxy, "litellm_request_debug": litellm_request_debug, + "stream_chunk_size": stream_chunk_size, } # Sparse extraction: only add kwargs keys that are actually present diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 1acac7de14d..3abe7c88fcd 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -278,7 +278,7 @@ class BedrockConverseLLM(BaseAWSLLM): ): ## SETUP ## stream: Final = optional_params.pop("stream", None) - stream_chunk_size: Final = optional_params.pop("stream_chunk_size", None) + stream_chunk_size: Final = litellm_params.get("stream_chunk_size") unencoded_model_id: Final = optional_params.pop("model_id", None) fake_stream = optional_params.pop("fake_stream", False) json_mode: Final = optional_params.get("json_mode", False) diff --git a/litellm/main.py b/litellm/main.py index 7f4b34d28a0..4c40d864169 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5640,6 +5640,7 @@ def completion( max_retries=max_retries, timeout=timeout, litellm_request_debug=kwargs.get("litellm_request_debug", False), + stream_chunk_size=kwargs.get("stream_chunk_size"), tpm=kwargs.get("tpm"), rpm=kwargs.get("rpm"), use_xai_oauth=kwargs.get("use_xai_oauth", False), diff --git a/litellm/types/utils.py b/litellm/types/utils.py index dfc98a9d89d..3f8471cbde1 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4043,6 +4043,7 @@ all_litellm_params = ( "no-log", "base_model", "stream_timeout", + "stream_chunk_size", "supports_system_message", "region_name", "allowed_model_region", diff --git a/tests/_support/__init__.py b/tests/_support/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/_support/stream_chunk_size.py b/tests/_support/stream_chunk_size.py new file mode 100644 index 00000000000..051f552e282 --- /dev/null +++ b/tests/_support/stream_chunk_size.py @@ -0,0 +1,31 @@ +from collections.abc import Mapping +from typing import Final + +import litellm +import pytest +from litellm.integrations.custom_logger import CustomLogger + + +class LitellmParamsRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.seen: tuple[Mapping[str, object], ...] = () + + def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None: + params: Final = kwargs["litellm_params"] + assert isinstance(params, Mapping) + self.seen = (*self.seen, params) + + +def record_litellm_params(monkeypatch: pytest.MonkeyPatch) -> LitellmParamsRecorder: + recorder: Final = LitellmParamsRecorder() + monkeypatch.setattr(litellm, "input_callback", [recorder]) + return recorder + + +def keys_at_every_depth(value: object) -> frozenset[str]: + if isinstance(value, Mapping): + return frozenset(value) | frozenset().union(*(keys_at_every_depth(item) for item in value.values())) + if isinstance(value, (list, tuple)): + return frozenset().union(*(keys_at_every_depth(item) for item in value)) + return frozenset() diff --git a/tests/integration/providers/test_stream_chunk_size_wire.py b/tests/integration/providers/test_stream_chunk_size_wire.py new file mode 100644 index 00000000000..3681da0e3d4 --- /dev/null +++ b/tests/integration/providers/test_stream_chunk_size_wire.py @@ -0,0 +1,316 @@ +import asyncio +import base64 +import json +import os +import struct +import zlib +from collections.abc import Callable, Mapping +from pathlib import Path +from typing import Final + +import litellm +import pytest +from integration._support.upstream import INTERNAL_FIELDS +from integration._support.wire import Reply, Request, wire_server +from tests._support.stream_chunk_size import keys_at_every_depth, record_litellm_params + +TEXT: Final = "wire control" +OPENAI_RESPONSE: Final = { + "id": "chatcmpl-wire", + "object": "chat.completion", + "created": 1, + "model": "gpt-4.1-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": TEXT}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14}, +} +ANTHROPIC_RESPONSE: Final = { + "id": "msg_wire", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": TEXT}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 4}, +} +GEMINI_RESPONSE: Final = { + "candidates": [{"content": {"role": "model", "parts": [{"text": TEXT}]}, "finishReason": "STOP", "index": 0}], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 4, "totalTokenCount": 14}, +} +CONVERSE_RESPONSE: Final = { + "output": {"message": {"role": "assistant", "content": [{"text": TEXT}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 10, "outputTokens": 4, "totalTokens": 14}, + "metrics": {"latencyMs": 1}, +} +OPENAI_STREAM_CHUNKS: Final = ( + { + "id": "chatcmpl-wire", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4.1-mini", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": TEXT}, "finish_reason": None}], + }, + { + "id": "chatcmpl-wire", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4.1-mini", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + }, +) +ANTHROPIC_STREAM_EVENTS: Final = ( + { + "type": "message_start", + "message": { + "id": "msg_wire", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 10, "output_tokens": 1}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": TEXT}}, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}}, + {"type": "message_stop"}, +) +GEMINI_STREAM_CHUNKS: Final = ( + {"candidates": [{"content": {"role": "model", "parts": [{"text": TEXT}]}, "index": 0}]}, + { + "candidates": [{"content": {"role": "model", "parts": [{"text": ""}]}, "finishReason": "STOP", "index": 0}], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 4, "totalTokenCount": 14}, + }, +) +CONVERSE_STREAM_EVENTS: Final = ( + ("contentBlockDelta", {"delta": {"text": TEXT}, "contentBlockIndex": 0}), + ("messageStop", {"stopReason": "end_turn"}), + ("metadata", {"usage": {"inputTokens": 10, "outputTokens": 4, "totalTokens": 14}, "metrics": {"latencyMs": 1}}), +) +NON_STREAM_BODIES: Final = { + "openai": OPENAI_RESPONSE, + "azure": OPENAI_RESPONSE, + "anthropic": ANTHROPIC_RESPONSE, + "gemini": GEMINI_RESPONSE, + "converse": CONVERSE_RESPONSE, + "invoke": ANTHROPIC_RESPONSE, +} +PROVIDERS: Final = ("openai", "azure", "anthropic", "gemini", "converse", "invoke") + + +def _aws_string_header(name: str, value: str) -> bytes: + name_bytes: Final = name.encode() + value_bytes: Final = value.encode() + return struct.pack("!B", len(name_bytes)) + name_bytes + b"\x07" + struct.pack("!H", len(value_bytes)) + value_bytes + + +def _aws_event_frame(event_type: str, payload: Mapping[str, object]) -> bytes: + body: Final = json.dumps(payload, separators=(",", ":")).encode() + headers: Final = ( + _aws_string_header(":event-type", event_type) + + _aws_string_header(":content-type", "application/json") + + _aws_string_header(":message-type", "event") + ) + prelude: Final = struct.pack("!II", 12 + len(headers) + len(body) + 4, len(headers)) + message: Final = prelude + struct.pack("!I", zlib.crc32(prelude) & 0xFFFFFFFF) + headers + body + return message + struct.pack("!I", zlib.crc32(message) & 0xFFFFFFFF) + + +def _sse_reply(frames: tuple[bytes, ...]) -> Reply: + return Reply(chunks=frames, content_type="text/event-stream") + + +def _stream_reply(provider: str) -> Reply: + match provider: + case "openai" | "azure": + return _sse_reply( + tuple( + f"data: {json.dumps(chunk, separators=(',', ':'))}\n\n".encode() for chunk in OPENAI_STREAM_CHUNKS + ) + + (b"data: [DONE]\n\n",) + ) + case "anthropic": + return _sse_reply( + tuple( + f"event: {event['type']}\ndata: {json.dumps(event, separators=(',', ':'))}\n\n".encode() + for event in ANTHROPIC_STREAM_EVENTS + ) + ) + case "gemini": + return _sse_reply( + tuple( + f"data: {json.dumps(chunk, separators=(',', ':'))}\r\n\r\n".encode() + for chunk in GEMINI_STREAM_CHUNKS + ) + ) + case "converse": + return Reply( + chunks=tuple(_aws_event_frame(event_type, payload) for event_type, payload in CONVERSE_STREAM_EVENTS), + content_type="application/vnd.amazon.eventstream", + ) + case "invoke": + return Reply( + chunks=tuple( + _aws_event_frame( + "chunk", {"bytes": base64.b64encode(json.dumps(event, separators=(",", ":")).encode()).decode()} + ) + for event in ANTHROPIC_STREAM_EVENTS + ), + content_type="application/vnd.amazon.eventstream", + ) + + +def _request_parameters(provider: str, wire_url: str) -> dict[str, object]: + common: Final = {"messages": [{"role": "user", "content": "synthetic chunk control"}]} + match provider: + case "openai": + return {**common, "model": "openai/gpt-4.1-mini", "api_key": "synthetic-openai-key", "api_base": wire_url} + case "azure": + return { + **common, + "model": "azure/gpt-4.1-mini", + "api_key": "synthetic-azure-key", + "api_base": wire_url, + "api_version": "2025-01-01-preview", + } + case "anthropic": + return { + **common, + "model": "anthropic/claude-sonnet-4-5", + "api_key": "synthetic-anthropic-key", + "api_base": wire_url, + } + case "gemini": + return { + **common, + "model": "gemini/gemini-2.5-flash", + "api_key": "synthetic-gemini-key", + "api_base": wire_url, + } + case "converse": + return { + **common, + "model": "bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_access_key_id": "fake", + "aws_secret_access_key": "fake", + "aws_region_name": "us-east-1", + "aws_bedrock_runtime_endpoint": wire_url, + } + case "invoke": + return { + **common, + "model": "bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_access_key_id": "fake", + "aws_secret_access_key": "fake", + "aws_region_name": "us-east-1", + "aws_bedrock_runtime_endpoint": wire_url, + } + + +def _expected_target(provider: str, streaming: bool) -> str: + match provider: + case "openai": + return "/chat/completions" + case "azure": + return "/openai/deployments/gpt-4.1-mini/chat/completions?api-version=2025-01-01-preview" + case "anthropic": + return "/v1/messages" + case "gemini": + return ":streamGenerateContent" if streaming else ":generateContent" + case "converse": + return "/converse-stream" if streaming else "/converse" + case "invoke": + return "/invoke-with-response-stream" if streaming else "/invoke" + + +def _at(value: object, *path: str) -> object: + if not path: + return value + assert isinstance(value, Mapping) + return _at(value[path[0]], *path[1:]) + + +def _custom_key(body: Mapping[str, object], provider: str) -> object: + match provider: + case "anthropic": + return _at(body, "extra_body", "custom_provider_key") + case "converse": + return _at(body, "additionalModelRequestFields", "extra_body", "custom_provider_key") + return _at(body, "custom_provider_key") + + +def _peer(provider: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + body: Final = json.loads(request.body) if request.body else {} + streaming: Final = ( + (isinstance(body, dict) and body.get("stream") is True) + or "streamGenerateContent" in request.target + or request.target.endswith(("-stream",)) + ) + expected: Final = _expected_target(provider, streaming) + assert expected in request.target, f"{provider}: expected {expected} in {request.target}" + return _stream_reply(provider) if streaming else Reply(body=json.dumps(NON_STREAM_BODIES[provider]).encode()) + + return respond + + +@pytest.fixture +def provider_wire_environment(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + empty: Final = tmp_path / "empty-aws-config" + empty.write_text("") + for name in tuple(name for name in os.environ if name.startswith("AWS_")): + monkeypatch.delenv(name, raising=False) + for name, value in { + "AWS_CONFIG_FILE": str(empty), + "AWS_SHARED_CREDENTIALS_FILE": str(empty), + "AWS_EC2_METADATA_DISABLED": "true", + "LITELLM_RUST": "false", + }.items(): + monkeypatch.setenv(name, value) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + +@pytest.mark.parametrize("provider", PROVIDERS) +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +async def test_stream_chunk_size_never_reaches_provider_body( + monkeypatch: pytest.MonkeyPatch, + provider_wire_environment: None, + provider: str, + asynchronous: bool, + stream: bool, +) -> None: + recorder: Final = record_litellm_params(monkeypatch) + with wire_server(_peer(provider)) as wire: + parameters: Final = { + **_request_parameters(provider, wire.url), + "stream": stream, + "stream_chunk_size": 64, + "extra_body": {"custom_provider_key": 1}, + "max_tokens": 16, + "timeout": 5, + "num_retries": 0, + } + result: Final = ( + await litellm.acompletion(**parameters) + if asynchronous + else await asyncio.to_thread(litellm.completion, **parameters) + ) + if stream: + chunks: Final = [chunk async for chunk in result] if asynchronous else [chunk for chunk in result] + text: Final = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) + assert text == TEXT + else: + assert result.choices[0].message.content == TEXT + requests: Final = wire.drain() + assert len(requests) == 1 + assert len(recorder.seen) == 1 + assert recorder.seen[0]["stream_chunk_size"] == 64 + body: Final = json.loads(requests[0].body) + keys: Final = keys_at_every_depth(body) + assert "stream_chunk_size" not in keys + assert not INTERNAL_FIELDS.intersection(keys) + assert _custom_key(body, provider) == 1 diff --git a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py index 7fb45e1b092..39bc2688ae0 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py +++ b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py @@ -90,6 +90,10 @@ class TestGetLitellmParamsKwargsExtraction: assert "s3_endpoint_url" not in result_without_s3_kwargs assert "s3_region_name" not in result_without_s3_kwargs + def test_stream_chunk_size_is_carried_as_a_litellm_param(self) -> None: + assert get_litellm_params(stream_chunk_size=64)["stream_chunk_size"] == 64 + assert get_litellm_params()["stream_chunk_size"] is None + def test_s3_credential_kwargs_are_forwarded_for_s3_signing(self): result = get_litellm_params(s3_access_key_id="s3-key", s3_secret_access_key="s3-secret") assert result["s3_access_key_id"] == "s3-key" @@ -265,5 +269,7 @@ class TestMetadataFallsBackToLitellmMetadata: "value, expected", [("true", True), ("false", False), (" TRUE ", True), (True, True), (None, None), ("os.environ/DROP_PARAMS", None)], ) -def test_drop_params_strings_reach_litellm_params_as_flags(value, expected): +def test_drop_params_strings_reach_litellm_params_as_flags( + value: str | bool | None, expected: bool | None +) -> None: assert get_litellm_params(drop_params=value)["drop_params"] is expected diff --git a/tests/unit/llms/chat/test_converse_handler.py b/tests/unit/llms/chat/test_converse_handler.py index 05debee0602..c342a7e1806 100644 --- a/tests/unit/llms/chat/test_converse_handler.py +++ b/tests/unit/llms/chat/test_converse_handler.py @@ -1,4 +1,6 @@ import json +from collections.abc import AsyncIterator +from typing import Final from unittest.mock import AsyncMock, MagicMock import httpx @@ -9,6 +11,11 @@ from litellm.llms.bedrock.chat import BedrockConverseLLM from litellm.llms.bedrock.chat.converse_handler import make_sync_call from litellm.llms.bedrock.common_utils import _get_all_bedrock_regions from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from tests._support.stream_chunk_size import ( + LitellmParamsRecorder, + keys_at_every_depth, + record_litellm_params, +) @@ -68,8 +75,8 @@ class TestBedrockRegionInModelPath: ], ) def test_region_and_model_id_extraction( - self, model, expected_model_id, expected_region - ): + self, model: str, expected_model_id: str, expected_region: str | None + ) -> None: """ Verify that completion() correctly extracts both modelId and aws_region_name from the bedrock/{region}/{model} path format. @@ -139,11 +146,11 @@ class TestBedrockRegionInModelPath: assert optional_params["aws_region_name"] == "eu-west-1" -def _stream_completion_with_spied_iter_bytes(model: str, **kwargs) -> MagicMock: - mock_response = MagicMock() +def _stream_completion_with_spied_iter_bytes(model: str, stream_chunk_size: int | None = None) -> MagicMock: + mock_response: Final = MagicMock() mock_response.status_code = 200 mock_response.iter_bytes = MagicMock(return_value=iter([])) - client = HTTPHandler() + client: Final = HTTPHandler() client.post = MagicMock(return_value=mock_response) litellm.completion( @@ -154,7 +161,7 @@ def _stream_completion_with_spied_iter_bytes(model: str, **kwargs) -> MagicMock: aws_access_key_id="fake", aws_secret_access_key="fake", aws_region_name="us-east-1", - **kwargs, + stream_chunk_size=stream_chunk_size, ) return mock_response.iter_bytes @@ -276,7 +283,7 @@ async def test_async_converse_completion_forwards_bedrock_response_headers(): @pytest.mark.asyncio async def test_async_converse_streaming_forwards_bedrock_response_headers(): - async def _no_bytes(chunk_size=None): + async def _no_bytes(chunk_size: int | None = None) -> AsyncIterator[bytes]: return yield b"" @@ -300,7 +307,7 @@ async def test_async_converse_streaming_forwards_bedrock_response_headers(): assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-def" -def test_completion_plumbs_stream_chunk_size_through_converse(): +def test_completion_plumbs_stream_chunk_size_through_converse() -> None: iter_bytes_spy = _stream_completion_with_spied_iter_bytes( model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" ) @@ -313,6 +320,145 @@ def test_completion_plumbs_stream_chunk_size_through_converse(): iter_bytes_spy.assert_called_once_with(chunk_size=2048) +def _stream_converse_completion_with_spied_client( + monkeypatch: pytest.MonkeyPatch, stream_chunk_size: int | None = None +) -> tuple[MagicMock, MagicMock, LitellmParamsRecorder]: + recorder: Final = record_litellm_params(monkeypatch) + mock_response: Final = MagicMock() + mock_response.status_code = 200 + mock_response.iter_bytes = MagicMock(return_value=iter([])) + client: Final = HTTPHandler() + client.post = MagicMock(return_value=mock_response) + + litellm.completion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + stream_chunk_size=stream_chunk_size, + ) + return mock_response.iter_bytes, client.post, recorder + + +def test_completion_stream_chunk_size_reaches_iter_bytes_but_not_converse_body( + monkeypatch: pytest.MonkeyPatch, +) -> None: + iter_bytes_spy, post_spy, recorder = _stream_converse_completion_with_spied_client( + monkeypatch, stream_chunk_size=64 + ) + + iter_bytes_spy.assert_called_once_with(chunk_size=64) + data: Final = post_spy.call_args.kwargs["data"] + assert "stream_chunk_size" not in keys_at_every_depth(json.loads(data)), data + assert len(recorder.seen) == 1 + assert recorder.seen[0]["stream_chunk_size"] == 64 + + +def test_completion_without_stream_chunk_size_uses_default_chunking(monkeypatch: pytest.MonkeyPatch) -> None: + iter_bytes_spy, _, recorder = _stream_converse_completion_with_spied_client(monkeypatch) + + iter_bytes_spy.assert_called_once_with(chunk_size=None) + assert len(recorder.seen) == 1 + assert recorder.seen[0]["stream_chunk_size"] is None + + +async def _astream_converse_completion_with_spied_client( + monkeypatch: pytest.MonkeyPatch, stream_chunk_size: int | None = None +) -> tuple[MagicMock, AsyncMock, LitellmParamsRecorder]: + async def _no_bytes(chunk_size: int | None = None) -> AsyncIterator[bytes]: + return + yield b"" + + mock_response: Final = MagicMock() + mock_response.status_code = 200 + recorder: Final = record_litellm_params(monkeypatch) + mock_response.aiter_bytes = MagicMock(return_value=_no_bytes()) + aiter_bytes_spy: Final = mock_response.aiter_bytes + client: Final = AsyncHTTPHandler() + client.post = AsyncMock(return_value=mock_response) + + await litellm.acompletion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + stream_chunk_size=stream_chunk_size, + ) + return aiter_bytes_spy, client.post, recorder + + +@pytest.mark.asyncio +async def test_acompletion_stream_chunk_size_reaches_aiter_bytes_but_not_converse_body( + monkeypatch: pytest.MonkeyPatch, +) -> None: + aiter_bytes_spy, post_spy, recorder = await _astream_converse_completion_with_spied_client( + monkeypatch, stream_chunk_size=64 + ) + + aiter_bytes_spy.assert_called_once_with(chunk_size=64) + data: Final = post_spy.call_args.kwargs["data"] + assert "stream_chunk_size" not in keys_at_every_depth(json.loads(data)), data + assert len(recorder.seen) == 1 + assert recorder.seen[0]["stream_chunk_size"] == 64 + + +@pytest.mark.asyncio +async def test_acompletion_without_stream_chunk_size_uses_default_chunking( + monkeypatch: pytest.MonkeyPatch, +) -> None: + aiter_bytes_spy, _, recorder = await _astream_converse_completion_with_spied_client(monkeypatch) + + aiter_bytes_spy.assert_called_once_with(chunk_size=None) + assert len(recorder.seen) == 1 + assert recorder.seen[0]["stream_chunk_size"] is None + + +@pytest.mark.parametrize("stream_chunk_size,expected_chunk_size", [(64, 64), (None, None)]) +def test_router_deployment_stream_chunk_size_reaches_iter_bytes( + monkeypatch: pytest.MonkeyPatch, stream_chunk_size: int | None, expected_chunk_size: int | None +) -> None: + recorder: Final = record_litellm_params(monkeypatch) + mock_response: Final = MagicMock() + mock_response.status_code = 200 + mock_response.iter_bytes = MagicMock(return_value=iter([])) + client: Final = HTTPHandler() + client.post = MagicMock(return_value=mock_response) + deployment_params: Final = { + "model": "bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_access_key_id": "fake", + "aws_secret_access_key": "fake", + "aws_region_name": "us-east-1", + } + router: Final = litellm.Router( + model_list=[ + { + "model_name": "converse-chunked", + "litellm_params": deployment_params + | ({} if stream_chunk_size is None else {"stream_chunk_size": stream_chunk_size}), + } + ] + ) + + router.completion( + model="converse-chunked", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + ) + + mock_response.iter_bytes.assert_called_once_with(chunk_size=expected_chunk_size) + data: Final = client.post.call_args.kwargs["data"] + assert "stream_chunk_size" not in keys_at_every_depth(json.loads(data)), data + assert len(recorder.seen) == 1 + assert recorder.seen[0]["stream_chunk_size"] == stream_chunk_size + + def _bedrock_error_response(status_code: int, request_id: str) -> httpx.Response: return httpx.Response( status_code=status_code,