diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index 2c495aa8db..cef723ae26 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -26,6 +26,54 @@ def _get_default_encoding() -> Any: _default_encoding = encoding return _default_encoding + +# Lazy loader for get_modified_max_tokens to avoid importing token_counter at module import time +_get_modified_max_tokens_func: Optional[Any] = None + + +def _get_modified_max_tokens() -> Any: + """ + Lazily load and cache the get_modified_max_tokens function. + + This avoids importing `litellm.litellm_core_utils.token_counter` at `litellm` import time. + The function is cached after the first import. + + This is used internally by utils.py functions that need the token counter but shouldn't + trigger its import during module load. + """ + global _get_modified_max_tokens_func + if _get_modified_max_tokens_func is None: + from litellm.litellm_core_utils.token_counter import ( + get_modified_max_tokens as _get_modified_max_tokens_imported, + ) + + _get_modified_max_tokens_func = _get_modified_max_tokens_imported + return _get_modified_max_tokens_func + + +# Lazy loader for token_counter to avoid importing token_counter module at module import time +_token_counter_new_func: Optional[Any] = None + + +def _get_token_counter_new() -> Any: + """ + Lazily load and cache the token_counter function (aliased as token_counter_new). + + This avoids importing `litellm.litellm_core_utils.token_counter` at `litellm` import time. + The function is cached after the first import. + + This is used internally by utils.py functions that need the token counter but shouldn't + trigger its import during module load. + """ + global _token_counter_new_func + if _token_counter_new_func is None: + from litellm.litellm_core_utils.token_counter import ( + token_counter as _token_counter_imported, + ) + + _token_counter_new_func = _token_counter_imported + return _token_counter_new_func + # Cost calculator names that support lazy loading via _lazy_import_cost_calculator COST_CALCULATOR_NAMES = ( "completion_cost", diff --git a/litellm/utils.py b/litellm/utils.py index 169a2cdcb1..4b648b9eb5 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -103,7 +103,11 @@ from litellm.litellm_core_utils.dot_notation_indexing import ( delete_nested_value, is_nested_path, ) -from litellm._lazy_imports import _get_default_encoding +from litellm._lazy_imports import ( + _get_default_encoding, + _get_modified_max_tokens, + _get_token_counter_new, +) from litellm.litellm_core_utils.exception_mapping_utils import ( _get_response_headers, exception_type, @@ -147,7 +151,6 @@ from litellm.litellm_core_utils.redact_messages import ( ) from litellm.litellm_core_utils.rules import Rules from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper -from litellm.litellm_core_utils.token_counter import get_modified_max_tokens from litellm.llms.base_llm.google_genai.transformation import ( BaseGoogleGenAIGenerateContentConfig, ) @@ -252,7 +255,6 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( update_response_metadata, ) from litellm.litellm_core_utils.thread_pool_executor import executor -from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -268,6 +270,7 @@ if TYPE_CHECKING: # Heavy types that are only needed for type checking; avoid importing # their modules at runtime during `litellm` import. from litellm.llms.base_llm.files.transformation import BaseFilesConfig + from litellm.proxy._types import AllowedModelRegion from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.llms.base_llm.completion.transformation import BaseTextCompletionConfig @@ -318,7 +321,6 @@ from .exceptions import ( UnprocessableEntityError, UnsupportedParamsError, ) -from .proxy._types import AllowedModelRegion, KeyManagementSystem from .types.llms.openai import ( ChatCompletionDeltaToolCallChunk, ChatCompletionToolCallChunk, @@ -1252,7 +1254,7 @@ def client(original_function): # noqa: PLR0915 elif kwargs.get("messages", None): messages = kwargs["messages"] user_max_tokens = kwargs.get("max_tokens") - modified_max_tokens = get_modified_max_tokens( + modified_max_tokens = _get_modified_max_tokens()( model=model, base_model=base_model, messages=messages, @@ -1489,7 +1491,7 @@ def client(original_function): # noqa: PLR0915 elif kwargs.get("messages", None): messages = kwargs["messages"] user_max_tokens = kwargs.get("max_tokens") - modified_max_tokens = get_modified_max_tokens( + modified_max_tokens = _get_modified_max_tokens()( model=model, base_model=base_model, messages=messages, @@ -1880,7 +1882,7 @@ def token_counter( if litellm.disable_token_counter is True: return 0 - return token_counter_new( + return _get_token_counter_new()( model, custom_tokenizer, text,