From 911aff27524282b1ced52a9b50d33a3127e80052 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Thu, 8 Oct 2026 19:07:57 -0700 Subject: [PATCH] refactor(sdk): separate core AWS and tokenizer dependencies (#44447) * 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> --- .../test-mcp-dependency-resolution.yml | 9 +- .../focus/destinations/s3_destination.py | 3 +- .../exception_mapping_utils.py | 8 +- .../litellm_core_utils/optional_imports.py | 13 ++ litellm/litellm_core_utils/token_counter.py | 11 +- litellm/litellm_core_utils/tokenizer.py | 60 +++++- .../text_to_speech/transformation.py | 9 +- litellm/llms/bedrock/base_aws_llm.py | 62 +++--- litellm/llms/bedrock/chat/invoke_handler.py | 6 +- litellm/llms/bedrock/common_utils.py | 10 +- litellm/llms/bedrock/image_edit/handler.py | 3 +- .../bedrock/image_generation/image_handler.py | 3 +- litellm/llms/bedrock_mantle/common_utils.py | 68 +++--- litellm/llms/sagemaker/chat/handler.py | 16 +- litellm/llms/sagemaker/completion/handler.py | 16 +- litellm/rust_bridge/tokenizer.py | 22 +- .../secret_managers/aws_secret_manager_v2.py | 10 +- litellm/types/llms/bedrock.py | 13 +- litellm/utils.py | 19 +- packaging/litellm-core/pyproject.toml | 4 +- .../base_sdk_tests/check_base_sdk_install.py | 48 ++++- tests/base_sdk_tests/check_sdk_http.py | 202 ++++++++++++++++++ .../base_sdk_tests/test_core_distribution.py | 15 +- .../integrations/focus/test_s3_destination.py | 3 +- .../test_exception_mapping_utils.py | 25 +++ .../test_optional_imports.py | 39 ++++ .../unit/litellm_core_utils/test_tokenizer.py | 135 ++++++++++++ .../text_to_speech/test_transformation.py | 19 ++ .../llms/bedrock/chat/test_invoke_handler.py | 16 ++ tests/unit/llms/bedrock/test_base_aws_llm.py | 182 ++++++++++------ .../llms/bedrock/test_bedrock_common_utils.py | 16 ++ .../test_bedrock_mantle_transformation.py | 15 ++ .../sagemaker/test_sagemaker_chat_handler.py | 20 ++ .../test_sagemaker_completion_handler.py | 26 +++ tests/unit/rust_bridge/test_tokenizer.py | 22 ++ .../test_aws_secret_manager_v2.py | 13 ++ tests/unit/test_main.py | 126 +++++++++++ 37 files changed, 1085 insertions(+), 202 deletions(-) create mode 100644 litellm/litellm_core_utils/optional_imports.py create mode 100644 tests/base_sdk_tests/check_sdk_http.py create mode 100644 tests/unit/litellm_core_utils/test_optional_imports.py create mode 100644 tests/unit/llms/aws_polly/text_to_speech/test_transformation.py diff --git a/.github/workflows/test-mcp-dependency-resolution.yml b/.github/workflows/test-mcp-dependency-resolution.yml index 7c88a8a579d..3f0d1260e49 100644 --- a/.github/workflows/test-mcp-dependency-resolution.yml +++ b/.github/workflows/test-mcp-dependency-resolution.yml @@ -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 diff --git a/litellm/integrations/focus/destinations/s3_destination.py b/litellm/integrations/focus/destinations/s3_destination.py index 661cf1933ff..0535523dd99 100644 --- a/litellm/integrations/focus/destinations/s3_destination.py +++ b/litellm/integrations/focus/destinations/s3_destination.py @@ -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, diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index b5afa962cdc..0ba466ec324 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -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__ diff --git a/litellm/litellm_core_utils/optional_imports.py b/litellm/litellm_core_utils/optional_imports.py new file mode 100644 index 00000000000..7b8bbb85d42 --- /dev/null +++ b/litellm/litellm_core_utils/optional_imports.py @@ -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 diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 4286d4d6f6a..151830f1589 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -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)) diff --git a/litellm/litellm_core_utils/tokenizer.py b/litellm/litellm_core_utils/tokenizer.py index 31cd8f63116..3037045c277 100644 --- a/litellm/litellm_core_utils/tokenizer.py +++ b/litellm/litellm_core_utils/tokenizer.py @@ -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: ... diff --git a/litellm/llms/aws_polly/text_to_speech/transformation.py b/litellm/llms/aws_polly/text_to_speech/transformation.py index 133e40dc1ab..01fec02c310 100644 --- a/litellm/llms/aws_polly/text_to_speech/transformation.py +++ b/litellm/llms/aws_polly/text_to_speech/transformation.py @@ -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) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 424e1dc2106..91d18c651da 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -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) diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 6276324ad7e..4682a47de6b 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -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 diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index c21cb61ce60..a3f80001868 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -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( diff --git a/litellm/llms/bedrock/image_edit/handler.py b/litellm/llms/bedrock/image_edit/handler.py index 76ec537b68d..53bb332f511 100644 --- a/litellm/llms/bedrock/image_edit/handler.py +++ b/litellm/llms/bedrock/image_edit/handler.py @@ -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 diff --git a/litellm/llms/bedrock/image_generation/image_handler.py b/litellm/llms/bedrock/image_generation/image_handler.py index 018e92fd79f..e81dfe76942 100644 --- a/litellm/llms/bedrock/image_generation/image_handler.py +++ b/litellm/llms/bedrock/image_generation/image_handler.py @@ -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 diff --git a/litellm/llms/bedrock_mantle/common_utils.py b/litellm/llms/bedrock_mantle/common_utils.py index 8f9cf5c4f6b..75e34cf3ae5 100644 --- a/litellm/llms/bedrock_mantle/common_utils.py +++ b/litellm/llms/bedrock_mantle/common_utils.py @@ -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, diff --git a/litellm/llms/sagemaker/chat/handler.py b/litellm/llms/sagemaker/chat/handler.py index 10be9ef384c..3e671ce5d07 100644 --- a/litellm/llms/sagemaker/chat/handler.py +++ b/litellm/llms/sagemaker/chat/handler.py @@ -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) diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index 46d204cf385..069518021fa 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -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) diff --git a/litellm/rust_bridge/tokenizer.py b/litellm/rust_bridge/tokenizer.py index 5ed89813620..1ae7f136a43 100644 --- a/litellm/rust_bridge/tokenizer.py +++ b/litellm/rust_bridge/tokenizer.py @@ -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), ) diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index a702c206cfd..98d12e497da 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -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 diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index cc420fb9509..78d03ecc6e1 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -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 diff --git a/litellm/utils.py b/litellm/utils.py index 7a334f8974b..a11a33561ca 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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) diff --git a/packaging/litellm-core/pyproject.toml b/packaging/litellm-core/pyproject.toml index d7702a02b51..e3e2ca3d1db 100644 --- a/packaging/litellm-core/pyproject.toml +++ b/packaging/litellm-core/pyproject.toml @@ -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", ] diff --git a/tests/base_sdk_tests/check_base_sdk_install.py b/tests/base_sdk_tests/check_base_sdk_install.py index ad4e659c470..e56d48f5e6e 100644 --- a/tests/base_sdk_tests/check_base_sdk_install.py +++ b/tests/base_sdk_tests/check_base_sdk_install.py @@ -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: diff --git a/tests/base_sdk_tests/check_sdk_http.py b/tests/base_sdk_tests/check_sdk_http.py new file mode 100644 index 00000000000..d773e252127 --- /dev/null +++ b/tests/base_sdk_tests/check_sdk_http.py @@ -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() diff --git a/tests/base_sdk_tests/test_core_distribution.py b/tests/base_sdk_tests/test_core_distribution.py index 5bcb7120d7d..e16226e1df8 100644 --- a/tests/base_sdk_tests/test_core_distribution.py +++ b/tests/base_sdk_tests/test_core_distribution.py @@ -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 { diff --git a/tests/unit/integrations/focus/test_s3_destination.py b/tests/unit/integrations/focus/test_s3_destination.py index 8e54b561f82..12ddf412ac3 100644 --- a/tests/unit/integrations/focus/test_s3_destination.py +++ b/tests/unit/integrations/focus/test_s3_destination.py @@ -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") diff --git a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py index 7e7f1b536f8..1bae82014cf 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -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") diff --git a/tests/unit/litellm_core_utils/test_optional_imports.py b/tests/unit/litellm_core_utils/test_optional_imports.py new file mode 100644 index 00000000000..8b6cc539885 --- /dev/null +++ b/tests/unit/litellm_core_utils/test_optional_imports.py @@ -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 diff --git a/tests/unit/litellm_core_utils/test_tokenizer.py b/tests/unit/litellm_core_utils/test_tokenizer.py index 9d08442b164..354c91920fc 100644 --- a/tests/unit/litellm_core_utils/test_tokenizer.py +++ b/tests/unit/litellm_core_utils/test_tokenizer.py @@ -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 diff --git a/tests/unit/llms/aws_polly/text_to_speech/test_transformation.py b/tests/unit/llms/aws_polly/text_to_speech/test_transformation.py new file mode 100644 index 00000000000..4778b6296fc --- /dev/null +++ b/tests/unit/llms/aws_polly/text_to_speech/test_transformation.py @@ -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"} diff --git a/tests/unit/llms/bedrock/chat/test_invoke_handler.py b/tests/unit/llms/bedrock/chat/test_invoke_handler.py index ba35b93270a..d5fe5d8cf65 100644 --- a/tests/unit/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/unit/llms/bedrock/chat/test_invoke_handler.py @@ -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 diff --git a/tests/unit/llms/bedrock/test_base_aws_llm.py b/tests/unit/llms/bedrock/test_base_aws_llm.py index f1c95aaf7a2..cc625dcf3f3 100644 --- a/tests/unit/llms/bedrock/test_base_aws_llm.py +++ b/tests/unit/llms/bedrock/test_base_aws_llm.py @@ -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"] diff --git a/tests/unit/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index b3e31b36373..3a2c851807c 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -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 diff --git a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index 9c35e3253e5..ea9cfc1fd39 100644 --- a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -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 diff --git a/tests/unit/llms/sagemaker/test_sagemaker_chat_handler.py b/tests/unit/llms/sagemaker/test_sagemaker_chat_handler.py index 598d8177e6e..4ae117b5ebf 100644 --- a/tests/unit/llms/sagemaker/test_sagemaker_chat_handler.py +++ b/tests/unit/llms/sagemaker/test_sagemaker_chat_handler.py @@ -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"] diff --git a/tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py b/tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py index 9a5f36e2081..339e02c20d5 100644 --- a/tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py +++ b/tests/unit/llms/sagemaker/test_sagemaker_completion_handler.py @@ -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"] diff --git a/tests/unit/rust_bridge/test_tokenizer.py b/tests/unit/rust_bridge/test_tokenizer.py index c5093cdb0ce..42a31a911fb 100644 --- a/tests/unit/rust_bridge/test_tokenizer.py +++ b/tests/unit/rust_bridge/test_tokenizer.py @@ -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 diff --git a/tests/unit/secret_managers/test_aws_secret_manager_v2.py b/tests/unit/secret_managers/test_aws_secret_manager_v2.py index 1972782bc01..ab358e0bf8b 100644 --- a/tests/unit/secret_managers/test_aws_secret_manager_v2.py +++ b/tests/unit/secret_managers/test_aws_secret_manager_v2.py @@ -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" diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index dd7073bb381..c1b72bf8ff1 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -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