diff --git a/litellm/__init__.py b/litellm/__init__.py index 3a715ff1dfb..b214db74a76 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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}") diff --git a/litellm/images/main.py b/litellm/images/main.py index 786136e6699..878fce83f17 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -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 diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 8186006f8c8..4ce818f0cef 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -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", diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5b4fa6b7e24..38decd8af99 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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 diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index f6678510493..4cf3e036d3f 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -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( diff --git a/litellm/main.py b/litellm/main.py index 3e51785ac63..afc7a36fb40 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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: