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 115

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_access-groups-pass-through-vector-stores-d3df

Browse source 
Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>
This commit is contained in:
oss-agent-shin 2026-05-07 01:47:38 +00:00
commit 805608ed0d
No known key found for this signature in database
12 changed files with 615 additions and 85 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.github/workflows/test-unit-proxy-db.yml vendored
View file

@ -141,6 +141,7 @@ jobs:
tests/proxy_unit_tests/test_server_root_path.py
tests/proxy_unit_tests/test_proxy_pass_user_config.py
tests/proxy_unit_tests/test_proxy_token_counter.py
tests/proxy_unit_tests/test_request_size_limit_middleware.py
workers: 4
dist: loadscope
timeout: 15

3
.gitignore vendored
View file

@ -100,4 +100,5 @@ STABILIZATION_TODO.md
**/playwright-report
**/*.storageState.json
**/coverage
test-config
test-config
.vscode

7
codecov.yaml
View file

@ -1,3 +1,8 @@
codecov:
require_ci_to_pass: false # post coverage status even if CI has unrelated failures
notify:
wait_for_ci: false # post as soon as expected uploads arrive, don't wait on CI
component_management:
individual_components:
- component_id: "Router"
@ -28,7 +33,7 @@ coverage:
project:
default:
target: auto
threshold: 1% # at maximum allow project coverage to drop by 1%
threshold: 0% # do not allow project coverage to drop
patch:
default:
target: auto

12
litellm/integrations/prometheus.py
View file

@ -1443,12 +1443,12 @@ class PrometheusLogger(CustomLogger):
)
remaining_tokens_variable_name = f"litellm-key-remaining-tokens-{model_group}"
remaining_requests = (
metadata.get(remaining_requests_variable_name, sys.maxsize) or sys.maxsize
)
remaining_tokens = (
metadata.get(remaining_tokens_variable_name, sys.maxsize) or sys.maxsize
)
remaining_requests = metadata.get(remaining_requests_variable_name)
if remaining_requests is None:
remaining_requests = sys.maxsize
remaining_tokens = metadata.get(remaining_tokens_variable_name)
if remaining_tokens is None:
remaining_tokens = sys.maxsize
enum_values = UserAPIKeyLabelValues(
hashed_api_key=user_api_key,

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

@ -299,29 +299,9 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM):
)
def _get_response_stream_shape(self):
"""Get the response stream shape for parsing, reusing existing logic."""
try:
# Try to reuse the cached shape from the existing decoder
from litellm.llms.bedrock.chat.invoke_handler import (
get_response_stream_shape,
)
from litellm.llms.bedrock.common_utils import BEDROCK_RESPONSE_STREAM_SHAPE
return get_response_stream_shape()
except ImportError:
# Fallback: create our own shape
try:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
bedrock_service_dict = loader.load_service_model(
"bedrock-runtime", "service-2"
)
bedrock_service_model = ServiceModel(bedrock_service_dict)
return bedrock_service_model.shape_for("ResponseStream")
except Exception as e:
verbose_logger.warning(f"Could not load response stream shape: {e}")
return None
return BEDROCK_RESPONSE_STREAM_SHAPE
def _extract_response_content(self, events: InvokeAgentEventList) -> str:
"""Extract the final response content from parsed events."""

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

@ -67,9 +67,13 @@ from litellm.types.utils import (
from litellm.utils import CustomStreamWrapper, get_secret
from ..base_aws_llm import BaseAWSLLM
from ..common_utils import BedrockError, ModelResponseIterator, get_bedrock_tool_name
from ..common_utils import (
BEDROCK_RESPONSE_STREAM_SHAPE,
BedrockError,
ModelResponseIterator,
get_bedrock_tool_name,
)
_response_stream_shape_cache = None
bedrock_tool_name_mappings: InMemoryCache = InMemoryCache(
max_size_in_memory=50, default_ttl=600
)
@ -1391,20 +1395,6 @@ class BedrockLLM(BaseAWSLLM):
return None
def get_response_stream_shape():
global _response_stream_shape_cache
if _response_stream_shape_cache is None:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
bedrock_service_dict = loader.load_service_model("bedrock-runtime", "service-2")
bedrock_service_model = ServiceModel(bedrock_service_dict)
_response_stream_shape_cache = bedrock_service_model.shape_for("ResponseStream")
return _response_stream_shape_cache
class AWSEventStreamDecoder:
def __init__(self, model: str, json_mode: Optional[bool] = False) -> None:
from botocore.parsers import EventStreamJSONParser
@ -1838,8 +1828,18 @@ 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:
raise BedrockError(
status_code=500,
message=(
"Bedrock event-stream shape could not be loaded from botocore. "
"Ensure botocore is correctly installed."
),
)
response_dict = event.to_response_dict()
parsed_response = self.parser.parse(response_dict, get_response_stream_shape())
parsed_response = self.parser.parse(
response_dict, BEDROCK_RESPONSE_STREAM_SHAPE
)
if response_dict["status_code"] != 200:
decoded_body = response_dict["body"].decode()

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

@ -14,6 +14,7 @@ if TYPE_CHECKING:
import httpx
import litellm
from litellm import verbose_logger
from litellm.llms.base_llm.anthropic_messages.transformation import (
BaseAnthropicMessagesConfig,
)
@ -917,38 +918,57 @@ def get_bedrock_chat_config(model: str):
return litellm.AmazonInvokeConfig()
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.
"""
try:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
service_dict = loader.load_service_model("bedrock-runtime", "service-2")
return ServiceModel(service_dict).shape_for("ResponseStream")
except Exception as e:
verbose_logger.warning(
"litellm: could not pre-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()
class BedrockEventStreamDecoderBase:
"""
Base class for event stream decoding for Bedrock
"""
_response_stream_shape_cache = None
def __init__(self):
from botocore.parsers import EventStreamJSONParser
self.parser = EventStreamJSONParser()
def get_response_stream_shape(self):
if self._response_stream_shape_cache is None:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
bedrock_service_dict = loader.load_service_model(
"bedrock-runtime", "service-2"
)
bedrock_service_model = ServiceModel(bedrock_service_dict)
self._response_stream_shape_cache = bedrock_service_model.shape_for(
"ResponseStream"
)
return self._response_stream_shape_cache
def _parse_message_from_event(self, event) -> Optional[str]:
if BEDROCK_RESPONSE_STREAM_SHAPE is None:
raise BedrockError(
status_code=500,
message=(
"Bedrock event-stream shape could not be loaded from botocore. "
"Ensure botocore is correctly installed."
),
)
response_dict = event.to_response_dict()
parsed_response = self.parser.parse(
response_dict, self.get_response_stream_shape()
response_dict, BEDROCK_RESPONSE_STREAM_SHAPE
)
if response_dict["status_code"] != 200:

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

@ -9,7 +9,27 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.utils import GenericStreamingChunk as GChunk
from litellm.types.utils import StreamingChatCompletionChunk
_response_stream_shape_cache = None
def _load_sagemaker_response_stream_shape():
try:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
service_dict = loader.load_service_model("sagemaker-runtime", "service-2")
return ServiceModel(service_dict).shape_for(
"InvokeEndpointWithResponseStreamOutput"
)
except Exception as e:
verbose_logger.warning(
"litellm: could not pre-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()
class SagemakerError(BaseLLMException):
@ -187,8 +207,18 @@ 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:
raise SagemakerError(
status_code=500,
message=(
"SageMaker event-stream shape could not be loaded from botocore. "
"Ensure botocore is correctly installed."
),
)
response_dict = event.to_response_dict()
parsed_response = self.parser.parse(response_dict, get_response_stream_shape())
parsed_response = self.parser.parse(
response_dict, SAGEMAKER_RESPONSE_STREAM_SHAPE
)
if response_dict["status_code"] != 200:
raise ValueError(f"Bad response code, expected 200: {response_dict}")
@ -204,20 +234,3 @@ class AWSEventStreamDecoder:
return None
return chunk.decode() # type: ignore[no-any-return]
def get_response_stream_shape():
global _response_stream_shape_cache
if _response_stream_shape_cache is None:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
sagemaker_service_dict = loader.load_service_model(
"sagemaker-runtime", "service-2"
)
sagemaker_service_model = ServiceModel(sagemaker_service_dict)
_response_stream_shape_cache = sagemaker_service_model.shape_for(
"InvokeEndpointWithResponseStreamOutput"
)
return _response_stream_shape_cache

276
scripts/mock_bedrock_passthrough_target.py Normal file
View file

@ -0,0 +1,276 @@
#!/usr/bin/env python3
"""
Minimal HTTP target for testing LiteLLM **Bedrock pass-through** (`/bedrock/...` on the proxy).
What it does
- Serves a tiny Converse-shaped JSON (and optional invoke-shaped) response so the proxy can
complete a round trip without calling AWS.
- Does **not** verify SigV4 (Bedrock does); any Authorization header is accepted.
How to run
uv run python scripts/mock_bedrock_passthrough_target.py --host 127.0.0.1 --port 9999
Wire LiteLLM to this host (use **one** of these patterns):
1) model_list (recommended) — set the Bedrock runtime base to the mock:
model_list:
- model_name: mock-bedrock-claude
litellm_params:
model: bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0
custom_llm_provider: bedrock
aws_region_name: us-west-2
api_base: "http://127.0.0.1:9999"
2) Environment (see litellm BaseAWSLLM.get_runtime_endpoint)::
export AWS_BEDROCK_RUNTIME_ENDPOINT="http://127.0.0.1:9999"
Then call the proxy, e.g. (model_name must match config)::
curl -sS -X POST "http://127.0.0.1:4000/bedrock/model/mock-bedrock-claude/converse" \
-H "Authorization: Bearer $LITELLM_KEY" -H "Content-Type: application/json" \
-d '{"messages":[{"role":"user","content":[{"text":"hi"}]}]}'
The proxy will forward to: {api_base}/model/<resolved model id>/converse (SigV4-signed).
This mock implements POST .../converse and returns a minimal valid Converse response.
Notes
- `invoke-with-response-stream` returns a real **binary** AWS event stream
(`application/vnd.amazon.eventstream`) with Anthropic-style JSON payloads inside each
`PayloadPart`, matching Bedrock's InvokeModelWithResponseStream wire format. See
https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InvokeModelWithResponseStream.html
and https://docs.aws.amazon.com/awstreams/latest/devguide/message-formats.html
- `converse-stream` is still JSON-only placeholder (different inner event shapes).
- Use real (or any non-empty) AWS creds in the environment of the **proxy**; signing still runs.
"""
from __future__ import annotations
import argparse
import base64
import json
from binascii import crc32
from struct import pack
from typing import Any, Dict, Iterator, List
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from starlette.responses import StreamingResponse
app = FastAPI(title="Mock Bedrock runtime (pass-through test target)")
# Minimal structure compatible with Converse: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Converse.html
def _converse_response_body() -> Dict[str, Any]:
return {
"output": {
"message": {
"role": "assistant",
"content": [
{"text": "mock: ok from mock_bedrock_passthrough_target.py"}
],
}
},
"stopReason": "end_turn",
"usage": {
"inputTokens": 1,
"outputTokens": 2,
"totalTokens": 3,
},
}
# Minimal invoke (Anthropic messages on bedrock) style — adjust if you test /invoke
def _invoke_response_body() -> Dict[str, Any]:
return {
"id": "msg_mock",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "mock invoke response"}],
"model": "mock",
"stop_reason": "end_turn",
"usage": {"input_tokens": 1, "output_tokens": 2},
}
def _encode_event_stream_message(headers: Dict[str, str], payload: bytes) -> bytes:
"""Single AWS binary event-stream frame (same layout botocore's ``EventStreamBuffer`` parses)."""
header_blob = b""
for name, value in headers.items():
nb = name.encode("utf-8")
vb = value.encode("utf-8")
header_blob += bytes([len(nb)]) + nb + bytes([7]) + pack("!H", len(vb)) + vb
headers_length = len(header_blob)
payload_length = len(payload)
total_length = 12 + headers_length + payload_length + 4
prelude_wo_crc = pack("!II", total_length, headers_length)
prelude_crc_val = crc32(prelude_wo_crc) & 0xFFFFFFFF
prelude = prelude_wo_crc + pack("!I", prelude_crc_val)
wo_msg_crc = prelude + header_blob + payload
msg_crc_val = crc32(wo_msg_crc[8:], prelude_crc_val) & 0xFFFFFFFF
return wo_msg_crc + pack("!I", msg_crc_val)
def _bedrock_payload_part(inner_event: Dict[str, Any]) -> bytes:
"""Outer JSON expected by bedrock-runtime ``ResponseStream`` / ``PayloadPart``."""
inner_bytes = json.dumps(inner_event, separators=(",", ":")).encode("utf-8")
outer = {
"chunk": {
"bytes": base64.b64encode(inner_bytes).decode("ascii"),
}
}
return json.dumps(outer, separators=(",", ":")).encode("utf-8")
def _anthropic_invoke_stream_events(
model_id: str, assistant_text: str
) -> List[Dict[str, Any]]:
"""
Minimal Anthropic Messages stream events as returned inside Bedrock stream chunks.
Mirrors the sequence Amazon emits for Claude on ``invoke-with-response-stream``.
"""
msg_id = "msg_mock_bedrock_stream"
input_tokens = 3
output_tokens = max(1, len(assistant_text) // 4)
events: List[Dict[str, Any]] = [
{
"type": "message_start",
"message": {
"model": model_id,
"id": msg_id,
"type": "message",
"role": "assistant",
"content": [],
"stop_reason": None,
"stop_sequence": None,
"usage": {
"input_tokens": input_tokens,
"output_tokens": 1,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0,
"cache_creation": {
"ephemeral_5m_input_tokens": 0,
"ephemeral_1h_input_tokens": 0,
},
},
},
},
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "text", "text": ""},
},
]
# Split text into small deltas so downstream streaming behavior is visible.
step = 24
for i in range(0, len(assistant_text), step):
events.append(
{
"type": "content_block_delta",
"index": 0,
"delta": {
"type": "text_delta",
"text": assistant_text[i : i + step],
},
}
)
events.append({"type": "content_block_stop", "index": 0})
events.append(
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
"usage": {
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0,
},
}
)
events.append(
{
"type": "message_stop",
"amazon-bedrock-invocationMetrics": {
"inputTokenCount": input_tokens,
"outputTokenCount": output_tokens,
"invocationLatency": 42,
"firstByteLatency": 10,
},
}
)
return events
def _iter_invoke_with_response_stream(model_id: str) -> Iterator[bytes]:
text = (
"mock streaming: ok from scripts/mock_bedrock_passthrough_target.py "
"(invoke-with-response-stream)."
)
headers = {
":event-type": "chunk",
":content-type": "application/json",
":message-type": "event",
}
for ev in _anthropic_invoke_stream_events(model_id, text):
yield _encode_event_stream_message(headers, _bedrock_payload_part(ev))
@app.get("/health")
def health() -> Dict[str, str]:
return {"status": "ok"}
@app.post("/model/{model_path:path}/converse")
async def converse(model_path: str, request: Request) -> JSONResponse:
# Optional: log body for debugging
_ = await request.body()
return JSONResponse(content=_converse_response_body())
@app.post("/model/{model_path:path}/converse-stream")
async def converse_stream(model_path: str, request: Request) -> JSONResponse:
"""
Not a real AWS event stream — returns JSON for quick smoke tests only.
"""
_ = await request.body()
return JSONResponse(
content={
"note": "This mock does not implement application/vnd.amazon.eventstream; use /converse for basic tests."
}
)
@app.post("/model/{model_path:path}/invoke")
async def invoke(model_path: str, request: Request) -> JSONResponse:
_ = await request.body()
return JSONResponse(content=_invoke_response_body())
@app.post("/model/{model_path:path}/invoke-with-response-stream")
async def invoke_with_response_stream(
model_path: str, request: Request
) -> StreamingResponse:
"""
Binary ``application/vnd.amazon.eventstream`` body compatible with boto3/botocore
``InvokeModelWithResponseStream`` / LiteLLM's Bedrock invoke streaming path.
"""
_ = await request.body()
return StreamingResponse(
_iter_invoke_with_response_stream(model_id=model_path),
media_type="application/vnd.amazon.eventstream",
)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", type=int, default=9999)
args = parser.parse_args()
import uvicorn
uvicorn.run(app, host=args.host, port=args.port, log_level="info")
if __name__ == "__main__":
main()

34
tests/test_litellm/integrations/test_prometheus_custom_metadata_label_counts.py
View file

@ -1,4 +1,5 @@
import logging
import sys
import pytest
from prometheus_client import REGISTRY
@ -123,3 +124,36 @@ def test_virtual_key_rate_limit_metrics_accept_custom_metadata_labels(
and sample.value == 3
for sample in samples
)
def test_virtual_key_rate_limit_metrics_preserve_zero_remaining_values(
monkeypatch: pytest.MonkeyPatch,
):
prometheus_logger = _create_prometheus_logger_with_custom_labels(monkeypatch)
metadata = {
"model_group": "gpt-4o-mini",
"litellm-key-remaining-requests-gpt-4o-mini": 0,
"litellm-key-remaining-tokens-gpt-4o-mini": 0,
}
kwargs = {
"litellm_params": {
"metadata": metadata,
},
"standard_logging_object": _standard_logging_payload_with_requester_metadata(),
}
prometheus_logger._set_virtual_key_rate_limit_metrics(
user_api_key="test-hash",
user_api_key_alias="test-alias",
kwargs=kwargs,
metadata=metadata,
model_id="model-123",
)
request_samples = _metric_samples("litellm_remaining_api_key_requests_for_model")
token_samples = _metric_samples("litellm_remaining_api_key_tokens_for_model")
assert any(sample.value == 0 for sample in request_samples)
assert any(sample.value == 0 for sample in token_samples)
assert not any(sample.value == sys.maxsize for sample in request_samples)
assert not any(sample.value == sys.maxsize for sample in token_samples)

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

@ -13,6 +13,110 @@ sys.path.insert(
from litellm.llms.bedrock.common_utils import BedrockModelInfo
# --------------------------------------------------------------------------- #
# BEDROCK_RESPONSE_STREAM_SHAPE eager-load tests #
# --------------------------------------------------------------------------- #
def test_bedrock_response_stream_shape_loaded_at_import():
"""
BEDROCK_RESPONSE_STREAM_SHAPE is resolved at module import time.
In a standard environment with botocore installed it must be non-None.
"""
from litellm.llms.bedrock.common_utils import BEDROCK_RESPONSE_STREAM_SHAPE
assert BEDROCK_RESPONSE_STREAM_SHAPE is not None
def test_bedrock_response_stream_shape_load_failure_returns_none():
"""
If botocore's Loader raises (e.g. missing data files), _load_bedrock_response_stream_shape
should return None rather than propagating the exception, so the module
still imports cleanly.
"""
from unittest.mock import patch
import litellm.llms.bedrock.common_utils as mod
with patch(
"botocore.loaders.Loader.load_service_model",
side_effect=Exception("no data"),
):
shape = mod._load_bedrock_response_stream_shape()
assert shape is None
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.
"""
from botocore.model import StructureShape
from litellm.llms.bedrock.common_utils import 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
assert isinstance(shape, StructureShape)
assert shape.name == "ResponseStream"
def test_bedrock_response_stream_shape_same_object_across_imports():
"""
Both bedrock modules that use the shape must reference the identical object —
confirming the constant is not re-loaded per import.
"""
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,
)
assert common_shape is invoke_shape
def test_bedrock_event_stream_decoder_base_uses_module_shape():
"""
BedrockEventStreamDecoderBase instances no longer carry their own
per-instance cache — _parse_message_from_event uses the module constant
directly, so there is no instance-level _response_stream_shape_cache attr.
"""
from litellm.llms.bedrock.common_utils import BedrockEventStreamDecoderBase
decoder_a = BedrockEventStreamDecoderBase()
decoder_b = BedrockEventStreamDecoderBase()
assert "_response_stream_shape_cache" not in decoder_a.__dict__
assert "_response_stream_shape_cache" not in decoder_b.__dict__
def test_bedrock_parse_message_from_event_raises_on_none_shape():
"""
When BEDROCK_RESPONSE_STREAM_SHAPE is 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
decoder = BedrockEventStreamDecoderBase()
mock_event = MagicMock()
with patch.object(mod, "BEDROCK_RESPONSE_STREAM_SHAPE", None):
with pytest.raises(BedrockError) as exc_info:
decoder._parse_message_from_event(mock_event)
assert exc_info.value.status_code == 500
assert "botocore" in str(exc_info.value.message).lower()
# The botocore parser must never have been called
mock_event.to_response_dict.assert_not_called()
def test_deepseek_cris():
"""
Test that DeepSeek models with cross-region inference prefix use converse route

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

@ -11,6 +11,102 @@ from litellm.llms.sagemaker.common_utils import AWSEventStreamDecoder
from litellm.llms.sagemaker.completion.transformation import SagemakerConfig
# --------------------------------------------------------------------------- #
# SAGEMAKER_RESPONSE_STREAM_SHAPE eager-load tests #
# --------------------------------------------------------------------------- #
def test_sagemaker_response_stream_shape_loaded_at_import():
"""
SAGEMAKER_RESPONSE_STREAM_SHAPE is resolved at module import time.
In a standard environment with botocore installed it must be non-None.
"""
from litellm.llms.sagemaker.common_utils import SAGEMAKER_RESPONSE_STREAM_SHAPE
assert SAGEMAKER_RESPONSE_STREAM_SHAPE is not None
def test_sagemaker_response_stream_shape_load_failure_returns_none():
"""
If botocore's Loader raises (e.g. missing data files), _load_sagemaker_response_stream_shape
should return None rather than propagating the exception, so the module
still imports cleanly.
"""
from unittest.mock import patch
import litellm.llms.sagemaker.common_utils as mod
with patch(
"botocore.loaders.Loader.load_service_model",
side_effect=Exception("no data"),
):
shape = mod._load_sagemaker_response_stream_shape()
assert shape is None
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.
"""
from botocore.model import StructureShape
from litellm.llms.sagemaker.common_utils import 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
assert isinstance(shape, StructureShape)
assert shape.name == "InvokeEndpointWithResponseStreamOutput"
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.
"""
from litellm.llms.sagemaker.common_utils import SAGEMAKER_RESPONSE_STREAM_SHAPE
decoder_a = AWSEventStreamDecoder(model="test-model-a")
decoder_b = AWSEventStreamDecoder(model="test-model-b")
# 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
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
botocore parser — not an opaque AttributeError from inside botocore.
"""
from unittest.mock import MagicMock, patch
import litellm.llms.sagemaker.common_utils as mod
from litellm.llms.sagemaker.common_utils import SagemakerError
decoder = AWSEventStreamDecoder(model="test-model")
mock_event = MagicMock()
with patch.object(mod, "SAGEMAKER_RESPONSE_STREAM_SHAPE", None):
with pytest.raises(SagemakerError) as exc_info:
decoder._parse_message_from_event(mock_event)
assert exc_info.value.status_code == 500
assert "botocore" in str(exc_info.value.message).lower()
# The botocore parser must never have been called
mock_event.to_response_dict.assert_not_called()
@pytest.mark.asyncio
async def test_aiter_bytes_unicode_decode_error():
"""