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

[Refactor#2] litellm/init – Lazy-load utils to reduce memory + import time (#17171)

Browse source 
* fix: lazy load utils.py imports

Lazy-load most functions and response types from utils.py to avoid loading
tiktoken and other heavy dependencies at import time. This significantly
reduces memory usage when importing completion from litellm.
This commit is contained in:
Alexsander Hamir 2025-12-03 11:40:16 -08:00 • committed by GitHub
parent dcb7552b79
commit 56328e6535
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 354 additions and 130 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

147
litellm/__init__.py
View file

@ -1056,57 +1056,10 @@ from .timeout import timeout
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.core_helpers import remove_index_from_tool_calls
from litellm.litellm_core_utils.token_counter import get_modified_max_tokens
from .utils import (
client,
exception_type,
get_optional_params,
get_response_string,
token_counter,
create_pretrained_tokenizer,
create_tokenizer,
supports_function_calling,
supports_web_search,
supports_url_context,
supports_response_schema,
supports_parallel_function_calling,
supports_vision,
supports_audio_input,
supports_audio_output,
supports_system_messages,
supports_reasoning,
get_litellm_params,
acreate,
get_max_tokens,
get_model_info,
register_prompt_template,
validate_environment,
check_valid_key,
register_model,
encode,
decode,
_calculate_retry_after,
_should_retry,
get_supported_openai_params,
get_api_base,
get_first_chars_messages,
ModelResponse,
ModelResponseStream,
EmbeddingResponse,
ImageResponse,
TranscriptionResponse,
TextCompletionResponse,
get_provider_fields,
ModelResponseListIterator,
get_valid_models,
)
ALL_LITELLM_RESPONSE_TYPES = [
ModelResponse,
EmbeddingResponse,
ImageResponse,
TranscriptionResponse,
TextCompletionResponse,
]
# client must be imported immediately as it's used as a decorator at function definition time
from .utils import client
# Note: Most other utils imports are lazy-loaded via __getattr__ to avoid loading utils.py
# (which imports tiktoken) at import time
from .llms.bytez.chat.transformation import BytezChatConfig
from .llms.custom_llm import CustomLLM
@ -1538,56 +1491,6 @@ def set_global_gitlab_config(config: Dict[str, Any]) -> None:
# Lazy loading system for heavy modules to reduce initial import time and memory usage
def _lazy_import_cost_calculator(name: str) -> Any:
"""Lazy import for cost_calculator functions."""
from .cost_calculator import (
completion_cost as _completion_cost,
cost_per_token as _cost_per_token,
response_cost_calculator as _response_cost_calculator,
)
_cost_functions = {
"completion_cost": _completion_cost,
"cost_per_token": _cost_per_token,
"response_cost_calculator": _response_cost_calculator,
}
func = _cost_functions[name]
globals()[name] = func
return func
def _lazy_import_litellm_logging(name: str) -> Any:
"""Lazy import for litellm_logging module."""
try:
from litellm.litellm_core_utils.litellm_logging import (
Logging as _Logging,
modify_integration as _modify_integration,
)
_logging_objects = {
"Logging": _Logging,
"modify_integration": _modify_integration,
}
obj = _logging_objects[name]
globals()[name] = obj
return obj
except Exception as e:
raise AttributeError(
f"module {__name__!r} has no attribute {name!r}. "
f"Lazy import failed: {e}"
) from e
_LAZY_LOAD_REGISTRY: Dict[str, Callable[[str], Any]] = {
"completion_cost": _lazy_import_cost_calculator,
"cost_per_token": _lazy_import_cost_calculator,
"response_cost_calculator": _lazy_import_cost_calculator,
"Logging": _lazy_import_litellm_logging,
"modify_integration": _lazy_import_litellm_logging,
}
if TYPE_CHECKING:
cost_per_token: Callable[..., Tuple[float, float]]
@ -1598,7 +1501,45 @@ if TYPE_CHECKING:
def __getattr__(name: str) -> Any:
"""Lazy import handler for cost_calculator and litellm_logging functions."""
if name in _LAZY_LOAD_REGISTRY:
return _LAZY_LOAD_REGISTRY[name](name)
# Lazy load cost_calculator functions
_cost_calculator_names = (
"completion_cost",
"cost_per_token",
"response_cost_calculator",
)
if name in _cost_calculator_names:
from ._lazy_imports import _lazy_import_cost_calculator
return _lazy_import_cost_calculator(name)
# Lazy load litellm_logging functions
_litellm_logging_names = (
"Logging",
"modify_integration",
)
if name in _litellm_logging_names:
from ._lazy_imports import _lazy_import_litellm_logging
return _lazy_import_litellm_logging(name)
# Lazy load utils functions
_utils_names = (
"exception_type", "get_optional_params", "get_response_string", "token_counter",
"create_pretrained_tokenizer", "create_tokenizer", "supports_function_calling",
"supports_web_search", "supports_url_context", "supports_response_schema",
"supports_parallel_function_calling", "supports_vision", "supports_audio_input",
"supports_audio_output", "supports_system_messages", "supports_reasoning",
"get_litellm_params", "acreate", "get_max_tokens", "get_model_info",
"register_prompt_template", "validate_environment", "check_valid_key",
"register_model", "encode", "decode", "_calculate_retry_after", "_should_retry",
"get_supported_openai_params", "get_api_base", "get_first_chars_messages",
"ModelResponse", "ModelResponseStream", "EmbeddingResponse", "ImageResponse",
"TranscriptionResponse", "TextCompletionResponse", "get_provider_fields",
"ModelResponseListIterator", "get_valid_models",
)
if name in _utils_names:
from ._lazy_imports import _lazy_import_utils
return _lazy_import_utils(name)
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
# ALL_LITELLM_RESPONSE_TYPES is lazy-loaded via __getattr__ to avoid loading utils at import time

259
litellm/_lazy_imports.py Normal file
View file

@ -0,0 +1,259 @@
from typing import Any
import sys
def _get_litellm_globals() -> dict:
"""Helper to get the globals dictionary of the litellm module."""
return sys.modules["litellm"].__dict__
# Lazy import for utils module - imports only the requested item by name.
# Note: PLR0915 (too many statements) is suppressed because the many if statements
# are intentional - each attribute is imported individually only when requested,
# ensuring true lazy imports rather than importing the entire utils module.
def _lazy_import_utils(name: str) -> Any: # noqa: PLR0915
"""Lazy import for utils module - imports only the requested item by name."""
_globals = _get_litellm_globals()
if name == "exception_type":
from .utils import exception_type as _exception_type
_globals["exception_type"] = _exception_type
return _exception_type
if name == "get_optional_params":
from .utils import get_optional_params as _get_optional_params
_globals["get_optional_params"] = _get_optional_params
return _get_optional_params
if name == "get_response_string":
from .utils import get_response_string as _get_response_string
_globals["get_response_string"] = _get_response_string
return _get_response_string
if name == "token_counter":
from .utils import token_counter as _token_counter
_globals["token_counter"] = _token_counter
return _token_counter
if name == "create_pretrained_tokenizer":
from .utils import create_pretrained_tokenizer as _create_pretrained_tokenizer
_globals["create_pretrained_tokenizer"] = _create_pretrained_tokenizer
return _create_pretrained_tokenizer
if name == "create_tokenizer":
from .utils import create_tokenizer as _create_tokenizer
_globals["create_tokenizer"] = _create_tokenizer
return _create_tokenizer
if name == "supports_function_calling":
from .utils import supports_function_calling as _supports_function_calling
_globals["supports_function_calling"] = _supports_function_calling
return _supports_function_calling
if name == "supports_web_search":
from .utils import supports_web_search as _supports_web_search
_globals["supports_web_search"] = _supports_web_search
return _supports_web_search
if name == "supports_url_context":
from .utils import supports_url_context as _supports_url_context
_globals["supports_url_context"] = _supports_url_context
return _supports_url_context
if name == "supports_response_schema":
from .utils import supports_response_schema as _supports_response_schema
_globals["supports_response_schema"] = _supports_response_schema
return _supports_response_schema
if name == "supports_parallel_function_calling":
from .utils import supports_parallel_function_calling as _supports_parallel_function_calling
_globals["supports_parallel_function_calling"] = _supports_parallel_function_calling
return _supports_parallel_function_calling
if name == "supports_vision":
from .utils import supports_vision as _supports_vision
_globals["supports_vision"] = _supports_vision
return _supports_vision
if name == "supports_audio_input":
from .utils import supports_audio_input as _supports_audio_input
_globals["supports_audio_input"] = _supports_audio_input
return _supports_audio_input
if name == "supports_audio_output":
from .utils import supports_audio_output as _supports_audio_output
_globals["supports_audio_output"] = _supports_audio_output
return _supports_audio_output
if name == "supports_system_messages":
from .utils import supports_system_messages as _supports_system_messages
_globals["supports_system_messages"] = _supports_system_messages
return _supports_system_messages
if name == "supports_reasoning":
from .utils import supports_reasoning as _supports_reasoning
_globals["supports_reasoning"] = _supports_reasoning
return _supports_reasoning
if name == "get_litellm_params":
from .utils import get_litellm_params as _get_litellm_params
_globals["get_litellm_params"] = _get_litellm_params
return _get_litellm_params
if name == "acreate":
from .utils import acreate as _acreate
_globals["acreate"] = _acreate
return _acreate
if name == "get_max_tokens":
from .utils import get_max_tokens as _get_max_tokens
_globals["get_max_tokens"] = _get_max_tokens
return _get_max_tokens
if name == "get_model_info":
from .utils import get_model_info as _get_model_info
_globals["get_model_info"] = _get_model_info
return _get_model_info
if name == "register_prompt_template":
from .utils import register_prompt_template as _register_prompt_template
_globals["register_prompt_template"] = _register_prompt_template
return _register_prompt_template
if name == "validate_environment":
from .utils import validate_environment as _validate_environment
_globals["validate_environment"] = _validate_environment
return _validate_environment
if name == "check_valid_key":
from .utils import check_valid_key as _check_valid_key
_globals["check_valid_key"] = _check_valid_key
return _check_valid_key
if name == "register_model":
from .utils import register_model as _register_model
_globals["register_model"] = _register_model
return _register_model
if name == "encode":
from .utils import encode as _encode
_globals["encode"] = _encode
return _encode
if name == "decode":
from .utils import decode as _decode
_globals["decode"] = _decode
return _decode
if name == "_calculate_retry_after":
from .utils import _calculate_retry_after as __calculate_retry_after
_globals["_calculate_retry_after"] = __calculate_retry_after
return __calculate_retry_after
if name == "_should_retry":
from .utils import _should_retry as __should_retry
_globals["_should_retry"] = __should_retry
return __should_retry
if name == "get_supported_openai_params":
from .utils import get_supported_openai_params as _get_supported_openai_params
_globals["get_supported_openai_params"] = _get_supported_openai_params
return _get_supported_openai_params
if name == "get_api_base":
from .utils import get_api_base as _get_api_base
_globals["get_api_base"] = _get_api_base
return _get_api_base
if name == "get_first_chars_messages":
from .utils import get_first_chars_messages as _get_first_chars_messages
_globals["get_first_chars_messages"] = _get_first_chars_messages
return _get_first_chars_messages
if name == "ModelResponse":
from .utils import ModelResponse as _ModelResponse
_globals["ModelResponse"] = _ModelResponse
return _ModelResponse
if name == "ModelResponseStream":
from .utils import ModelResponseStream as _ModelResponseStream
_globals["ModelResponseStream"] = _ModelResponseStream
return _ModelResponseStream
if name == "EmbeddingResponse":
from .utils import EmbeddingResponse as _EmbeddingResponse
_globals["EmbeddingResponse"] = _EmbeddingResponse
return _EmbeddingResponse
if name == "ImageResponse":
from .utils import ImageResponse as _ImageResponse
_globals["ImageResponse"] = _ImageResponse
return _ImageResponse
if name == "TranscriptionResponse":
from .utils import TranscriptionResponse as _TranscriptionResponse
_globals["TranscriptionResponse"] = _TranscriptionResponse
return _TranscriptionResponse
if name == "TextCompletionResponse":
from .utils import TextCompletionResponse as _TextCompletionResponse
_globals["TextCompletionResponse"] = _TextCompletionResponse
return _TextCompletionResponse
if name == "get_provider_fields":
from .utils import get_provider_fields as _get_provider_fields
_globals["get_provider_fields"] = _get_provider_fields
return _get_provider_fields
if name == "ModelResponseListIterator":
from .utils import ModelResponseListIterator as _ModelResponseListIterator
_globals["ModelResponseListIterator"] = _ModelResponseListIterator
return _ModelResponseListIterator
if name == "get_valid_models":
from .utils import get_valid_models as _get_valid_models
_globals["get_valid_models"] = _get_valid_models
return _get_valid_models
raise AttributeError(f"Utils lazy import: unknown attribute {name!r}")
def _lazy_import_cost_calculator(name: str) -> Any:
"""Lazy import for cost_calculator functions."""
_globals = _get_litellm_globals()
from .cost_calculator import (
completion_cost as _completion_cost,
cost_per_token as _cost_per_token,
response_cost_calculator as _response_cost_calculator,
)
_cost_functions = {
"completion_cost": _completion_cost,
"cost_per_token": _cost_per_token,
"response_cost_calculator": _response_cost_calculator,
}
func = _cost_functions[name]
_globals[name] = func
return func
def _lazy_import_litellm_logging(name: str) -> Any:
"""Lazy import for litellm_logging module."""
_globals = _get_litellm_globals()
try:
from litellm.litellm_core_utils.litellm_logging import (
Logging as _Logging,
modify_integration as _modify_integration,
)
_logging_objects = {
"Logging": _Logging,
"modify_integration": _modify_integration,
}
obj = _logging_objects[name]
_globals[name] = obj
return obj
except Exception as e:
raise AttributeError(
f"module 'litellm' has no attribute {name!r}. "
f"Lazy import failed: {e}"
) from e

4
litellm/images/main.py
View file

@ -6,7 +6,9 @@ from typing import Any, Coroutine, Dict, List, Literal, Optional, Union, cast, o
import httpx
import litellm
from litellm import client, exception_type, get_litellm_params
from litellm.utils import exception_type, get_litellm_params
# client is imported from litellm as it's a decorator
from litellm import client
from litellm.constants import DEFAULT_IMAGE_ENDPOINT_MODEL
from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT
from litellm.exceptions import LiteLLMUnknownProvider

2
litellm/llms/azure/azure.py
View file

@ -1020,7 +1020,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
headers: dict,
client=None,
timeout=None,
) -> litellm.ImageResponse:
) -> ImageResponse:
response: Optional[dict] = None
try:

4
litellm/llms/azure_ai/embed/handler.py
View file

@ -58,7 +58,7 @@ class AzureAIEmbedding(OpenAIChatCompletion):
data: ImageEmbeddingRequest,
timeout: float,
logging_obj,
model_response: litellm.EmbeddingResponse,
model_response: EmbeddingResponse,
optional_params: dict,
api_key: Optional[str],
api_base: Optional[str],
@ -138,7 +138,7 @@ class AzureAIEmbedding(OpenAIChatCompletion):
input: List,
timeout: float,
logging_obj,
model_response: litellm.EmbeddingResponse,
model_response: EmbeddingResponse,
optional_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,

3
litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py
View file

@ -10,7 +10,6 @@ from typing import Any, List, Optional
import httpx
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.types.llms.bedrock import BedrockInvokeNovaRequest
from litellm.types.llms.openai import AllMessageValues
@ -80,7 +79,7 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig):
encoding: Any,
api_key: Optional[str] = None,
json_mode: Optional[bool] = None,
) -> litellm.ModelResponse:
) -> ModelResponse:
return AmazonConverseConfig.transform_response(
self,
model,

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

@ -27,6 +27,25 @@ class BedrockError(BaseLLMException):
pass
# Lazy import cache to avoid circular imports and performance impact
_get_model_info = None
def get_cached_model_info():
"""
Lazy import and cache get_model_info to avoid circular imports.
This function is used by bedrock transformation classes that need get_model_info
but cannot import it at module level due to circular import issues.
The function is cached after first use to avoid performance impact.
"""
global _get_model_info
if _get_model_info is None:
from litellm import get_model_info
_get_model_info = get_model_info
return _get_model_info
class AmazonBedrockGlobalConfig:
def __init__(self):
pass

3
litellm/llms/bedrock/image/amazon_nova_canvas_transformation.py
View file

@ -3,7 +3,6 @@ from typing import Any, Dict, List, Optional
from openai.types.image import Image
from litellm import get_model_info
from litellm.types.llms.bedrock import (
AmazonNovaCanvasColorGuidedGenerationParams,
AmazonNovaCanvasColorGuidedRequest,
@ -15,6 +14,7 @@ from litellm.types.llms.bedrock import (
AmazonNovaCanvasTextToImageRequest,
AmazonNovaCanvasTextToImageResponse,
)
from litellm.llms.bedrock.common_utils import get_cached_model_info
from litellm.types.utils import ImageResponse
@ -207,6 +207,7 @@ class AmazonNovaCanvasConfig:
size: Optional[str] = None,
optional_params: Optional[dict] = None,
) -> float:
get_model_info = get_cached_model_info()
model_info = get_model_info(
model=model,
custom_llm_provider="bedrock",

3
litellm/llms/bedrock/image/amazon_stability1_transformation.py
View file

@ -5,7 +5,7 @@ from typing import List, Optional
from openai.types.image import Image
from litellm import get_model_info
from litellm.llms.bedrock.common_utils import get_cached_model_info
from litellm.types.utils import ImageResponse
@ -151,6 +151,7 @@ class AmazonStabilityConfig:
size = size or "1024-x-1024"
model = f"{size}/{steps}/{model}"
get_model_info = get_cached_model_info()
model_info = get_model_info(
model=model,
custom_llm_provider="bedrock",

3
litellm/llms/bedrock/image/amazon_stability3_transformation.py
View file

@ -3,12 +3,12 @@ from typing import List, Optional
from openai.types.image import Image
from litellm import get_model_info
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.types.llms.bedrock import (
AmazonStability3TextToImageRequest,
AmazonStability3TextToImageResponse,
)
from litellm.llms.bedrock.common_utils import get_cached_model_info
from litellm.types.utils import ImageResponse
@ -115,6 +115,7 @@ class AmazonStability3Config:
size: Optional[str] = None,
optional_params: Optional[dict] = None,
) -> float:
get_model_info = get_cached_model_info()
model_info = get_model_info(
model=model,
custom_llm_provider="bedrock",

2
litellm/llms/bedrock/image/amazon_titan_transformation.py
View file

@ -7,7 +7,7 @@ from typing import List, Optional
from openai.types.image import Image
from litellm import get_model_info
from litellm.utils import get_model_info
from litellm.types.llms.bedrock import (
AmazonNovaCanvasImageGenerationConfig,
AmazonTitanImageGenerationRequestBody,

4
litellm/llms/ovhcloud/chat/transformation.py
View file

@ -7,7 +7,9 @@ More information on our website: https://endpoints.ai.cloud.ovh.net
from typing import Optional, Union, List
import httpx
from litellm import ModelResponseStream, OpenAIGPTConfig, get_model_info, verbose_logger
from litellm.utils import ModelResponseStream, get_model_info
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
from litellm._logging import verbose_logger
from litellm.llms.ovhcloud.utils import OVHCloudException
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.llms.base_llm.chat.transformation import BaseLLMException

3
litellm/llms/together_ai/chat.py
View file

@ -8,7 +8,8 @@ Docs: https://docs.together.ai/reference/completions-1
from typing import Optional
from litellm import get_model_info, verbose_logger
from litellm.utils import get_model_info
from litellm._logging import verbose_logger
from ..openai.chat.gpt_transformation import OpenAIGPTConfig

3
litellm/llms/vertex_ai/common_utils.py
View file

@ -5,7 +5,8 @@ from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union, get_ty
import httpx
import litellm
from litellm import supports_response_schema, supports_system_messages, verbose_logger
from litellm.utils import supports_response_schema, supports_system_messages
from litellm._logging import verbose_logger
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_defs
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter

2
litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py
View file

@ -8,7 +8,7 @@ from typing import Any, Literal, Optional, Union
import httpx
import litellm
from litellm import EmbeddingResponse
from litellm.types.utils import EmbeddingResponse
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,

2
litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py
View file

@ -6,7 +6,7 @@ Why separate file? Make it easy to see how transformation works
from typing import List
from litellm import EmbeddingResponse
from litellm.types.utils import EmbeddingResponse
from litellm.types.llms.openai import EmbeddingInput
from litellm.types.llms.vertex_ai import (
ContentType,

2
litellm/llms/vertex_ai/image_generation/image_generation_handler.py
View file

@ -176,7 +176,7 @@ class VertexImageGeneration(VertexLLM):
vertex_project: Optional[str],
vertex_location: Optional[str],
vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES],
model_response: litellm.ImageResponse,
model_response: ImageResponse,
logging_obj: Any,
model: str = "imagegeneration", # vertex ai uses imagegeneration as the default model
client: Optional[AsyncHTTPHandler] = None,

4
litellm/llms/vertex_ai/multimodal_embeddings/embedding_handler.py
View file

@ -147,13 +147,13 @@ class VertexMultimodalEmbedding(VertexLLM):
optional_params: dict,
litellm_params: dict,
data: dict,
model_response: litellm.EmbeddingResponse,
model_response: EmbeddingResponse,
timeout: Optional[Union[float, httpx.Timeout]],
logging_obj: LiteLLMLoggingObj,
headers={},
client: Optional[AsyncHTTPHandler] = None,
api_key: Optional[str] = None,
) -> litellm.EmbeddingResponse:
) -> EmbeddingResponse:
if client is None:
_params = {}
if timeout is not None:

4
litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py
View file

@ -137,7 +137,7 @@ class VertexEmbedding(VertexBase):
self,
model: str,
input: Union[list, str],
model_response: litellm.EmbeddingResponse,
model_response: EmbeddingResponse,
logging_obj: LiteLLMLoggingObject,
optional_params: dict,
custom_llm_provider: Literal[
@ -152,7 +152,7 @@ class VertexEmbedding(VertexBase):
gemini_api_key: Optional[str] = None,
extra_headers: Optional[dict] = None,
encoding=None,
) -> litellm.EmbeddingResponse:
) -> EmbeddingResponse:
"""
Async embedding implementation
"""

11
litellm/main.py
View file

@ -52,13 +52,10 @@ from pydantic import BaseModel
from typing_extensions import overload
import litellm
from litellm import ( # type: ignore
client,
exception_type,
get_litellm_params,
get_optional_params,
)
# client must be imported from litellm as it's a decorator used at function definition time
from litellm import client
# Other utils are imported directly to avoid circular imports
from litellm.utils import exception_type, get_litellm_params, get_optional_params
# Logging is imported lazily when needed to avoid loading litellm_logging at import time
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging