scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 99

Merge branch 'litellm_internal_staging' into pragnyan/fix-xiaomi-output-config

Browse source
This commit is contained in:
Pragnyan Ramtha 2026-05-19 13:26:25 +05:30 • committed by GitHub
commit a5c5f31c2e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 215 additions and 72 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/llms/bedrock/chat/invoke_agent/transformation.py
View file

@ -299,9 +299,9 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM):
)
def _get_response_stream_shape(self):
from litellm.llms.bedrock.common_utils import BEDROCK_RESPONSE_STREAM_SHAPE
from litellm.llms.bedrock.common_utils import get_bedrock_response_stream_shape
return BEDROCK_RESPONSE_STREAM_SHAPE
return get_bedrock_response_stream_shape()
def _extract_response_content(self, events: InvokeAgentEventList) -> str:
"""Extract the final response content from parsed events."""

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

@ -68,9 +68,9 @@ from litellm.utils import CustomStreamWrapper, get_secret
from ..base_aws_llm import BaseAWSLLM
from ..common_utils import (
BEDROCK_RESPONSE_STREAM_SHAPE,
BedrockError,
ModelResponseIterator,
get_bedrock_response_stream_shape,
get_bedrock_tool_name,
)
@ -1828,7 +1828,8 @@ class AWSEventStreamDecoder:
yield self._chunk_parser(chunk_data=_data)
def _parse_message_from_event(self, event) -> Optional[str]:
if BEDROCK_RESPONSE_STREAM_SHAPE is None:
response_stream_shape = get_bedrock_response_stream_shape()
if response_stream_shape is None:
raise BedrockError(
status_code=500,
message=(
@ -1837,9 +1838,7 @@ class AWSEventStreamDecoder:
),
)
response_dict = event.to_response_dict()
parsed_response = self.parser.parse(
response_dict, BEDROCK_RESPONSE_STREAM_SHAPE
)
parsed_response = self.parser.parse(response_dict, response_stream_shape)
if response_dict["status_code"] != 200:
decoded_body = response_dict["body"].decode()

25
litellm/llms/bedrock/common_utils.py
View file

@ -4,6 +4,7 @@ from __future__ import annotations
Common utilities used across bedrock chat/embedding/image generation
"""
import functools
import json
import os
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
@ -963,10 +964,8 @@ def _load_bedrock_response_stream_shape():
"""
Load the ResponseStream shape from botocore's bundled bedrock-runtime schema.
Called once at module import time; the result is stored in
``BEDROCK_RESPONSE_STREAM_SHAPE`` and reused for the process lifetime.
Returns ``None`` if botocore is unavailable or the service model cannot be
loaded, so the module still imports cleanly.
loaded.
"""
try:
from botocore.loaders import Loader
@ -977,15 +976,22 @@ def _load_bedrock_response_stream_shape():
return ServiceModel(service_dict).shape_for("ResponseStream")
except Exception as e:
verbose_logger.warning(
"litellm: could not pre-load bedrock-runtime response stream shape "
"litellm: could not load bedrock-runtime response stream shape "
"— Bedrock event-stream decoding will be unavailable. Error: %s",
e,
)
return None
# Eagerly resolved once per process — avoids per-instance or per-request disk I/O.
BEDROCK_RESPONSE_STREAM_SHAPE = _load_bedrock_response_stream_shape()
@functools.lru_cache(maxsize=1)
def get_bedrock_response_stream_shape():
"""
Lazily load and cache the bedrock-runtime ResponseStream shape for the process.
Avoids importing botocore (and logging warnings) unless Bedrock event-stream
decoding is actually needed.
"""
return _load_bedrock_response_stream_shape()
class BedrockEventStreamDecoderBase:
@ -999,7 +1005,8 @@ class BedrockEventStreamDecoderBase:
self.parser = EventStreamJSONParser()
def _parse_message_from_event(self, event) -> Optional[str]:
if BEDROCK_RESPONSE_STREAM_SHAPE is None:
response_stream_shape = get_bedrock_response_stream_shape()
if response_stream_shape is None:
raise BedrockError(
status_code=500,
message=(
@ -1008,9 +1015,7 @@ class BedrockEventStreamDecoderBase:
),
)
response_dict = event.to_response_dict()
parsed_response = self.parser.parse(
response_dict, BEDROCK_RESPONSE_STREAM_SHAPE
)
parsed_response = self.parser.parse(response_dict, response_stream_shape)
if response_dict["status_code"] != 200:
decoded_body = response_dict["body"].decode()

20
litellm/llms/sagemaker/common_utils.py
View file

@ -1,3 +1,4 @@
import functools
import json
from typing import AsyncIterator, Iterator, List, Optional, Union
@ -22,14 +23,22 @@ def _load_sagemaker_response_stream_shape():
)
except Exception as e:
verbose_logger.warning(
"litellm: could not pre-load sagemaker-runtime response stream shape "
"litellm: could not load sagemaker-runtime response stream shape "
"— SageMaker event-stream decoding will be unavailable. Error: %s",
e,
)
return None
SAGEMAKER_RESPONSE_STREAM_SHAPE = _load_sagemaker_response_stream_shape()
@functools.lru_cache(maxsize=1)
def get_sagemaker_response_stream_shape():
"""
Lazily load and cache the sagemaker-runtime stream shape for the process.
Avoids importing botocore (and logging warnings) unless SageMaker event-stream
decoding is actually needed.
"""
return _load_sagemaker_response_stream_shape()
class SagemakerError(BaseLLMException):
@ -207,7 +216,8 @@ class AWSEventStreamDecoder:
verbose_logger.error(f"Final error parsing accumulated JSON: {e}")
def _parse_message_from_event(self, event) -> Optional[str]:
if SAGEMAKER_RESPONSE_STREAM_SHAPE is None:
response_stream_shape = get_sagemaker_response_stream_shape()
if response_stream_shape is None:
raise SagemakerError(
status_code=500,
message=(
@ -216,9 +226,7 @@ class AWSEventStreamDecoder:
),
)
response_dict = event.to_response_dict()
parsed_response = self.parser.parse(
response_dict, SAGEMAKER_RESPONSE_STREAM_SHAPE
)
parsed_response = self.parser.parse(response_dict, response_stream_shape)
if response_dict["status_code"] != 200:
raise ValueError(f"Bad response code, expected 200: {response_dict}")

9
litellm/proxy/proxy_server.py
View file

@ -6917,6 +6917,15 @@ async def async_data_generator( # noqa: PLR0915
if isinstance(chunk, BaseModel):
chunk = _serialize_streaming_chunk(chunk)
elif isinstance(chunk, bytes):
# Some upstream streaming iterators (e.g. AsyncGoogleGenAIGenerateContentStreamingIterator
# for /v1beta/.../streamGenerateContent) yield raw SSE bytes from Gemini.
# Decode to str so the f-string below does not emit a Python b'...' literal,
# and pass already-formatted SSE through unchanged to avoid double "data:" prefix.
chunk = chunk.decode("utf-8", errors="replace")
if chunk.startswith(("data:", "event:", ":")):
yield chunk if chunk.endswith("\n\n") else chunk + "\n\n"
continue
elif isinstance(chunk, str) and chunk.startswith("data: "):
error_message = chunk
break

79
tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py
View file

@ -14,18 +14,46 @@ from litellm.llms.bedrock.common_utils import BedrockModelInfo
# --------------------------------------------------------------------------- #
# BEDROCK_RESPONSE_STREAM_SHAPE eager-load tests #
# get_bedrock_response_stream_shape lazy-load tests #
# --------------------------------------------------------------------------- #
def test_bedrock_response_stream_shape_loaded_at_import():
@pytest.fixture(autouse=True)
def _reset_bedrock_response_stream_shape_cache():
"""Prevent lru_cache leakage between tests in this module."""
import litellm.llms.bedrock.common_utils as mod
mod.get_bedrock_response_stream_shape.cache_clear()
yield
mod.get_bedrock_response_stream_shape.cache_clear()
def test_bedrock_response_stream_shape_lazy_loads_once():
"""
BEDROCK_RESPONSE_STREAM_SHAPE is resolved at module import time.
get_bedrock_response_stream_shape() loads from botocore at most once per process.
"""
from unittest.mock import MagicMock, patch
import litellm.llms.bedrock.common_utils as mod
sentinel = MagicMock()
with patch.object(
mod, "_load_bedrock_response_stream_shape", return_value=sentinel
) as mock_load:
assert mod.get_bedrock_response_stream_shape() is sentinel
assert mod.get_bedrock_response_stream_shape() is sentinel
mock_load.assert_called_once()
def test_bedrock_response_stream_shape_loaded_on_first_access():
"""
get_bedrock_response_stream_shape() loads once on first use.
In a standard environment with botocore installed it must be non-None.
"""
from litellm.llms.bedrock.common_utils import BEDROCK_RESPONSE_STREAM_SHAPE
pytest.importorskip("botocore")
from litellm.llms.bedrock.common_utils import get_bedrock_response_stream_shape
assert BEDROCK_RESPONSE_STREAM_SHAPE is not None
assert get_bedrock_response_stream_shape() is not None
def test_bedrock_response_stream_shape_load_failure_returns_none():
@ -38,6 +66,7 @@ def test_bedrock_response_stream_shape_load_failure_returns_none():
import litellm.llms.bedrock.common_utils as mod
pytest.importorskip("botocore")
with patch(
"botocore.loaders.Loader.load_service_model",
side_effect=Exception("no data"),
@ -51,31 +80,29 @@ def test_bedrock_response_stream_shape_is_structure_shape():
The loaded shape should be the botocore StructureShape for ResponseStream,
not a plain dict or any other type.
"""
pytest.importorskip("botocore")
from botocore.model import StructureShape
from litellm.llms.bedrock.common_utils import BEDROCK_RESPONSE_STREAM_SHAPE
from litellm.llms.bedrock.common_utils import get_bedrock_response_stream_shape
assert BEDROCK_RESPONSE_STREAM_SHAPE is not None, (
"BEDROCK_RESPONSE_STREAM_SHAPE is None — botocore may not be installed"
)
shape: StructureShape = BEDROCK_RESPONSE_STREAM_SHAPE # remove Optional
loaded_shape = get_bedrock_response_stream_shape()
assert (
loaded_shape is not None
), "get_bedrock_response_stream_shape() is None — botocore may not be installed"
shape: StructureShape = loaded_shape
assert isinstance(shape, StructureShape)
assert shape.name == "ResponseStream"
def test_bedrock_response_stream_shape_same_object_across_imports():
def test_bedrock_response_stream_shape_same_object_across_calls():
"""
Both bedrock modules that use the shape must reference the identical object —
confirming the constant is not re-loaded per import.
Repeated calls must return the identical cached object.
"""
from litellm.llms.bedrock.chat.invoke_handler import (
BEDROCK_RESPONSE_STREAM_SHAPE as invoke_shape,
)
from litellm.llms.bedrock.common_utils import (
BEDROCK_RESPONSE_STREAM_SHAPE as common_shape,
)
from litellm.llms.bedrock.common_utils import get_bedrock_response_stream_shape
assert common_shape is invoke_shape
first = get_bedrock_response_stream_shape()
second = get_bedrock_response_stream_shape()
assert first is second
def test_bedrock_event_stream_decoder_base_uses_module_shape():
@ -95,19 +122,23 @@ def test_bedrock_event_stream_decoder_base_uses_module_shape():
def test_bedrock_parse_message_from_event_raises_on_none_shape():
"""
When BEDROCK_RESPONSE_STREAM_SHAPE is None (botocore unavailable),
When get_bedrock_response_stream_shape() returns None (botocore unavailable),
_parse_message_from_event must raise BedrockError before touching the
botocore parser — not an opaque AttributeError from inside botocore.
"""
from unittest.mock import MagicMock, patch
import litellm.llms.bedrock.common_utils as mod
from litellm.llms.bedrock.common_utils import BedrockError, BedrockEventStreamDecoderBase
from litellm.llms.bedrock.common_utils import (
BedrockError,
BedrockEventStreamDecoderBase,
)
decoder = BedrockEventStreamDecoderBase()
decoder = BedrockEventStreamDecoderBase.__new__(BedrockEventStreamDecoderBase)
decoder.parser = MagicMock()
mock_event = MagicMock()
with patch.object(mod, "BEDROCK_RESPONSE_STREAM_SHAPE", None):
with patch.object(mod, "get_bedrock_response_stream_shape", return_value=None):
with pytest.raises(BedrockError) as exc_info:
decoder._parse_message_from_event(mock_event)

81
tests/test_litellm/llms/sagemaker/test_sagemaker_common_utils.py
View file

@ -12,18 +12,46 @@ from litellm.llms.sagemaker.completion.transformation import SagemakerConfig
# --------------------------------------------------------------------------- #
# SAGEMAKER_RESPONSE_STREAM_SHAPE eager-load tests #
# get_sagemaker_response_stream_shape lazy-load tests #
# --------------------------------------------------------------------------- #
def test_sagemaker_response_stream_shape_loaded_at_import():
@pytest.fixture(autouse=True)
def _reset_sagemaker_response_stream_shape_cache():
"""Prevent lru_cache leakage between tests in this module."""
import litellm.llms.sagemaker.common_utils as mod
mod.get_sagemaker_response_stream_shape.cache_clear()
yield
mod.get_sagemaker_response_stream_shape.cache_clear()
def test_sagemaker_response_stream_shape_lazy_loads_once():
"""
SAGEMAKER_RESPONSE_STREAM_SHAPE is resolved at module import time.
get_sagemaker_response_stream_shape() loads from botocore at most once per process.
"""
from unittest.mock import MagicMock, patch
import litellm.llms.sagemaker.common_utils as mod
sentinel = MagicMock()
with patch.object(
mod, "_load_sagemaker_response_stream_shape", return_value=sentinel
) as mock_load:
assert mod.get_sagemaker_response_stream_shape() is sentinel
assert mod.get_sagemaker_response_stream_shape() is sentinel
mock_load.assert_called_once()
def test_sagemaker_response_stream_shape_loaded_on_first_access():
"""
get_sagemaker_response_stream_shape() loads once on first use.
In a standard environment with botocore installed it must be non-None.
"""
from litellm.llms.sagemaker.common_utils import SAGEMAKER_RESPONSE_STREAM_SHAPE
pytest.importorskip("botocore")
from litellm.llms.sagemaker.common_utils import get_sagemaker_response_stream_shape
assert SAGEMAKER_RESPONSE_STREAM_SHAPE is not None
assert get_sagemaker_response_stream_shape() is not None
def test_sagemaker_response_stream_shape_load_failure_returns_none():
@ -36,6 +64,7 @@ def test_sagemaker_response_stream_shape_load_failure_returns_none():
import litellm.llms.sagemaker.common_utils as mod
pytest.importorskip("botocore")
with patch(
"botocore.loaders.Loader.load_service_model",
side_effect=Exception("no data"),
@ -49,14 +78,16 @@ def test_sagemaker_response_stream_shape_is_structure_shape():
The loaded shape should be the botocore StructureShape for
InvokeEndpointWithResponseStreamOutput, not a plain dict or any other type.
"""
pytest.importorskip("botocore")
from botocore.model import StructureShape
from litellm.llms.sagemaker.common_utils import SAGEMAKER_RESPONSE_STREAM_SHAPE
from litellm.llms.sagemaker.common_utils import get_sagemaker_response_stream_shape
assert SAGEMAKER_RESPONSE_STREAM_SHAPE is not None, (
"SAGEMAKER_RESPONSE_STREAM_SHAPE is None — botocore may not be installed"
)
shape: StructureShape = SAGEMAKER_RESPONSE_STREAM_SHAPE # remove Optional
shape = get_sagemaker_response_stream_shape()
assert (
shape is not None
), "get_sagemaker_response_stream_shape() is None — botocore may not be installed"
shape: StructureShape = shape # remove Optional
assert isinstance(shape, StructureShape)
assert shape.name == "InvokeEndpointWithResponseStreamOutput"
@ -64,29 +95,25 @@ def test_sagemaker_response_stream_shape_is_structure_shape():
def test_sagemaker_response_stream_shape_not_reloaded_on_new_decoder():
"""
Creating multiple AWSEventStreamDecoder instances must not trigger
additional botocore Loader calls — the shape is resolved once at import
time and reused.
additional botocore Loader calls — the shape is cached after first access.
"""
from litellm.llms.sagemaker.common_utils import SAGEMAKER_RESPONSE_STREAM_SHAPE
from litellm.llms.sagemaker.common_utils import get_sagemaker_response_stream_shape
decoder_a = AWSEventStreamDecoder(model="test-model-a")
decoder_b = AWSEventStreamDecoder(model="test-model-b")
decoder_a = AWSEventStreamDecoder.__new__(AWSEventStreamDecoder)
decoder_b = AWSEventStreamDecoder.__new__(AWSEventStreamDecoder)
# Both decoders should use the same pre-loaded shape object (identity check)
assert "_response_stream_shape_cache" not in decoder_a.__dict__
assert "_response_stream_shape_cache" not in decoder_b.__dict__
# The module constant is still the same object
from litellm.llms.sagemaker.common_utils import (
SAGEMAKER_RESPONSE_STREAM_SHAPE as shape_after,
)
assert SAGEMAKER_RESPONSE_STREAM_SHAPE is shape_after
first = get_sagemaker_response_stream_shape()
second = get_sagemaker_response_stream_shape()
assert first is second
def test_sagemaker_parse_message_from_event_raises_on_none_shape():
"""
When SAGEMAKER_RESPONSE_STREAM_SHAPE is None (botocore unavailable),
_parse_message_from_event must raise ValueError before touching the
When get_sagemaker_response_stream_shape() returns None (botocore unavailable),
_parse_message_from_event must raise SagemakerError before touching the
botocore parser — not an opaque AttributeError from inside botocore.
"""
from unittest.mock import MagicMock, patch
@ -94,10 +121,14 @@ def test_sagemaker_parse_message_from_event_raises_on_none_shape():
import litellm.llms.sagemaker.common_utils as mod
from litellm.llms.sagemaker.common_utils import SagemakerError
decoder = AWSEventStreamDecoder(model="test-model")
decoder = AWSEventStreamDecoder.__new__(AWSEventStreamDecoder)
decoder.model = "test-model"
decoder.parser = MagicMock()
decoder.content_blocks = []
decoder.is_messages_api = None
mock_event = MagicMock()
with patch.object(mod, "SAGEMAKER_RESPONSE_STREAM_SHAPE", None):
with patch.object(mod, "get_sagemaker_response_stream_shape", return_value=None):
with pytest.raises(SagemakerError) as exc_info:
decoder._parse_message_from_event(mock_event)

60
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -5065,6 +5065,66 @@ async def test_async_data_generator_uses_direct_stream_fast_path_without_callbac
mock_response.aclose.assert_awaited_once()
@pytest.mark.asyncio
async def test_async_data_generator_passes_through_google_native_sse_bytes():
"""
Google-native streamGenerateContent yields raw SSE bytes; they must not be
re-wrapped as data: b'data: {...}'.
"""
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.proxy_server import async_data_generator
from litellm.proxy.utils import ProxyLogging
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
mock_request_data = {
"model": "gemini-2.0-flash",
"messages": [{"role": "user", "content": "test"}],
}
gemini_event = b'data: {"candidates": [{"content": "hi"}]}\n\n'
gemini_event_without_terminator = b'data: {"candidates": [{"content": "there"}]}'
raw_payload = b'{"partial": true}'
class MockStream:
def __aiter__(self):
return self._stream()
async def _stream(self):
yield gemini_event
yield gemini_event_without_terminator
yield raw_payload
async def aclose(self):
pass
mock_response = MockStream()
mock_response.aclose = AsyncMock()
mock_proxy_logging_obj = MagicMock(spec=ProxyLogging)
mock_proxy_logging_obj.has_streaming_callbacks.return_value = False
mock_proxy_logging_obj.needs_iterator_wrap.return_value = False
mock_proxy_logging_obj.needs_per_chunk_streaming_hook.return_value = False
mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = MagicMock()
mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock()
mock_proxy_logging_obj.post_call_failure_hook = AsyncMock()
with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj):
with patch.object(ProxyLogging, "_fire_deferred_stream_logging"):
yielded_data = []
async for data in async_data_generator(
mock_response, mock_user_api_key_dict, mock_request_data
):
yielded_data.append(data)
yielded_text = [
chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk
for chunk in yielded_data
]
assert yielded_text[0] == gemini_event.decode("utf-8")
assert yielded_text[1] == gemini_event_without_terminator.decode("utf-8") + "\n\n"
assert yielded_text[2] == f'data: {raw_payload.decode("utf-8")}\n\n'
assert "b'data:" not in "".join(yielded_text)
assert yielded_text[-1] == "data: [DONE]\n\n"
@pytest.mark.asyncio
async def test_async_data_generator_cleanup_on_normal_completion():
"""