[Refactor] litellm/init.py: lazy-load heavy client decorator imports (#18064)

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",

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,