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

refactor(sdk): separate core AWS and tokenizer dependencies (#44447)

Browse source 
* refactor(sdk): separate core AWS and tokenizer dependencies

* fix(sdk): preserve runtime tokenizer alias compatibility

* test(sdk): compare tokenizer installs through existing entry point

* fix(sdk): preserve bearer headers and optional dependency interfaces

* fix(sdk): retain safe tokenizer fallback diagnostics

* test(sdk): inject tokenizer dependency for fallback diagnostics

* fix(core): preserve AWS dependency errors with retries

* fix(core): centralize optional AWS dependency handling

* refactor(core): reuse optional import helper for Invoke streams

* test(core): compare complete shared dependency requirements

* fix(core): keep tokenizer logging import compatible with main

* fix(core): preserve native token decoding and Responses dependency errors

* test(core): isolate optional dependency import failures

* fix(sdk): retain typed exception message access

* test(sdk): scope HTTP verification environment changes

* fix(sdk): preserve bearer request typing across Bedrock callers

---------

Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
This commit is contained in:
joshua-berri 2026-10-08 19:07:57 -07:00 • committed by GitHub
parent 79c3de46cf
commit 911aff2752
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
37 changed files with 1085 additions and 202 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
.github/workflows/test-mcp-dependency-resolution.yml vendored
View file

@ -89,6 +89,7 @@ jobs:
uv pip check --python ".venv-$extra"
if [ "$extra" = core ]; then
checker=("$GITHUB_WORKSPACE/tests/base_sdk_tests/check_base_sdk_install.py")
(cd "$RUNNER_TEMP" && "$GITHUB_WORKSPACE/.venv-$extra/bin/python" -I "$GITHUB_WORKSPACE/tests/base_sdk_tests/check_sdk_http.py")
else
checker=("$GITHUB_WORKSPACE/scripts/check_mcp_sdk_install.py" --extra "$extra")
fi
@ -162,10 +163,14 @@ jobs:
uv venv --python ${{ matrix.python-version }} ".venv-core-$resolution"
uv pip install --python ".venv-core-$resolution" --resolution "$resolution" "$wheel"
uv pip check --python ".venv-core-$resolution"
(cd "$RUNNER_TEMP" && "$GITHUB_WORKSPACE/.venv-core-$resolution/bin/python" -I "$GITHUB_WORKSPACE/tests/base_sdk_tests/check_base_sdk_install.py")
(cd "$RUNNER_TEMP" && "$GITHUB_WORKSPACE/.venv-core-$resolution/bin/python" -I "$GITHUB_WORKSPACE/tests/base_sdk_tests/check_base_sdk_install.py" --profile core)
(cd "$RUNNER_TEMP" && "$GITHUB_WORKSPACE/.venv-core-$resolution/bin/python" -I "$GITHUB_WORKSPACE/tests/base_sdk_tests/check_sdk_http.py")
done
uv pip install --python .venv-core-highest coverage==7.14.0
.venv-core-highest/bin/python -I -m coverage run --rcfile=/dev/null --append --branch --include="*/tests/base_sdk_tests/check_base_sdk_install.py" tests/base_sdk_tests/check_base_sdk_install.py
.venv-core-highest/bin/python -I -m coverage run --rcfile=/dev/null --append --branch --include="*/tests/base_sdk_tests/check_base_sdk_install.py" tests/base_sdk_tests/check_base_sdk_install.py --profile core
uv pip install --python .venv-core-highest boto3 tokenizers huggingface-hub
(cd "$RUNNER_TEMP" && "$GITHUB_WORKSPACE/.venv-core-highest/bin/python" -I "$GITHUB_WORKSPACE/tests/base_sdk_tests/check_base_sdk_install.py" --profile dependencies)
(cd "$RUNNER_TEMP" && "$GITHUB_WORKSPACE/.venv-core-highest/bin/python" -I "$GITHUB_WORKSPACE/tests/base_sdk_tests/check_sdk_http.py")
uv run --no-project --with coverage==7.14.0 python -m coverage xml --rcfile=/dev/null -o coverage-core.xml
- name: Upload core packaging coverage
uses: codecov/codecov-action@0fb7174895f61a3b6b78fc075e0cd60383518dac # v5.5.5

3
litellm/integrations/focus/destinations/s3_destination.py
View file

@ -7,7 +7,6 @@ from collections.abc import Mapping
from datetime import timezone
from typing import Final, TypedDict
import boto3
from typing_extensions import ReadOnly
from .base import FocusDestination, FocusTimeWindow
@ -75,6 +74,8 @@ class FocusS3Destination(FocusDestination):
}
def _upload(self, content: bytes, object_key: str) -> None:
import boto3
s3_client: Final = boto3.client("s3", **self._client_kwargs())
s3_client.put_object(
Bucket=self.bucket_name,

8
litellm/litellm_core_utils/exception_mapping_utils.py
View file

@ -2376,6 +2376,10 @@ def exception_type(
return original_exception
if _is_guardrail_block(original_exception):
return original_exception
if isinstance(original_exception, ImportError) and (
original_exception.name in ("boto3", "botocore") or custom_llm_provider in ("bedrock", "bedrock_mantle")
):
return original_exception
exception_mapping_worked = False
exception_provider = custom_llm_provider
mappable_exception: Final[_ProviderHTTPException] = cast("_ProviderHTTPException", original_exception)
@ -2398,9 +2402,9 @@ def exception_type(
if model or custom_llm_provider:
if hasattr(original_exception, "message"):
error_str = (
redact_secret_string(str(original_exception.message))
redact_secret_string(str(mappable_exception.message))
if _ENABLE_SECRET_REDACTION
else str(original_exception.message)
else str(mappable_exception.message)
)
if isinstance(original_exception, BaseException):
exception_type = type(original_exception).__name__

13
litellm/litellm_core_utils/optional_imports.py Normal file
View file

@ -0,0 +1,13 @@
from typing import Final
def ensure_optional_import(module: str) -> None:
try:
__import__(module)
except ModuleNotFoundError as error:
if error.name != module:
raise
package: Final = "boto3" if module == "botocore" else module
raise ModuleNotFoundError(
f"Missing optional dependency '{module}'. Run 'pip install {package}'.", name=module
) from error

11
litellm/litellm_core_utils/token_counter.py
View file

@ -5,7 +5,7 @@ import io
import struct
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from itertools import accumulate
from typing import Final, Literal, cast
from typing import TYPE_CHECKING, Final, Literal, cast
import anyio
import anyio.lowlevel
@ -30,7 +30,10 @@ from litellm.constants import (
TOKEN_COUNTER_MAX_EXACT_CHARS,
)
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.tokenizer import Encoding, HuggingFace, HuggingFaceTokenizer, OpenAIEncoding
from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEncoding
if TYPE_CHECKING:
from litellm.litellm_core_utils.tokenizer import Encoding, HuggingFace
from litellm.litellm_core_utils.url_utils import safe_get
from litellm.llms.custom_httpx.http_handler import get_httpx_client
from litellm.rust_bridge.tokenizer import get_encoding
@ -680,13 +683,13 @@ def _get_exact_count_function(
raise ValueError("Unsupported tokenizer type")
def _encoding_count(encoding: Encoding, text: str) -> int:
def _encoding_count(encoding: "Encoding", text: str) -> int:
if isinstance(encoding, OpenAIEncoding):
return encoding.count(text)
return len(encoding.encode(text, disallowed_special=()))
def openai_tokenizer_encoding(model: str) -> Encoding:
def openai_tokenizer_encoding(model: str) -> "Encoding":
"""The encoding `token_counter` uses for a model on the `openai_tokenizer` path."""
return get_encoding(openai_tokenizer_encoding_name(model))

60
litellm/litellm_core_utils/tokenizer.py
View file

@ -14,16 +14,15 @@ from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from functools import partial
from pathlib import Path
from types import MappingProxyType
from types import MappingProxyType, UnionType
from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, runtime_checkable
import tiktoken
from tokenizers import AddedToken
from tokenizers import Tokenizer as PythonHuggingFaceTokenizer
if TYPE_CHECKING:
import numpy as np
import numpy.typing as npt
from tokenizers import Tokenizer as PythonHuggingFaceTokenizer
from litellm.rust_bridge._native import HuggingFaceEncoding
from litellm.rust_bridge._native import Tokenizer as NativeTokenizer
@ -217,6 +216,19 @@ class OpenAIEncoding:
return allowed
@dataclass(frozen=True, slots=True)
class _NativeAddedToken:
content: str
single_word: bool
lstrip: bool
rstrip: bool
normalized: bool
special: bool
def __str__(self) -> str:
return self.content
@dataclass(frozen=True, slots=True)
class HuggingFaceTokenizer:
"""The read-only ``tokenizers.Tokenizer`` surface over the Rust Hugging Face codec."""
@ -269,7 +281,14 @@ class HuggingFaceTokenizer:
def get_vocab_size(self, with_added_tokens: bool = True) -> int:
return self._native.get_vocab_size(with_added_tokens)
def get_added_tokens_decoder(self) -> dict[int, AddedToken]: # mutable-ok: [LIT001] SDK return type
def get_added_tokens_decoder(self) -> dict[int, _AddedToken]: # mutable-ok: [LIT001] SDK return type
try:
from tokenizers import AddedToken
except ModuleNotFoundError as error:
if error.name != "tokenizers":
raise
return {token_id: _NativeAddedToken(*data) for token_id, data in self._native.added_tokens_decoder()}
return {
token_id: AddedToken(
content, single_word=single_word, lstrip=lstrip, rstrip=rstrip, normalized=normalized, special=special
@ -361,11 +380,40 @@ def _batch_input(
Encoding: TypeAlias = tiktoken.Encoding | OpenAIEncoding
HuggingFace: TypeAlias = PythonHuggingFaceTokenizer | HuggingFaceTokenizer
Tokenizer: TypeAlias = Encoding | HuggingFace
if TYPE_CHECKING:
HuggingFace: TypeAlias = PythonHuggingFaceTokenizer | HuggingFaceTokenizer
Tokenizer: TypeAlias = Encoding | HuggingFace
def __getattr__(name: str) -> UnionType | type[HuggingFaceTokenizer]:
if name not in {"HuggingFace", "Tokenizer"}:
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
available: Final = HuggingFaceTokenizer if name == "HuggingFace" else Encoding | HuggingFaceTokenizer
try:
from tokenizers import Tokenizer as PythonTokenizer
except ModuleNotFoundError as error:
if error.name == "tokenizers":
return available
raise
return available | PythonTokenizer
class _AddedToken(Protocol):
@property
def content(self) -> str: ...
@property
def single_word(self) -> bool: ...
@property
def lstrip(self) -> bool: ...
@property
def rstrip(self) -> bool: ...
@property
def normalized(self) -> bool: ...
@property
def special(self) -> bool: ...

9
litellm/llms/aws_polly/text_to_speech/transformation.py
View file

@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, Final, Union
import httpx
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.litellm_core_utils.optional_imports import ensure_optional_import
from litellm.llms.base_llm.text_to_speech.transformation import (
BaseTextToSpeechConfig,
TextToSpeechRequestData,
@ -262,11 +263,9 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM):
Returns:
Tuple of (signed_headers, json_body_string)
"""
try:
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError("Missing boto3 to call AWS Polly. Run 'pip install boto3'.")
ensure_optional_import("botocore")
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
# Get AWS region
aws_region_name: Final = litellm_params.get("aws_region_name", self.DEFAULT_REGION)

62
litellm/llms/bedrock/base_aws_llm.py
View file

@ -32,9 +32,10 @@ from litellm.constants import (
)
from litellm.litellm_core_utils.aws_partition import contains_bedrock_arn, get_aws_dns_suffix
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.optional_imports import ensure_optional_import
from litellm.secret_managers.main import get_secret, get_secret_str
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams, AwsSessionTag
from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams, AwsSessionTag, BearerPreparedRequest
if TYPE_CHECKING:
from botocore.awsrequest import AWSPreparedRequest
@ -446,6 +447,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
# iam_cache: static keys, ambient env (including skip-AssumeRole path), web identity, and
# AssumeRole. Do not cache profile / explicit session-token paths here.
#########################################################
ensure_optional_import("botocore")
if self._is_auth_with_web_identity_token(
aws_web_identity_token,
aws_role_name,
@ -775,6 +777,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
aws_region_name = standard_aws_region_name
if aws_region_name is None:
try:
ensure_optional_import("boto3")
import boto3
with tracer.trace("boto3.Session()"):
@ -958,6 +961,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
# For ECS/EC2: call sts:GetCallerIdentity to check if already running as the role
try:
ensure_optional_import("boto3")
import boto3
with tracer.trace("boto3.client(sts).get_caller_identity"):
@ -1017,6 +1021,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
"""
Authenticate with AWS Web Identity Token
"""
ensure_optional_import("boto3")
import boto3
verbose_logger.debug(
@ -1116,6 +1121,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
aws_session_tags: Sequence[AwsSessionTag] | None = None,
) -> dict:
"""Handle cross-account role assumption for IRSA."""
ensure_optional_import("boto3")
import boto3
verbose_logger.debug("Cross-account role assumption detected")
@ -1181,6 +1187,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
aws_session_tags: Sequence[AwsSessionTag] | None = None,
) -> dict:
"""Handle same-account role assumption for IRSA."""
ensure_optional_import("boto3")
import boto3
irsa_sts_kwargs: Final = self._build_sts_client_kwargs(
@ -1287,6 +1294,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
"""
Authenticate with AWS Role
"""
ensure_optional_import("boto3")
import boto3
from botocore.credentials import Credentials
@ -1410,6 +1418,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
"""
Authenticate with AWS profile
"""
ensure_optional_import("boto3")
import boto3
# uses auth values from AWS profile usually stored in ~/.aws/credentials
@ -1448,6 +1457,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
"""
Authenticate with AWS Access Key and Secret Key
"""
ensure_optional_import("boto3")
import boto3
# Check if credentials are already in cache. These credentials have no expiry time.
@ -1468,6 +1478,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
"""
Authenticate with AWS Environment Variables
"""
ensure_optional_import("boto3")
import boto3
with tracer.trace("boto3.Session()"):
@ -1561,10 +1572,6 @@ class BaseAWSLLM(SignsRequestsWithAWS):
Returns:
Credentials: Boto3 credentials object
"""
try:
from botocore.credentials import Credentials
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
aws_region_name: Final = self._get_aws_region_name(optional_params, model)
optional_params.pop("aws_region_name", None)
auth_params: Final = pop_aws_auth_params(optional_params)
@ -1594,23 +1601,26 @@ class BaseAWSLLM(SignsRequestsWithAWS):
headers: dict,
api_key: str | None = None,
supports_bearer_token: bool = True,
) -> AWSPreparedRequest:
) -> AWSPreparedRequest | BearerPreparedRequest:
aws_bearer_token: Final = bedrock_bearer_token(api_key) if supports_bearer_token else None
if aws_bearer_token is not None:
try:
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
headers["Authorization"] = f"Bearer {aws_bearer_token}"
request = AWSRequest(method="POST", url=endpoint_url, data=data, headers=headers)
bearer_request: Final = httpx.Request("POST", endpoint_url, content=data, headers=headers)
return BearerPreparedRequest(
method="POST",
url=str(bearer_request.url),
headers={
name.decode("ascii"): value.decode(bearer_request.headers.encoding)
for name, value in bearer_request.headers.raw
},
body=bearer_request.content,
)
else:
try:
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.exceptions import NoCredentialsError
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
ensure_optional_import("botocore")
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.exceptions import NoCredentialsError
if credentials is None:
raise NoCredentialsError()
@ -1703,12 +1713,10 @@ class BaseAWSLLM(SignsRequestsWithAWS):
return headers, json.dumps(request_data).encode()
# If no bearer token is set, proceed with the existing SigV4 authentication
try:
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.credentials import Credentials
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
ensure_optional_import("botocore")
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.credentials import Credentials
auth_params: Final = AwsAuthParams.model_validate(optional_params)
aws_region_name: Final = self._get_aws_region_name(optional_params=optional_params, model=model)
@ -1754,11 +1762,9 @@ def sign_aws_json_post(
body: str,
headers: Mapping[str, str],
) -> AWSPreparedRequest:
try:
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError(f"Missing boto3 to call {service_name}. Run 'pip install boto3'.")
ensure_optional_import("botocore")
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
aws_request: Final = AWSRequest(method="POST", url=url, data=body, headers=headers)
SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request)

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

@ -12,6 +12,7 @@ from litellm.caching.caching import InMemoryCache
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
from litellm.litellm_core_utils.core_helpers import map_finish_reason
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.optional_imports import ensure_optional_import
from litellm.llms.anthropic.chat.handler import (
ModelResponseIterator as AnthropicModelResponseIterator,
)
@ -264,7 +265,7 @@ async def make_call(
)
return completion_stream, response.headers
except BedrockError:
except (BedrockError, ImportError):
raise
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
@ -355,7 +356,7 @@ def make_sync_call(
)
return completion_stream, response.headers
except BedrockError:
except (BedrockError, ImportError):
raise
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
@ -440,6 +441,7 @@ class _EventStreamTally:
class AWSEventStreamDecoder:
def __init__(self, model: str, json_mode: bool | None = False) -> None:
ensure_optional_import("botocore")
from botocore.parsers import EventStreamJSONParser
self.model = model

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

@ -27,6 +27,7 @@ from pydantic import ConfigDict, TypeAdapter, ValidationError
import litellm
from litellm import verbose_logger
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.litellm_core_utils.optional_imports import ensure_optional_import
from litellm.llms.base_llm.anthropic_messages.transformation import (
BaseAnthropicMessagesConfig,
)
@ -1843,6 +1844,7 @@ class BedrockEventStreamDecoderBase:
"""
def __init__(self):
ensure_optional_import("botocore")
from botocore.parsers import EventStreamJSONParser
self.parser = EventStreamJSONParser()
@ -2063,11 +2065,9 @@ class CommonBatchFilesUtils:
Returns:
Tuple of (signed_headers, signed_data)
"""
try:
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
ensure_optional_import("botocore")
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
aws_region_name: Final = self._base_aws.get_aws_region_name(optional_params=optional_params, model="")
credentials: Final = self._base_aws.resolve_credentials(

3
litellm/llms/bedrock/image_edit/handler.py
View file

@ -27,6 +27,7 @@ from litellm.llms.custom_httpx.http_handler import (
get_httpx_client,
)
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.bedrock import BearerPreparedRequest
from litellm.types.utils import ImageResponse
from ..base_aws_llm import BaseAWSLLM, bedrock_bearer_token
@ -44,7 +45,7 @@ class BedrockImageEditPreparedRequest(LiteLLMBaseModel):
"""
endpoint_url: str
prepped: AWSPreparedRequest
prepped: AWSPreparedRequest | BearerPreparedRequest
body: bytes
data: dict

3
litellm/llms/bedrock/image_generation/image_handler.py
View file

@ -27,6 +27,7 @@ from litellm.llms.custom_httpx.http_handler import (
get_httpx_client,
)
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.bedrock import BearerPreparedRequest
from litellm.types.utils import ImageResponse
from ..base_aws_llm import BaseAWSLLM, bedrock_bearer_token
@ -44,7 +45,7 @@ class BedrockImagePreparedRequest(LiteLLMBaseModel):
"""
endpoint_url: str
prepped: AWSPreparedRequest
prepped: AWSPreparedRequest | BearerPreparedRequest
body: bytes
data: dict

68
litellm/llms/bedrock_mantle/common_utils.py
View file

@ -14,15 +14,10 @@ global state.
import re
from collections.abc import Mapping
from functools import partial
from typing import Final, Literal
from botocore.exceptions import (
CredentialRetrievalError,
NoCredentialsError,
PartialCredentialsError,
ProfileNotFound,
)
from litellm.litellm_core_utils.optional_imports import ensure_optional_import
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS
from litellm.llms.bedrock.common_utils import AmazonBedrockGlobalConfig
from litellm.secret_managers.main import get_secret_str
@ -72,7 +67,7 @@ class BedrockMantleAuthMixin(SignsRequestsWithAWS):
return resolve_mantle_bearer_token(api_key)
@staticmethod
def _resolve_region(params: dict) -> str:
def _resolve_region(params: Mapping[str, object]) -> str:
return resolve_mantle_region(params)
def sign_request(
@ -87,32 +82,39 @@ class BedrockMantleAuthMixin(SignsRequestsWithAWS):
fake_stream: bool | None = None,
) -> tuple[dict, bytes | None]:
bearer: Final = self._resolve_bearer_token(api_key)
if not bearer:
# Pin the credential-scope region to the region of the actual signing URL
# so the SigV4 scope and URL host can never disagree, even when a stale
# api_base and aws_region_name point at different regions.
host_match: Final = MANTLE_HOST_RE.match(api_base.rstrip("/"))
optional_params = {
**optional_params,
"aws_region_name": (
host_match.group(1)
if host_match
else self._resolve_region({**optional_params, "api_base": api_base})
),
}
headers = {k: v for k, v in headers.items() if k.lower() != "authorization"}
sign: Final[partial[tuple[dict[str, str | bytes], bytes | None]]] = partial(
self._aws_signer._sign_request,
service_name="bedrock",
request_data=request_data,
api_base=api_base,
api_key=bearer,
model=model,
stream=stream,
fake_stream=fake_stream,
)
if bearer:
return sign(headers=headers, optional_params=optional_params)
ensure_optional_import("botocore")
from botocore.exceptions import (
CredentialRetrievalError,
NoCredentialsError,
PartialCredentialsError,
ProfileNotFound,
)
# Pin the credential-scope region to the region of the actual signing URL
# so the SigV4 scope and URL host can never disagree, even when a stale
# api_base and aws_region_name point at different regions.
host_match: Final = MANTLE_HOST_RE.match(api_base.rstrip("/"))
optional_params = {
**optional_params,
"aws_region_name": (
host_match.group(1) if host_match else self._resolve_region({**optional_params, "api_base": api_base})
),
}
headers = {k: v for k, v in headers.items() if k.lower() != "authorization"}
try:
return self._aws_signer._sign_request(
service_name="bedrock",
headers=headers,
optional_params=optional_params,
request_data=request_data,
api_base=api_base,
api_key=bearer,
model=model,
stream=stream,
fake_stream=fake_stream,
)
return sign(headers=headers, optional_params=optional_params)
except (
NoCredentialsError,
PartialCredentialsError,

16
litellm/llms/sagemaker/chat/handler.py
View file

@ -6,6 +6,7 @@ from typing import Final
import httpx
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.litellm_core_utils.optional_imports import ensure_optional_import
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, pop_aws_auth_params
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.utils import ModelResponse, get_secret
@ -19,10 +20,9 @@ class SagemakerChatHandler(BaseAWSLLM):
self,
optional_params: dict,
):
try:
from botocore.credentials import Credentials
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
ensure_optional_import("botocore")
from botocore.credentials import Credentials
auth_params: Final = pop_aws_auth_params(optional_params)
aws_region_name = optional_params.pop("aws_region_name", None)
optional_params.pop("aws_bedrock_runtime_endpoint", None)
@ -54,11 +54,9 @@ class SagemakerChatHandler(BaseAWSLLM):
aws_region_name: str,
extra_headers: dict | None = None,
):
try:
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
ensure_optional_import("botocore")
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
sigv4: Final = SigV4Auth(credentials, "sagemaker", aws_region_name)
dns_suffix: Final = get_aws_dns_suffix(aws_region_name)

16
litellm/llms/sagemaker/completion/handler.py
View file

@ -10,6 +10,7 @@ from litellm._logging import verbose_logger
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.optional_imports import ensure_optional_import
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, pop_aws_auth_params
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
@ -42,10 +43,9 @@ class SagemakerLLM(BaseAWSLLM):
self,
optional_params: dict,
):
try:
from botocore.credentials import Credentials
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
ensure_optional_import("botocore")
from botocore.credentials import Credentials
auth_params: Final = pop_aws_auth_params(optional_params)
aws_region_name = optional_params.pop("aws_region_name", None)
optional_params.pop("aws_bedrock_runtime_endpoint", None)
@ -79,11 +79,9 @@ class SagemakerLLM(BaseAWSLLM):
aws_region_name: str,
extra_headers: dict | None = None,
):
try:
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
ensure_optional_import("botocore")
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
sigv4: Final = SigV4Auth(credentials, "sagemaker", aws_region_name)
dns_suffix: Final = get_aws_dns_suffix(aws_region_name)

22
litellm/rust_bridge/tokenizer.py
View file

@ -4,14 +4,16 @@ from functools import lru_cache
from typing import TYPE_CHECKING, Final, cast # noqa: TID251 # native class is validated at the binding boundary
import tiktoken
from tokenizers import Tokenizer as PythonHuggingFaceTokenizer
from litellm.litellm_core_utils.tokenizer import Encoding, HuggingFace, HuggingFaceTokenizer, OpenAIEncoding
from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer, OpenAIEncoding
from litellm.rust_bridge import runtime
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.catalog import Route, RouteContext
if TYPE_CHECKING:
from tokenizers import Tokenizer as PythonHuggingFaceTokenizer
from litellm.litellm_core_utils.tokenizer import Encoding, HuggingFace
from litellm.rust_bridge._native import Tokenizer as NativeTokenizer
@ -76,6 +78,16 @@ def get_encoding(name: str) -> Encoding:
)
def _python_tokenizer() -> type[PythonHuggingFaceTokenizer]:
try:
from tokenizers import Tokenizer
except ModuleNotFoundError as error:
if error.name != "tokenizers":
raise
raise ImportError("Python tokenization requires tokenizers. Run 'pip install tokenizers'.") from error
return Tokenizer
def anthropic() -> HuggingFace:
"""The packaged Anthropic tokenizer on the selected backend."""
from litellm.utils import claude_json_str
@ -84,7 +96,7 @@ def anthropic() -> HuggingFace:
HUGGINGFACE_CONTEXT,
binding=TOKENIZER,
native=lambda factory: HuggingFaceTokenizer(_native_anthropic(factory)),
python=lambda: PythonHuggingFaceTokenizer.from_str(claude_json_str),
python=lambda: _python_tokenizer().from_str(claude_json_str),
)
@ -93,7 +105,7 @@ def from_str(json: str) -> HuggingFace:
HUGGINGFACE_CONTEXT,
binding=TOKENIZER,
native=lambda factory: HuggingFaceTokenizer(factory.from_json(json)),
python=lambda: PythonHuggingFaceTokenizer.from_str(json),
python=lambda: _python_tokenizer().from_str(json),
)
@ -104,5 +116,5 @@ def from_pretrained(identifier: str, revision: str = "main", token: str | None =
native=lambda factory: HuggingFaceTokenizer(
factory.from_pretrained(identifier, revision=revision, token=token)
),
python=lambda: PythonHuggingFaceTokenizer.from_pretrained(identifier, revision=revision, token=token),
python=lambda: _python_tokenizer().from_pretrained(identifier, revision=revision, token=token),
)

10
litellm/secret_managers/aws_secret_manager_v2.py
View file

@ -25,6 +25,7 @@ import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.litellm_core_utils.optional_imports import ensure_optional_import
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
@ -627,11 +628,10 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
request_data: dict | None = None,
) -> tuple[str, "HTTPHeaders", bytes]:
"""Prepare the AWS Secrets Manager request"""
try:
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
ensure_optional_import("botocore")
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
optional_params = optional_params or {}
# Build optional_params from instance settings if not provided

13
litellm/types/llms/bedrock.py
View file

@ -1,5 +1,6 @@
import json
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from enum import Enum
from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias
@ -1009,13 +1010,21 @@ else:
AWSPreparedRequest = Any
@dataclass(frozen=True)
class BearerPreparedRequest:
method: str
url: str
headers: Mapping[str, str]
body: bytes
class BedrockPreparedRequest(TypedDict):
"""
Internal/Helper class for preparing the request for bedrock image generation
"""
endpoint_url: str
prepped: AWSPreparedRequest
prepped: AWSPreparedRequest | BearerPreparedRequest
body: bytes
data: dict

19
litellm/utils.py
View file

@ -59,6 +59,7 @@ from litellm._lazy_imports import (
get_messages_reach_token_count,
get_token_counter_new,
)
from litellm._logging import redact_secrets
from litellm._uuid import uuid
from litellm.constants import (
DEFAULT_CHAT_COMPLETION_PARAM_VALUES,
@ -90,7 +91,7 @@ from litellm.litellm_core_utils.fallback_generalizations import (
match_fill_missing_generalizations,
)
from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload
from litellm.litellm_core_utils.tokenizer import Encoding, HuggingFace, strip_special_tokens
from litellm.litellm_core_utils.tokenizer import strip_special_tokens
from litellm.rust_bridge import tokenizer as tokenizer_dispatch
from litellm.rust_bridge.catalog import decision
from litellm.rust_bridge.configuration import Decision
@ -369,6 +370,7 @@ if TYPE_CHECKING:
from litellm.litellm_core_utils.rules import Rules
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.litellm_core_utils.thread_pool_executor import BoundedLoggingThreadPoolExecutor
from litellm.litellm_core_utils.tokenizer import Encoding, HuggingFace
from litellm.llms.base_llm.anthropic_messages.transformation import (
BaseAnthropicMessagesConfig,
)
@ -2239,7 +2241,7 @@ def client(original_function):
is_acompletion_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True)
if (
num_retries and not is_acompletion_litellm_router_call
num_retries and not is_acompletion_litellm_router_call and not isinstance(e, ImportError)
): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying
try:
litellm.num_retries = None # set retries to None to prevent infinite loops
@ -2270,7 +2272,7 @@ def client(original_function):
is_aresponses_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True)
if (
num_retries and not is_aresponses_litellm_router_call
num_retries and not is_aresponses_litellm_router_call and not isinstance(e, ImportError)
): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying
try:
litellm.num_retries = None # set retries to None to prevent infinite loops
@ -2423,7 +2425,12 @@ def _select_tokenizer_helper(model: str) -> SelectTokenizerResponse:
if isinstance(e, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking)):
raise
verbose_logger.debug("Error selecting tokenizer: %s", e)
verbose_logger.warning(
"Falling back to tiktoken for %s; token counts may be approximate. "
"For Python Hugging Face tokenization, install tokenizers and huggingface-hub. Error: %s",
json.dumps(redact_secrets(model)),
json.dumps(redact_secrets(str(e))),
)
# default - tiktoken
return _return_openai_tokenizer(model)
@ -2494,7 +2501,7 @@ def encode(model="", text="", custom_tokenizer: dict | None = None):
tokenizer_json: Final = custom_tokenizer or select_tokenizer(model=model)
if tokenizer_json["type"] == "openai_tokenizer":
openai_tokenizer: Final = cast( # cast-ok: [LIT006] caller's explicit type tag selects this interface
Encoding, tokenizer_json["tokenizer"]
"Encoding", tokenizer_json["tokenizer"]
)
return openai_tokenizer.encode(text, disallowed_special=())
encoded: Final = tokenizer_json["tokenizer"].encode(text)
@ -2519,7 +2526,7 @@ def decode(
if tokenizer_json["type"] == "huggingface_tokenizer":
ids: Final = strip_special_tokens(tokenizer_json["tokenizer"], tokens) if skip_special_tokens else tokens
hf_tokenizer: Final = cast( # cast-ok: [LIT006] caller's explicit type tag selects this interface
HuggingFace, tokenizer_json["tokenizer"]
"HuggingFace", tokenizer_json["tokenizer"]
)
return hf_tokenizer.decode(ids, skip_special_tokens=skip_special_tokens)
return tokenizer_json["tokenizer"].decode(tokens)

4
packaging/litellm-core/pyproject.toml
View file

@ -14,14 +14,13 @@ dependencies = [
"filelock>=3.16.1,<4.0",
"httpx[http2]>=0.28.0,<1.0",
"openai>=2.20.0,<3.0.0",
"python-dateutil>=2.8.2,<3.0",
"python-dotenv>=1.0.0,<2.0",
"pyyaml>=6.0.3,<7.0",
"packaging>=24.0",
"importlib-metadata>=8.0.0,<9.0",
"tiktoken>=0.8.0,<1.0; python_version < '3.14'",
"tiktoken>=0.12.0,<1.0; python_version >= '3.14'",
"tokenizers>=0.21.0,<1.0",
"huggingface-hub>=0.34.0,<2.0",
"click>=8.0.0,<9.0",
"jinja2>=3.1.6,<4.0",
"aiohttp>=3.14.2,<4.0",
@ -30,7 +29,6 @@ dependencies = [
"pydantic>=2.12.0,<3.0.0; python_version >= '3.14'",
"pydantic-settings>=2.14.1,<3.0",
"jsonschema>=4.0.0,<5.0",
"boto3>=1.43.1,<2.0",
"typing-extensions>=4.13.0,<5.0",
]

48
tests/base_sdk_tests/check_base_sdk_install.py
View file

@ -6,6 +6,7 @@ pull ``packaging``, ``pluggy`` and ``iniconfig`` into the environment and could
the very class of undeclared-dependency bug this guards against.
"""
import argparse
import importlib.util
import sys
import traceback
@ -33,7 +34,9 @@ def check_import() -> str:
from importlib.metadata import version
import litellm
from litellm.llms.brave.search.transformation import BraveSearchConfig
_require(callable(BraveSearchConfig), "Brave search configuration unavailable")
_require(bool(litellm.__file__), "litellm has no __file__")
from litellm._version import version as sdk_version
@ -109,12 +112,47 @@ def check_token_counter() -> str:
return f"token_counter returned {count}"
def check_tokenizer_dependencies() -> str:
import litellm
from litellm.rust_bridge import tokenizer
from litellm.litellm_core_utils.tokenizer import HuggingFace, HuggingFaceTokenizer, Tokenizer
from litellm.utils import claude_json_str
_require(isinstance(litellm.encoding, Tokenizer), "runtime alias rejects the default encoding")
native = tokenizer.native_anthropic()
if native is not None:
_require(bool(native.encode("hello")), "native tokenizer returned no tokens")
_require(isinstance(HuggingFaceTokenizer(native), HuggingFace), "runtime alias rejects native tokenizers")
if importlib.util.find_spec("tokenizers") is not None:
python_tokenizer = tokenizer.from_str(claude_json_str)
_require(bool(python_tokenizer.encode("hello").ids), "Python tokenizer returned no tokens")
_require(isinstance(python_tokenizer, HuggingFace), "runtime alias rejects Python tokenizers")
return "installed Python tokenizer and available native tokenizer work"
try:
tokenizer.from_str(claude_json_str)
except ImportError as error:
_require("pip install tokenizers" in str(error), f"missing tokenizer guidance: {error}")
else:
raise AssertionError("Python tokenizer loaded without tokenizers")
return "native tokenizer works; Python tokenizer reports its missing dependency"
def check_bedrock_credential_resolution() -> str:
import os
from unittest import mock
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
if importlib.util.find_spec("boto3") is None:
try:
BaseAWSLLM()._sign_request(
service_name="bedrock", headers={}, optional_params={"aws_region_name": "us-east-1"},
request_data={}, api_base="https://bedrock-runtime.us-east-1.amazonaws.com", api_key="",
)
except ImportError as error:
_require("pip install boto3" in str(error), f"missing installation guidance: {error}")
return "AWS signing explains how to install boto3"
raise AssertionError("AWS signing unexpectedly worked without boto3")
non_aws_environ = {k: v for k, v in os.environ.items() if not k.startswith("AWS_")}
with mock.patch.dict(os.environ, non_aws_environ, clear=True):
credentials = BaseAWSLLM().get_credentials(
@ -137,6 +175,7 @@ CHECKS: tuple[tuple[str, Callable[[], str]], ...] = (
("embedding", check_embedding),
("bundled model metadata", check_bundled_model_metadata),
("token counter", check_token_counter),
("tokenizer dependencies", check_tokenizer_dependencies),
("bedrock credential resolution", check_bedrock_credential_resolution),
)
@ -149,7 +188,14 @@ def _run(check: Callable[[], str]) -> tuple[bool, str]:
def main() -> int:
print(f"base SDK smoke check on {sys.executable}")
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--profile", choices=("legacy", "core", "dependencies"), default="legacy")
profile = parser.parse_args().profile
for module in ("boto3", "botocore", "tokenizers", "huggingface_hub"):
present = importlib.util.find_spec(module) is not None
_require(present == (profile != "core"), f"{profile}: unexpected presence of {module}: {present}")
_require(importlib.util.find_spec("jsonschema") is not None, "jsonschema must remain mandatory")
print(f"{profile} SDK smoke check on {sys.executable}")
for label, check in CHECKS:
passed, detail = _run(check)
if not passed:

202
tests/base_sdk_tests/check_sdk_http.py Normal file
View file

@ -0,0 +1,202 @@
"""Check installed SDK HTTP behavior against a recording loopback upstream."""
import asyncio
import importlib.util
import json
import os
import struct
import zlib
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from queue import Queue
from threading import Event, Thread
from typing import Final
from unittest.mock import patch
ENVIRONMENT: Final = {"LITELLM_LOCAL_MODEL_COST_MAP": "True", "PYTHON_DOTENV_DISABLED": "1"}
with patch.dict(os.environ, ENVIRONMENT):
import litellm
RESPONSES: Final[Queue[tuple[int, bytes]]] = Queue()
REQUESTS: Final[Queue[tuple[str, dict[str, str], bytes]]] = Queue()
ARRIVED: Final = Event()
RELEASE: Final = Event()
MESSAGES: Final = [{"role": "user", "content": "ping"}]
CHAT: Final = {
"id": "chat-http-check",
"object": "chat.completion",
"created": 1,
"model": "gpt-4o",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "pong"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
}
class RecordingUpstream(BaseHTTPRequestHandler):
def log_message(self, format: str, *args: object) -> None:
pass
def do_POST(self) -> None:
REQUESTS.put((self.path, dict(self.headers), self.rfile.read(int(self.headers["Content-Length"]))))
status, body = RESPONSES.get(timeout=10)
ARRIVED.set()
if status == 0:
RELEASE.wait(timeout=10)
return
self.send_response(status)
self.send_header("Content-Type", "text/event-stream" if body.startswith(b"data:") else "application/vnd.amazon.eventstream" if body[:1] == b"\x00" else "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def enqueue(body: object, status: int = 200) -> None:
RESPONSES.put((status, json.dumps(body).encode()))
def enqueue_stream() -> None:
chunk: Final = {**CHAT, "object": "chat.completion.chunk", "choices": [{"index": 0, "delta": {"content": "pong"}}]}
RESPONSES.put((200, f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode()))
def bedrock_event(event: str, payload: object) -> bytes:
def encode_header(name: str, value: str) -> bytes:
return bytes([len(name)]) + name.encode() + b"\x07" + struct.pack(">H", len(value)) + value.encode()
headers: Final = encode_header(":message-type", "event") + encode_header(":event-type", event)
content: Final = json.dumps(payload).encode()
prelude: Final = struct.pack(">II", len(headers) + len(content) + 16, len(headers))
frame: Final = prelude + struct.pack(">I", zlib.crc32(prelude)) + headers + content
return frame + struct.pack(">I", zlib.crc32(frame))
def check_http(base: str) -> None:
arguments: Final = dict(model="openai/gpt-4o", messages=MESSAGES, api_key="test-key", api_base=base + "/v1")
enqueue(CHAT)
response: Final = litellm.completion(**arguments)
assert response.choices[0].message.content == "pong"
assert response.usage.total_tokens == 5
path, headers, body = REQUESTS.get(timeout=10)
assert path == "/v1/chat/completions"
assert headers["Authorization"] == "Bearer test-key"
assert json.loads(body)["messages"] == MESSAGES
enqueue_stream()
assert "".join(part.choices[0].delta.content or "" for part in litellm.completion(**arguments, stream=True)) == "pong"
REQUESTS.get(timeout=10)
for status, exception in ((401, litellm.AuthenticationError), (429, litellm.RateLimitError), (500, litellm.InternalServerError)):
enqueue({"error": {"message": "controlled upstream failure"}}, status)
try:
litellm.completion(**arguments, num_retries=0, max_retries=0)
except exception:
REQUESTS.get(timeout=10)
else:
raise AssertionError(f"HTTP {status} did not raise {exception.__name__}")
enqueue({"error": {"message": "retry once"}}, 429)
enqueue(CHAT)
assert litellm.completion(**arguments, num_retries=1).choices[0].message.content == "pong"
REQUESTS.get(timeout=10)
REQUESTS.get(timeout=10)
async def check_async() -> None:
enqueue(CHAT)
result: Final = await litellm.acompletion(**arguments)
assert result.choices[0].message.content == "pong"
REQUESTS.get(timeout=10)
enqueue_stream()
stream: Final = await litellm.acompletion(**arguments, stream=True)
assert "".join([part.choices[0].delta.content or "" async for part in stream]) == "pong"
REQUESTS.get(timeout=10)
ARRIVED.clear()
RESPONSES.put((0, b""))
task: Final = asyncio.create_task(litellm.acompletion(**arguments, num_retries=0, max_retries=0))
assert await asyncio.to_thread(ARRIVED.wait, 10), "request did not reach upstream"
task.cancel()
try:
await task
except asyncio.CancelledError:
REQUESTS.get(timeout=10)
else:
raise AssertionError("cancellation did not propagate")
finally:
RELEASE.set()
asyncio.run(check_async())
enqueue({"object": "list", "model": "text-embedding-3-small",
"data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}],
"usage": {"prompt_tokens": 3, "total_tokens": 3}})
embedding: Final = litellm.embedding(model="text-embedding-3-small", input=["ping"], api_key="test-key", api_base=base + "/v1")
assert embedding.data[0]["embedding"] == [0.1, 0.2]
assert REQUESTS.get(timeout=10)[0] == "/v1/embeddings"
RELEASE.clear()
RESPONSES.put((0, b""))
try:
litellm.completion(**arguments, timeout=0.1, num_retries=0, max_retries=0)
except litellm.Timeout:
REQUESTS.get(timeout=10)
else:
raise AssertionError("stalled upstream did not time out")
finally:
RELEASE.set()
for signed in ((False, True) if importlib.util.find_spec("boto3") else (False,)):
for asynchronous in (False, True):
enqueue({
"output": {"message": {"role": "assistant", "content": [{"text": "pong"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 3, "outputTokens": 2, "totalTokens": 5},
"metrics": {"latencyMs": 1},
})
bedrock: Final = dict(
model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", messages=MESSAGES,
aws_region_name="us-east-1", aws_bedrock_runtime_endpoint=base,
**({"aws_access_key_id": "test-key", "aws_secret_access_key": "test-secret", "api_key": ""}
if signed else {"api_key": "bearer-key"}),
)
result: Final = asyncio.run(litellm.acompletion(**bedrock)) if asynchronous else litellm.completion(**bedrock)
assert result.choices[0].message.content == "pong"
assert result.usage.total_tokens == 5
path, headers, body = REQUESTS.get(timeout=10)
assert path.endswith("/converse")
assert {key.lower(): value for key, value in headers.items()}["authorization"].startswith(
"AWS4-HMAC-SHA256" if signed else "Bearer bearer-key"
)
assert json.loads(body)["messages"][0]["content"][0]["text"] == "ping"
if importlib.util.find_spec("botocore") is not None:
RESPONSES.put((200, b"".join((
bedrock_event("messageStart", {"role": "assistant"}),
bedrock_event("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "pong"}}),
bedrock_event("messageStop", {"stopReason": "end_turn"}),
bedrock_event("metadata", {"usage": {"inputTokens": 3, "outputTokens": 2, "totalTokens": 5}}),
))))
streamed: Final = litellm.completion(
model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", messages=MESSAGES,
api_key="bearer-key", aws_region_name="us-east-1", aws_bedrock_runtime_endpoint=base, stream=True,
)
assert "".join(part.choices[0].delta.content or "" for part in streamed) == "pong"
assert REQUESTS.get(timeout=10)[0].endswith("/converse-stream")
enqueue({"id": "msg-http-check", "type": "message", "role": "assistant", "model": "claude-3-sonnet-20240229",
"content": [{"type": "text", "text": "pong"}], "stop_reason": "end_turn", "stop_sequence": None,
"usage": {"input_tokens": 3, "output_tokens": 2}})
anthropic: Final = litellm.completion(model="anthropic/claude-3-sonnet-20240229", messages=MESSAGES,
api_key="test-key", api_base=base, max_tokens=16)
assert anthropic.choices[0].message.content == "pong"
assert anthropic.usage.total_tokens == 5
assert REQUESTS.get(timeout=10)[0] == "/v1/messages"
print("PASS installed HTTP: sync/async, streaming, usage, error mapping, retries, cancellation, Bedrock auth")
def main() -> None:
with ThreadingHTTPServer(("127.0.0.1", 0), RecordingUpstream) as server:
worker: Final = Thread(target=server.serve_forever, daemon=True)
worker.start()
try:
check_http(f"http://127.0.0.1:{server.server_port}")
assert RESPONSES.empty() and REQUESTS.empty(), "unconsumed HTTP exchanges"
finally:
RELEASE.set()
server.shutdown()
worker.join(timeout=10)
if __name__ == "__main__":
with patch.dict(os.environ, ENVIRONMENT):
main()

15
tests/base_sdk_tests/test_core_distribution.py
View file

@ -59,7 +59,18 @@ def test_core_manifest_preserves_runtime_dependencies_without_extras() -> None:
core: Final = tomllib.loads(manifest.read_text())
legacy: Final = tomllib.loads((ROOT / "pyproject.toml").read_text())
assert core["project"]["name"] == "litellm-core"
assert core["project"]["dependencies"] == legacy["project"]["dependencies"]
removed: Final = {"boto3", "tokenizers", "huggingface-hub"}
core_dependencies: Final = {Requirement(value).name for value in core["project"]["dependencies"]}
legacy_dependencies: Final = {Requirement(value).name for value in legacy["project"]["dependencies"]}
assert not core_dependencies & removed
assert removed <= legacy_dependencies
assert "jsonschema" in core_dependencies
assert legacy_dependencies - removed <= core_dependencies
core_requirements: Final = {Requirement(value) for value in core["project"]["dependencies"]}
retained_requirements: Final = {
Requirement(value) for value in legacy["project"]["dependencies"] if Requirement(value).name not in removed
}
assert retained_requirements <= core_requirements
assert not core["project"].get("optional-dependencies")
assert not core["project"].get("scripts")
@ -110,7 +121,7 @@ def test_core_wheel_metadata_and_resources(distributions: tuple[Path, Path]) ->
assert metadata["Version"] == tomllib.loads((ROOT / "pyproject.toml").read_text())["project"]["version"]
assert not metadata.get_all("Provides-Extra")
assert not any(n.endswith(".dist-info/entry_points.txt") for n in names)
requirements: Final = tomllib.loads((ROOT / "pyproject.toml").read_text())["project"]["dependencies"]
requirements: Final = tomllib.loads((ROOT / "packaging/litellm-core/pyproject.toml").read_text())["project"]["dependencies"]
for python_version in ("3.10", "3.11", "3.12", "3.13", "3.14"):
environment: Final = {"python_version": python_version, "python_full_version": python_version + ".0"}
assert {

3
tests/unit/integrations/focus/test_s3_destination.py
View file

@ -6,6 +6,7 @@ from datetime import datetime, timezone
from types import SimpleNamespace
from typing import Any, Dict
import boto3
import pytest
import litellm.integrations.focus.destinations.s3_destination as s3_module
@ -81,7 +82,7 @@ def test_should_upload_with_configured_client(monkeypatch: pytest.MonkeyPatch):
return SimpleNamespace(put_object=put_object)
monkeypatch.setattr(s3_module.boto3, "client", fake_client)
monkeypatch.setattr(boto3, "client", fake_client)
dest._upload(content=b"payload", object_key="path/file.bin")

25
tests/unit/litellm_core_utils/test_exception_mapping_utils.py
View file

@ -1624,6 +1624,16 @@ def test_guardrail_block_raised_inside_an_llm_call_is_returned_unmapped(block: E
assert returned is block
@pytest.mark.parametrize("provider", ["bedrock", "bedrock_mantle"])
@pytest.mark.parametrize(
"failure", [ImportError("Run 'pip install boto3'."), ModuleNotFoundError(name="unrelated_dependency")]
)
def test_bedrock_import_errors_preserve_the_original_exception(provider, failure):
assert exception_type(
model="test-model", original_exception=failure, custom_llm_provider=provider
) is failure
def test_guardrail_provider_failure_status_is_still_mapped():
upstream_failure = HTTPException(status_code=401, detail={"error": "guardrail provider rejected the key"})
@ -2374,3 +2384,18 @@ def _pre_call_utils_httpx(
original_function = litellm.atext_completion
return data, original_function, mapped_target
@pytest.mark.parametrize("provider", ["sagemaker", "sagemaker_chat", "aws_polly", "openai"])
@pytest.mark.parametrize("dependency", ["boto3", "botocore"])
def test_missing_aws_dependency_is_not_mapped_to_provider_failure(provider, dependency):
failure = ModuleNotFoundError(f"No module named '{dependency}'", name=dependency)
assert exception_type(
model="test-model", original_exception=failure, custom_llm_provider=provider
) is failure
@pytest.mark.parametrize("failure", [ImportError("broken import"), ModuleNotFoundError(name="unrelated_dependency")])
def test_non_aws_import_failure_keeps_provider_mapping(failure):
with pytest.raises(litellm.APIConnectionError):
exception_type(model="test-model", original_exception=failure, custom_llm_provider="openai")

39
tests/unit/litellm_core_utils/test_optional_imports.py Normal file
View file

@ -0,0 +1,39 @@
import builtins
from typing import Final
from unittest.mock import patch
import pytest
from litellm.litellm_core_utils.optional_imports import ensure_optional_import
@pytest.mark.parametrize("module", ["boto3", "botocore", "tokenizers"])
def test_missing_optional_dependency_names_the_installable_package(module: str) -> None:
with patch.dict("sys.modules", {module: None}):
with pytest.raises(ModuleNotFoundError) as caught:
ensure_optional_import(module)
package: Final = "boto3" if module == "botocore" else module
assert str(caught.value) == f"Missing optional dependency '{module}'. Run 'pip install {package}'."
assert caught.value.name == module
assert isinstance(caught.value.__cause__, ModuleNotFoundError)
@pytest.mark.parametrize(
"failure", [ModuleNotFoundError(name="unrelated_dependency"), ImportError("broken installation")]
)
def test_optional_import_preserves_unrelated_failure(failure: ImportError) -> None:
original_import: Final = builtins.__import__
def import_dependency(name, *args, **kwargs):
if name == "botocore":
raise failure
return original_import(name, *args, **kwargs)
with patch("builtins.__import__", side_effect=import_dependency):
with pytest.raises(ImportError) as caught:
ensure_optional_import("botocore")
assert caught.value is failure
def test_available_optional_dependency_does_not_raise() -> None:
assert ensure_optional_import("json") is None

135
tests/unit/litellm_core_utils/test_tokenizer.py
View file

@ -401,3 +401,138 @@ def test_huggingface_encoding_exposes_the_tokenizers_lookup_and_mutation_surface
assert merged.offsets == type(expected).merge([expected, reference.encode("more")]).offsets
with pytest.raises(ValueError, match="direction"):
actual.pad(8, direction="sideways")
@pytest.mark.parametrize("log_level", ["WARNING", "ERROR"])
def test_missing_python_tokenizer_warns_before_approximate_count(caplog, monkeypatch, log_level):
from unittest.mock import patch
from litellm.utils import _load_huggingface_tokenizer, _select_tokenizer_helper
monkeypatch.setenv("LITELLM_RUST", "false")
_load_huggingface_tokenizer.cache_clear()
with caplog.at_level(log_level, logger="LiteLLM"), patch.dict(sys.modules, {"tokenizers": None}):
result = _select_tokenizer_helper("llama-2")
assert result["type"] == "openai_tokenizer"
assert result["tokenizer"].encode("hello")
assert ("token counts may be approximate" in caplog.text) is (log_level == "WARNING")
assert ("install tokenizers" in caplog.text) is (log_level == "WARNING")
@pytest.mark.parametrize("python_installed", [False, True])
def test_runtime_aliases_accept_available_tokenizer_instances(python_installed):
from contextlib import nullcontext
from unittest.mock import patch
from litellm.litellm_core_utils import tokenizer as types
native = HuggingFaceTokenizer.from_str(TOKENIZER_JSON)
python = ReferenceTokenizer.from_str(TOKENIZER_JSON)
with nullcontext() if python_installed else patch.dict(sys.modules, {"tokenizers": None}):
assert isinstance(native, types.HuggingFace)
assert isinstance(native, types.Tokenizer)
assert isinstance(tiktoken.get_encoding("cl100k_base"), types.Tokenizer)
assert isinstance(OpenAIEncoding.from_tiktoken("cl100k_base"), types.Tokenizer)
assert isinstance(python, types.HuggingFace) is python_installed
assert isinstance(python, types.Tokenizer) is python_installed
def test_runtime_alias_does_not_hide_broken_tokenizer_installation():
from unittest.mock import patch
from litellm.litellm_core_utils import tokenizer as types
failure = ModuleNotFoundError("broken installation", name="tokenizer_dependency")
with patch("builtins.__import__", side_effect=failure):
with pytest.raises(ModuleNotFoundError) as error:
getattr(types, "Tokenizer")
assert error.value is failure
def test_unknown_tokenizer_export_raises_attribute_error():
from litellm.litellm_core_utils import tokenizer as types
with pytest.raises(AttributeError, match="unknown_tokenizer"):
getattr(types, "unknown_tokenizer")
@pytest.mark.parametrize("python_installed", [False, True])
def test_added_token_return_annotation_resolves_without_optional_import(python_installed):
from contextlib import nullcontext
from typing import get_type_hints
from unittest.mock import patch
with nullcontext() if python_installed else patch.dict(sys.modules, {"tokenizers": None}):
hints = get_type_hints(HuggingFaceTokenizer.get_added_tokens_decoder)
assert "return" in hints
decoder = HuggingFaceTokenizer.from_str(TOKENIZER_JSON).get_added_tokens_decoder()
for token in decoder.values():
assert isinstance(token.content, str)
assert isinstance(token.special, bool)
def test_tokenizer_fallback_logs_safe_diagnostic_context(caplog, monkeypatch):
from types import SimpleNamespace
from unittest.mock import patch
from litellm.utils import _load_huggingface_tokenizer, _select_tokenizer_helper
monkeypatch.setenv("LITELLM_RUST", "false")
_load_huggingface_tokenizer.cache_clear()
model = "llama-2\r\nforged-model\x1b[31m\u2028\u2029"
secret = "sk-" + "x" * 48
failure = OSError("download failed\r\nforged-error\x1b[31m api_key=" + secret)
def fail_download(*args, **kwargs):
raise failure
dependency = SimpleNamespace(Tokenizer=SimpleNamespace(from_pretrained=fail_download))
with caplog.at_level("WARNING", logger="LiteLLM"), patch.dict(sys.modules, {"tokenizers": dependency}):
result = _select_tokenizer_helper(model)
assert result["type"] == "openai_tokenizer"
assert result["tokenizer"].encode("hello")
message = next(record.getMessage() for record in caplog.records if "token counts may be approximate" in record.getMessage())
assert "llama-2" in message
assert "download failed" in message
assert "forged-model" in message and "forged-error" in message
assert message.isascii() and message.isprintable()
assert secret not in message
assert "REDACTED" in message
assert "install tokenizers and huggingface-hub" in message
@pytest.mark.parametrize("python_installed", [False, True])
def test_native_added_token_decoder_preserves_fields_without_python_dependency(monkeypatch, python_installed):
from tokenizers import AddedToken
native = HuggingFaceTokenizer.from_str(TOKENIZER_JSON)
expected = ReferenceTokenizer.from_str(TOKENIZER_JSON).get_added_tokens_decoder()
if not python_installed:
monkeypatch.setitem(sys.modules, "tokenizers", None)
actual = native.get_added_tokens_decoder()
attributes = ("content", "single_word", "lstrip", "rstrip", "normalized", "special")
assert {
token_id: tuple(getattr(token, name) for name in attributes) for token_id, token in actual.items()
} == {
token_id: tuple(getattr(token, name) for name in attributes) for token_id, token in expected.items()
}
assert {token_id: str(token) for token_id, token in actual.items()} == {
token_id: str(token) for token_id, token in expected.items()
}
if python_installed:
assert all(isinstance(token, AddedToken) for token in actual.values())
def test_native_added_token_decoder_preserves_unrelated_import_failure():
import builtins
from unittest.mock import patch
native = HuggingFaceTokenizer.from_str(TOKENIZER_JSON)
failure = ModuleNotFoundError(name="broken_tokenizer_dependency")
original_import = builtins.__import__
def import_dependency(name, *args, **kwargs):
if name == "tokenizers":
raise failure
return original_import(name, *args, **kwargs)
with patch("builtins.__import__", side_effect=import_dependency):
with pytest.raises(ModuleNotFoundError) as caught:
native.get_added_tokens_decoder()
assert caught.value is failure

19
tests/unit/llms/aws_polly/text_to_speech/test_transformation.py Normal file
View file

@ -0,0 +1,19 @@
import json
from typing import Final
from litellm.llms.aws_polly.text_to_speech.transformation import AWSPollyTextToSpeechConfig
def test_installed_botocore_signs_the_speech_request() -> None:
headers, body = AWSPollyTextToSpeechConfig()._sign_polly_request(
request_body={"Text": "ping", "VoiceId": "Joanna"},
endpoint_url="https://polly.us-west-2.amazonaws.com/v1/speech",
litellm_params={
"aws_access_key_id": "test-key",
"aws_secret_access_key": "test-secret",
"aws_region_name": "us-west-2",
},
)
authorization: Final = headers["Authorization"]
assert "/us-west-2/polly/aws4_request" in authorization
assert json.loads(body) == {"Text": "ping", "VoiceId": "Joanna"}

16
tests/unit/llms/bedrock/chat/test_invoke_handler.py
View file

@ -1426,3 +1426,19 @@ def test_filter_headers_for_aws_signature():
non_aws_headers = {"x-custom-trace": "trace-123", "x-user-context": "premium", "x-request-source": "mobile-app"}
filtered_non_aws = aws_llm._filter_headers_for_aws_signature(non_aws_headers)
assert filtered_non_aws == {}
@pytest.mark.parametrize("missing", ["botocore", "unrelated_dependency"])
def test_invoke_decoder_reports_only_missing_aws_dependency(missing):
from unittest.mock import patch
from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder
failure = ModuleNotFoundError("missing dependency", name=missing)
with patch("builtins.__import__", side_effect=failure):
with pytest.raises(ImportError) as error:
AWSEventStreamDecoder(model="anthropic.claude-3-sonnet-20240229-v1:0")
if missing == "botocore":
assert "pip install boto3" in str(error.value)
assert error.value.__cause__ is failure
else:
assert error.value is failure

182
tests/unit/llms/bedrock/test_base_aws_llm.py
View file

@ -745,40 +745,23 @@ def test_sign_request_with_api_key_bearer_token():
assert result_body == json.dumps(request_data).encode()
def test_get_request_headers_with_env_var_bearer_token():
# Setup
llm = BaseAWSLLM()
credentials = Credentials("test_key", "test_secret", "test_token")
headers = {"Content-Type": "application/json"}
headers_dict = headers.copy()
# Create mock request
mock_prepared_request = MagicMock(spec=AWSPreparedRequest)
mock_request = MagicMock(spec=AWSRequest)
mock_request.headers = headers_dict
mock_request.prepare.return_value = mock_prepared_request
def mock_aws_request_init(method, url, data, headers):
mock_request.headers.update(headers)
return mock_request
# Test with bearer token
with (
patch.dict(os.environ, {"AWS_BEARER_TOKEN_BEDROCK": "test_token"}),
patch("botocore.awsrequest.AWSRequest", side_effect=mock_aws_request_init),
):
result = llm.get_request_headers(
credentials=credentials,
@pytest.mark.parametrize("from_environment", [True, False])
def test_get_request_headers_preserves_bearer_payload(from_environment):
with patch.dict(os.environ, {"AWS_BEARER_TOKEN_BEDROCK": "test_token"} if from_environment else {}, clear=True):
result = BaseAWSLLM().get_request_headers(
credentials=None,
aws_region_name="us-west-2",
extra_headers=None,
endpoint_url="https://api.example.com",
data='{"prompt": "test"}',
headers=headers_dict,
headers={"Content-Type": "application/json"},
api_key=None if from_environment else "test_token",
)
# Assert
assert mock_request.headers["Authorization"] == "Bearer test_token"
assert result == mock_prepared_request
assert result.headers["Authorization"] == "Bearer test_token"
assert result.headers["Content-Type"] == "application/json"
assert result.body == b'{"prompt": "test"}'
assert result.url == "https://api.example.com"
assert result.method == "POST"
def test_get_request_headers_with_sigv4():
@ -856,46 +839,6 @@ def test_sigv4_matches_rust_golden_vector():
)
def test_get_request_headers_with_api_key_bearer_token():
"""
Test that get_request_headers uses the api_key parameter as a bearer token when provided
"""
# Setup
llm = BaseAWSLLM()
credentials = Credentials("test_key", "test_secret", "test_token")
headers = {"Content-Type": "application/json"}
headers_dict = headers.copy()
api_key = "test_api_key"
# Create mock request
mock_prepared_request = MagicMock(spec=AWSPreparedRequest)
mock_request = MagicMock(spec=AWSRequest)
mock_request.headers = headers_dict
mock_request.prepare.return_value = mock_prepared_request
def mock_aws_request_init(method, url, data, headers):
mock_request.headers.update(headers)
return mock_request
# Test with api_key parameter
with (
patch.dict(os.environ, {}, clear=True),
patch("botocore.awsrequest.AWSRequest", side_effect=mock_aws_request_init),
):
result = llm.get_request_headers(
credentials=credentials,
aws_region_name="us-west-2",
extra_headers=None,
endpoint_url="https://api.example.com",
data='{"prompt": "test"}',
headers=headers_dict,
api_key=api_key,
)
# Assert
assert mock_request.headers["Authorization"] == f"Bearer {api_key}"
assert result == mock_prepared_request
def test_role_assumption_without_session_name():
"""
@ -4171,3 +4114,104 @@ def test_dynamic_aws_params_propagation(model, param_name, param_value, expected
# We now assert that get_credentials() was called with the dynamic param.
assert dummy_get_credentials.called_kwargs.get(param_name) == expected_credentials_value
def test_bearer_request_preparation_does_not_require_botocore():
import httpx
with patch.dict("sys.modules", {"botocore.credentials": None, "botocore.awsrequest": None}):
target = BaseAWSLLM()._get_boto_credentials_from_optional_params(
{"aws_region_name": "us-west-2"}, bearer_token="test-token"
)
request = BaseAWSLLM().get_request_headers(
credentials=None,
aws_region_name=target.aws_region_name,
extra_headers=None,
endpoint_url="https://bedrock-runtime.us-west-2.amazonaws.com/model/test/invoke",
data='{"text":"café"}',
headers={"Content-Type": "application/json"},
api_key="test-token",
)
assert dict(request.headers)["Authorization"] == "Bearer test-token"
assert request.body == '{"text":"café"}'.encode()
assert int(httpx.Headers(request.headers)["Content-Length"]) == len(request.body)
@pytest.mark.parametrize("missing", ["botocore", "unrelated_dependency"])
def test_signing_preserves_unrelated_import_failure(missing):
llm = BaseAWSLLM()
failure = ModuleNotFoundError("missing dependency", name=missing)
with patch("builtins.__import__", side_effect=failure):
with pytest.raises(ImportError) as error:
llm.get_request_headers(
credentials=Credentials("key", "secret"), aws_region_name="us-east-1",
extra_headers=None, endpoint_url="https://bedrock-runtime.us-east-1.amazonaws.com",
data="{}", headers={}, supports_bearer_token=False,
)
if missing == "botocore":
assert "pip install boto3" in str(error.value)
assert error.value.__cause__ is failure
else:
assert error.value is failure
@pytest.mark.parametrize("missing", ["botocore", "unrelated_dependency"])
@pytest.mark.parametrize("shared_signer", [False, True])
def test_json_signers_report_only_missing_aws_dependency(missing, shared_signer):
from litellm.llms.bedrock.base_aws_llm import sign_aws_json_post
llm = BaseAWSLLM()
failure = ModuleNotFoundError("missing dependency", name=missing)
from functools import partial
sign = (
partial(sign_aws_json_post, lambda: Credentials("key", "secret"), "bedrock", "us-east-1",
"https://bedrock-runtime.us-east-1.amazonaws.com", "{}", {})
if shared_signer else
partial(llm._sign_request, service_name="bedrock", headers={}, optional_params={}, request_data={},
api_base="https://bedrock-runtime.us-east-1.amazonaws.com", api_key="")
)
with patch.dict(os.environ, {}, clear=True), patch("builtins.__import__", side_effect=failure):
with pytest.raises(ImportError) as error:
sign()
if missing == "botocore":
assert "pip install boto3" in str(error.value)
assert error.value.__cause__ is failure
else:
assert error.value is failure
@pytest.mark.parametrize("missing", ["botocore", "unrelated_dependency"])
def test_direct_credentials_report_missing_aws_without_masking_other_imports(missing):
import builtins
original_import = builtins.__import__
failure = ModuleNotFoundError("dependency unavailable", name=missing)
def import_dependency(name, *args, **kwargs):
if name.startswith("botocore"):
raise failure
return original_import(name, *args, **kwargs)
with patch("builtins.__import__", side_effect=import_dependency):
with pytest.raises(ImportError) as error:
BaseAWSLLM().get_credentials(
aws_access_key_id="test-key", aws_secret_access_key="test-secret", aws_session_token="test-session"
)
if missing == "botocore":
assert "pip install boto3" in str(error.value)
assert error.value.__cause__ is failure
else:
assert error.value is failure
def test_shared_json_signer_preserves_body_and_signs_for_the_requested_service():
from litellm.llms.bedrock.base_aws_llm import sign_aws_json_post
body = '{"message":"ping"}'
request = sign_aws_json_post(
lambda: Credentials("test-key", "test-secret"), "s3", "us-west-2",
"https://s3.us-west-2.amazonaws.com", body, {"Content-Type": "application/json"},
)
assert request.body == body
assert "/us-west-2/s3/aws4_request" in request.headers["Authorization"]

16
tests/unit/llms/bedrock/test_bedrock_common_utils.py
View file

@ -1989,3 +1989,19 @@ class TestBedrockGovCloudSupport:
"""Test that GovCloud Titan models use Invoke API"""
route = BedrockModelInfo.get_bedrock_route(model_name)
assert route == "invoke"
@pytest.mark.parametrize("missing", ["botocore", "unrelated_dependency"])
def test_event_decoder_reports_only_missing_aws_dependency(missing):
from unittest.mock import patch
from litellm.llms.bedrock.common_utils import BedrockEventStreamDecoderBase
failure = ModuleNotFoundError("missing dependency", name=missing)
with patch("builtins.__import__", side_effect=failure):
with pytest.raises(ImportError) as error:
BedrockEventStreamDecoderBase()
if missing == "botocore":
assert "pip install boto3" in str(error.value)
assert error.value.__cause__ is failure
else:
assert error.value is failure

15
tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py
View file

@ -1042,3 +1042,18 @@ async def test_mantle_signing_runs_off_the_event_loop():
assert "Authorization" in signed
assert probe.served_during_refresh is True
@pytest.mark.parametrize("missing", ["botocore", "unrelated_dependency"])
def test_mantle_signing_reports_only_missing_aws_dependency(missing):
config = BedrockMantleChatConfig()
failure = ModuleNotFoundError("missing dependency", name=missing)
with patch.dict("os.environ", {}, clear=True), patch("builtins.__import__", side_effect=failure):
with pytest.raises(ImportError) as error:
config.sign_request(headers={}, optional_params={}, request_data={},
api_base="https://bedrock-mantle.us-east-1.api.aws/v1/chat/completions", api_key="")
if missing == "botocore":
assert "pip install boto3" in str(error.value)
assert error.value.__cause__ is failure
else:
assert error.value is failure

20
tests/unit/llms/sagemaker/test_sagemaker_chat_handler.py
View file

@ -163,3 +163,23 @@ async def test_completion_sagemaker_messages_api(sync_mode):
assert json_data["max_tokens"] == 80
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_missing_botocore_keeps_dependency_identity():
import pytest
with patch.dict("sys.modules", {"botocore": None}):
with pytest.raises(ModuleNotFoundError, match="pip install boto3") as caught:
SagemakerChatHandler()._load_credentials({})
assert caught.value.name == "botocore"
def test_installed_botocore_signs_the_chat_request():
from botocore.credentials import Credentials
request = SagemakerChatHandler()._prepare_request(
credentials=Credentials("test-key", "test-secret"), model="test-endpoint", data={"inputs": "ping"},
optional_params={}, aws_region_name="us-west-2",
)
assert request.body == b'{"inputs": "ping"}'
assert "/us-west-2/sagemaker/aws4_request" in request.headers["Authorization"]

26
tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py
View file

@ -267,3 +267,29 @@ def test_load_credentials_assumes_role_with_session_tags(monkeypatch):
assert credentials.access_key == "ASIASMCOMPTAGGED"
assert aws_region_name == "us-east-1"
assert "aws_session_tags" not in optional_params
def test_missing_botocore_keeps_dependency_identity():
from unittest.mock import patch
import pytest
from litellm.llms.sagemaker.completion.handler import SagemakerLLM
with patch.dict("sys.modules", {"botocore": None}):
with pytest.raises(ModuleNotFoundError, match="pip install boto3") as caught:
SagemakerLLM()._load_credentials({})
assert caught.value.name == "botocore"
def test_installed_botocore_signs_the_completion_request():
from botocore.credentials import Credentials
from litellm.llms.sagemaker.completion.handler import SagemakerLLM
request = SagemakerLLM()._prepare_request(
credentials=Credentials("test-key", "test-secret"), model="test-endpoint", data={"inputs": "ping"},
messages=[], litellm_params={}, optional_params={}, aws_region_name="us-west-2",
)
assert request.body == b'{"inputs": "ping"}'
assert "/us-west-2/sagemaker/aws4_request" in request.headers["Authorization"]

22
tests/unit/rust_bridge/test_tokenizer.py
View file

@ -1,4 +1,5 @@
from typing import Final
from unittest.mock import patch
import pytest
import tiktoken
@ -47,3 +48,24 @@ def test_native_custom_tokenizer_matches_python() -> None:
assert native.encode("Hello World").ids == reference.encode("Hello World").ids
assert native.decode(reference.encode("Hello World").ids) == reference.decode(reference.encode("Hello World").ids)
def test_python_tokenizer_missing_dependency_is_actionable() -> None:
with patch.dict("sys.modules", {"tokenizers": None}):
with pytest.raises(ImportError, match="pip install tokenizers") as error:
tokenizer._python_tokenizer()
assert isinstance(error.value.__cause__, ModuleNotFoundError)
assert error.value.__cause__.name == "tokenizers"
def test_python_tokenizer_factory_preserves_installed_interface() -> None:
result: Final = tokenizer._python_tokenizer().from_str(TOKENIZER_JSON)
assert result.encode("Hello World").ids == Tokenizer.from_str(TOKENIZER_JSON).encode("Hello World").ids
def test_python_tokenizer_preserves_unrelated_import_failure() -> None:
failure: Final = ModuleNotFoundError("broken tokenizer installation", name="unrelated_dependency")
with patch("builtins.__import__", side_effect=failure):
with pytest.raises(ModuleNotFoundError) as error:
tokenizer._python_tokenizer()
assert error.value is failure

13
tests/unit/secret_managers/test_aws_secret_manager_v2.py
View file

@ -664,3 +664,16 @@ async def test_end_to_end_iam_role_secret_write():
print("Delete Response:", delete_response)
except Exception as e:
print(f"Cleanup failed: {e}")
def test_missing_botocore_keeps_dependency_identity():
from unittest.mock import patch
import pytest
from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2
with patch.dict("sys.modules", {"botocore": None}):
with pytest.raises(ModuleNotFoundError, match="pip install boto3") as caught:
AWSSecretsManagerV2()._prepare_request(action="GetSecretValue", secret_name="test-secret")
assert caught.value.name == "botocore"

126
tests/unit/test_main.py
View file

@ -4630,6 +4630,45 @@ def test_completion_rejects_an_invalid_stream_chunk_size_before_the_mcp_gateway(
assert exc_info.value.param == "stream_chunk_size"
@pytest.mark.asyncio
@pytest.mark.parametrize("use_async", [False, True])
@pytest.mark.parametrize("missing_tenacity", [False, True])
@pytest.mark.parametrize("route", ["bedrock", "bedrock/invoke"])
async def test_bedrock_stream_missing_dependency_remains_actionable_with_retries(
monkeypatch, use_async, missing_tenacity, route
):
import builtins
original_import = builtins.__import__
def import_without_aws_or_retry_dependencies(name, *args, **kwargs):
if name.split(".")[0] == "botocore" or (name == "tenacity" and missing_tenacity):
raise ModuleNotFoundError(name=name.split(".")[0])
return original_import(name, *args, **kwargs)
monkeypatch.setattr(builtins, "__import__", import_without_aws_or_retry_dependencies)
monkeypatch.setattr(litellm, "num_retries", None)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
with respx.mock as upstream:
response = upstream.post(url__regex=r"https://bedrock-test\.invalid/.*").respond(200, content=b"")
arguments = dict(
model=f"{route}/anthropic.claude-3-sonnet-20240229-v1:0",
messages=[{"role": "user", "content": "ping"}],
api_key="test-bearer",
aws_region_name="us-east-1",
aws_bedrock_runtime_endpoint="https://bedrock-test.invalid",
stream=True,
num_retries=1,
)
if use_async:
with pytest.raises(ImportError, match="pip install boto3"):
await litellm.acompletion(**arguments)
else:
with pytest.raises(ImportError, match="pip install boto3"):
litellm.completion(**arguments)
assert response.call_count == 1
def test_drop_params_false_still_rejects_an_invalid_stream_chunk_size() -> None:
with pytest.raises(litellm.BadRequestError):
litellm.completion(
@ -5978,3 +6017,90 @@ async def test_transcription_model_names_pass_through(
assert response.text == "hello"
assert route.call_count == 1
assert f'name="model"\r\n\r\n{model}\r\n'.encode() in route.calls[0].request.content
@pytest.mark.asyncio
@pytest.mark.parametrize("use_async", [False, True])
@pytest.mark.parametrize("model", ["sagemaker/test-endpoint", "sagemaker_chat/test-endpoint"])
async def test_sagemaker_missing_dependency_remains_actionable_with_retries(monkeypatch, use_async, model):
import sys
monkeypatch.setattr(litellm, "num_retries", None)
for dependency in ("botocore", "boto3", "tenacity"):
monkeypatch.setitem(sys.modules, dependency, None)
if use_async:
with pytest.raises(ModuleNotFoundError, match="pip install boto3") as caught:
await litellm.acompletion(model=model, messages=[{"role": "user", "content": "ping"}], num_retries=1)
else:
with pytest.raises(ModuleNotFoundError, match="pip install boto3") as caught:
litellm.completion(model=model, messages=[{"role": "user", "content": "ping"}], num_retries=1)
assert caught.value.name == "botocore"
@pytest.mark.asyncio
@pytest.mark.parametrize("use_async", [False, True])
async def test_polly_missing_dependency_remains_actionable_with_retries(monkeypatch, use_async):
import sys
monkeypatch.setattr(litellm, "num_retries", None)
for dependency in ("botocore", "boto3", "tenacity"):
monkeypatch.setitem(sys.modules, dependency, None)
if use_async:
with pytest.raises(ModuleNotFoundError, match="pip install boto3") as caught:
await litellm.aspeech(model="aws_polly/standard", input="ping", voice="Joanna", num_retries=1)
else:
with pytest.raises(ModuleNotFoundError, match="pip install boto3") as caught:
litellm.speech(model="aws_polly/standard", input="ping", voice="Joanna", num_retries=1)
assert caught.value.name == "botocore"
@pytest.mark.asyncio
@pytest.mark.parametrize("missing_tenacity", [False, True])
@pytest.mark.parametrize("use_async", [False, True])
async def test_mantle_responses_missing_dependency_is_not_retried(monkeypatch, missing_tenacity, use_async):
import builtins
original_import = builtins.__import__
attempts = []
def import_without_aws(name, *args, **kwargs):
if name == "botocore":
attempts.append(name)
raise ModuleNotFoundError(name="botocore")
if name == "tenacity" and missing_tenacity:
attempts.append(name)
raise ModuleNotFoundError(name="tenacity")
return original_import(name, *args, **kwargs)
monkeypatch.setattr(builtins, "__import__", import_without_aws)
monkeypatch.setattr(litellm, "num_retries", None)
for name in ("AWS_BEARER_TOKEN_BEDROCK", "BEDROCK_MANTLE_API_KEY"):
monkeypatch.delenv(name, raising=False)
if use_async:
with pytest.raises(ModuleNotFoundError, match="pip install boto3"):
await litellm.aresponses(model="bedrock_mantle/openai.gpt-oss-120b", input="ping", num_retries=1)
else:
with pytest.raises(ModuleNotFoundError, match="pip install boto3"):
litellm.responses(model="bedrock_mantle/openai.gpt-oss-120b", input="ping", num_retries=1)
assert attempts == ["botocore"]
@pytest.mark.asyncio
async def test_async_responses_still_retries_provider_server_errors(monkeypatch):
monkeypatch.setattr(litellm, "num_retries", None)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
with respx.mock as upstream:
response = upstream.post("https://openai-test.invalid/v1/responses").mock(side_effect=[
httpx.Response(500, json={"error": {"message": "temporary provider failure", "type": "server_error"}}),
httpx.Response(200, json={
"id": "resp-retry", "object": "response", "created_at": 1, "status": "completed",
"model": "test-model", "output": [],
"usage": {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5},
}),
])
result = await litellm.aresponses(
model="openai/test-model", input="ping", api_key="test-key",
api_base="https://openai-test.invalid/v1", num_retries=1, max_retries=0,
)
assert result.status == "completed"
assert response.call_count == 2