mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
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 <shreshth@berri.ai>
This commit is contained in:
parent
b41e6c966a
commit
e73f949fbb
9 changed files with 513 additions and 10 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -4043,6 +4043,7 @@ all_litellm_params = (
|
|||
"no-log",
|
||||
"base_model",
|
||||
"stream_timeout",
|
||||
"stream_chunk_size",
|
||||
"supports_system_message",
|
||||
"region_name",
|
||||
"allowed_model_region",
|
||||
|
|
|
|||
0
tests/_support/__init__.py
Normal file
0
tests/_support/__init__.py
Normal file
31
tests/_support/stream_chunk_size.py
Normal file
31
tests/_support/stream_chunk_size.py
Normal file
|
|
@ -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()
|
||||
316
tests/integration/providers/test_stream_chunk_size_wire.py
Normal file
316
tests/integration/providers/test_stream_chunk_size_wire.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue