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

Merge pull request #17089 from BerriAI/litellm_refactor_01

Browse source 
[Refactor#1] `litellm/init` – Lazy-load `cost_calculator` & `logging` to reduce memory + import time
This commit is contained in:
Alexsander Hamir 2025-11-26 14:41:46 -08:00 • committed by GitHub
commit c0d9facd96
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 132 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

73
litellm/__init__.py
View file

@ -20,6 +20,9 @@ from typing import (
Literal,
get_args,
TYPE_CHECKING,
Tuple,
overload,
Type,
)
from litellm.types.integrations.datadog_llm_obs import DatadogLLMObsInitParams
from litellm.types.integrations.datadog import DatadogInitParams
@ -1044,8 +1047,6 @@ openai_image_generation_models = ["dall-e-2", "dall-e-3"]
openai_video_generation_models = ["sora-2"]
from .timeout import timeout
from .cost_calculator import completion_cost
from litellm.litellm_core_utils.litellm_logging import Logging, modify_integration
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
@ -1462,7 +1463,6 @@ from .vector_store_files.main import (
update as vector_store_file_update,
)
from .scheduler import *
from .cost_calculator import response_cost_calculator, cost_per_token
### ADAPTERS ###
from .types.adapter import AdapterItem
@ -1520,3 +1520,70 @@ def set_global_gitlab_config(config: Dict[str, Any]) -> None:
"""Set global BitBucket configuration for prompt management."""
global global_gitlab_config
global_gitlab_config = config
# 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]]
completion_cost: Callable[..., float]
response_cost_calculator: Any
modify_integration: Any
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)
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

4
litellm/images/main.py
View file

@ -6,11 +6,11 @@ from typing import Any, Coroutine, Dict, List, Literal, Optional, Union, cast, o
import httpx
import litellm
from litellm import Logging, client, exception_type, get_litellm_params
from litellm import client, exception_type, get_litellm_params
from litellm.constants import DEFAULT_IMAGE_ENDPOINT_MODEL
from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT
from litellm.exceptions import LiteLLMUnknownProvider
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import Logging, Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.mock_functions import mock_image_generation
from litellm.llms.base_llm import BaseImageEditConfig, BaseImageGenerationConfig
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler

24
litellm/integrations/prometheus.py
View file

@ -24,13 +24,29 @@ from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth
from litellm.types.integrations.prometheus import *
from litellm.types.integrations.prometheus import _sanitize_prometheus_label_name
from litellm.types.utils import StandardLoggingPayload
from litellm.utils import get_end_user_id_for_cost_tracking
if TYPE_CHECKING:
from apscheduler.schedulers.asyncio import AsyncIOScheduler
else:
AsyncIOScheduler = Any
# Cached lazy import for get_end_user_id_for_cost_tracking
# Module-level cache to avoid repeated imports while preserving memory benefits
_get_end_user_id_for_cost_tracking = None
def _get_cached_end_user_id_for_cost_tracking():
"""
Get cached get_end_user_id_for_cost_tracking function.
Lazy imports on first call to avoid loading utils.py at import time (60MB saved).
Subsequent calls use cached function for better performance.
"""
global _get_end_user_id_for_cost_tracking
if _get_end_user_id_for_cost_tracking is None:
from litellm.utils import get_end_user_id_for_cost_tracking
_get_end_user_id_for_cost_tracking = get_end_user_id_for_cost_tracking
return _get_end_user_id_for_cost_tracking
class PrometheusLogger(CustomLogger):
# Class variables or attributes
@ -778,6 +794,8 @@ class PrometheusLogger(CustomLogger):
model = kwargs.get("model", "")
litellm_params = kwargs.get("litellm_params", {}) or {}
_metadata = litellm_params.get("metadata", {})
get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking()
end_user_id = get_end_user_id_for_cost_tracking(
litellm_params, service_type="prometheus"
)
@ -1164,6 +1182,8 @@ class PrometheusLogger(CustomLogger):
"standard_logging_object", {}
)
litellm_params = kwargs.get("litellm_params", {}) or {}
get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking()
end_user_id = get_end_user_id_for_cost_tracking(
litellm_params, service_type="prometheus"
)
@ -2249,6 +2269,8 @@ def prometheus_label_factory(
}
if UserAPIKeyLabelNames.END_USER.value in filtered_labels:
get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking()
filtered_labels["end_user"] = get_end_user_id_for_cost_tracking(
litellm_params={"user_api_key_end_user_id": enum_values.end_user},
service_type="prometheus",

23
litellm/litellm_core_utils/litellm_logging.py
View file

@ -58,7 +58,6 @@ from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.deepeval.deepeval import DeepEvalLogger
from litellm.integrations.mlflow import MlflowLogger
from litellm.integrations.prometheus import PrometheusLogger
from litellm.integrations.sqs import SQSLogger
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
@ -248,6 +247,23 @@ class ServiceTraceIDCache:
in_memory_trace_id_cache = ServiceTraceIDCache()
in_memory_dynamic_logger_cache = DynamicLoggingCache()
# Cached lazy import for PrometheusLogger
# Module-level cache to avoid repeated imports while preserving memory benefits
_PrometheusLogger = None
def _get_cached_prometheus_logger():
"""
Get cached PrometheusLogger class.
Lazy imports on first call to avoid loading prometheus.py and utils.py at import time (60MB saved).
Subsequent calls use cached class for better performance.
"""
global _PrometheusLogger
if _PrometheusLogger is None:
from litellm.integrations.prometheus import PrometheusLogger
_PrometheusLogger = PrometheusLogger
return _PrometheusLogger
class Logging(LiteLLMLoggingBaseClass):
global supabaseClient, promptLayerLogger, weightsBiasesLogger, logfireLogger, capture_exception, add_breadcrumb, lunaryLogger, logfireLogger, prometheusLogger, slack_app
@ -3462,6 +3478,8 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
_in_memory_loggers.append(_literalai_logger)
return _literalai_logger # type: ignore
elif logging_integration == "prometheus":
PrometheusLogger = _get_cached_prometheus_logger()
for callback in _in_memory_loggers:
if isinstance(callback, PrometheusLogger):
return callback # type: ignore
@ -3939,7 +3957,8 @@ def get_custom_logger_compatible_class( # noqa: PLR0915
for callback in _in_memory_loggers:
if isinstance(callback, LiteralAILogger):
return callback
elif logging_integration == "prometheus" and PrometheusLogger is not None:
elif logging_integration == "prometheus":
PrometheusLogger = _get_cached_prometheus_logger()
for callback in _in_memory_loggers:
if isinstance(callback, PrometheusLogger):
return callback

18
litellm/llms/vertex_ai/videos/transformation.py
View file

@ -7,7 +7,7 @@ Based on: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/model-refer
import base64
import time
from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union, cast
import httpx
from httpx._types import RequestFiles
@ -172,19 +172,11 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
"""
# Extract Vertex AI parameters using safe helpers from VertexBase
# Use safe_get_* methods that don't mutate litellm_params dict
if litellm_params is None:
litellm_params_dict: Dict[str, Any] = {}
elif isinstance(litellm_params, dict):
litellm_params_dict = litellm_params
else:
litellm_params_dict = litellm_params.model_dump()
# Ensure litellm_params is a dict for type checking
params_dict: Dict[str, Any] = cast(Dict[str, Any], litellm_params) if litellm_params is not None else {}
vertex_project = VertexBase.safe_get_vertex_ai_project(
litellm_params=litellm_params_dict
)
vertex_credentials = VertexBase.safe_get_vertex_ai_credentials(
litellm_params=litellm_params_dict
)
vertex_project = VertexBase.safe_get_vertex_ai_project(litellm_params=params_dict)
vertex_credentials = VertexBase.safe_get_vertex_ai_credentials(litellm_params=params_dict)
# Get access token from Vertex credentials
access_token, project_id = self.get_access_token(

16
litellm/main.py
View file

@ -53,12 +53,14 @@ from typing_extensions import overload
import litellm
from litellm import ( # type: ignore
Logging,
client,
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
from litellm.constants import (
DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT,
@ -1154,7 +1156,7 @@ def completion( # type: ignore # noqa: PLR0915
api_base = base_url
if num_retries is not None:
max_retries = num_retries
logging: Logging = cast(Logging, litellm_logging_obj)
logging: LiteLLMLoggingObj = cast(LiteLLMLoggingObj, litellm_logging_obj)
fallbacks = fallbacks or litellm.model_fallbacks
if fallbacks is not None:
return completion_with_fallbacks(**args)
@ -4162,7 +4164,7 @@ def embedding( # noqa: PLR0915
litellm_params_dict = get_litellm_params(**kwargs)
logging: Logging = litellm_logging_obj # type: ignore
logging: LiteLLMLoggingObj = litellm_logging_obj # type: ignore
logging.update_environment_variables(
model=model,
user=user,
@ -5878,7 +5880,7 @@ def speech( # noqa: PLR0915
kwargs=kwargs,
)
logging_obj: Logging = cast(Logging, kwargs.get("litellm_logging_obj"))
logging_obj: LiteLLMLoggingObj = cast(LiteLLMLoggingObj, kwargs.get("litellm_logging_obj"))
logging_obj.update_environment_variables(
model=model,
user=user,
@ -6235,6 +6237,10 @@ async def ahealth_check(
}
"""
from litellm.litellm_core_utils.health_check_helpers import HealthCheckHelpers
from litellm.litellm_core_utils.cached_imports import get_litellm_logging_class
# Use cached import helper to lazy-load Logging class (only loads when function is called)
Logging = get_litellm_logging_class()
# Map modes to their corresponding health check calls
#########################################################
@ -6423,7 +6429,7 @@ def stream_chunk_builder( # noqa: PLR0915
messages: Optional[list] = None,
start_time=None,
end_time=None,
logging_obj: Optional[Logging] = None,
logging_obj: Optional["Logging"] = None,
) -> Optional[Union[ModelResponse, TextCompletionResponse]]:
try:
if chunks is None: