refactor: litellm init file #1

litellm/__init__.py

@@ -1039,8 +1039,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
@@ -1449,7 +1447,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
@@ -1504,3 +1501,60 @@ 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 import for cost_calculator functions to avoid loading the module at import time
+# This significantly reduces memory usage when importing litellm
+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,
+    )
+    
+    # Map names to imported functions
+    _cost_functions = {
+        "completion_cost": _completion_cost,
+        "cost_per_token": _cost_per_token,
+        "response_cost_calculator": _response_cost_calculator,
+    }
+    
+    # Cache the imported function in the module namespace
+    func = _cost_functions[name]
+    globals()[name] = func
+    
+    return func
+
+
+# Lazy import for litellm_logging to avoid loading the module at import time
+# This significantly reduces memory usage when importing litellm
+def _lazy_import_litellm_logging(name: str) -> Any:
+    """Lazy import for litellm_logging module."""
+    from litellm.litellm_core_utils.litellm_logging import (
+        Logging as _Logging,
+        modify_integration as _modify_integration,
+    )
+    
+    # Map names to imported objects
+    _logging_objects = {
+        "Logging": _Logging,
+        "modify_integration": _modify_integration,
+    }
+    
+    # Cache the imported object in the module namespace
+    obj = _logging_objects[name]
+    globals()[name] = obj
+    
+    return obj
+
+
+def __getattr__(name: str) -> Any:
+    """Lazy import for cost_calculator and litellm_logging functions."""
+    if name in ("completion_cost", "response_cost_calculator", "cost_per_token"):
+        return _lazy_import_cost_calculator(name)
+    
+    if name in ("Logging", "modify_integration"):
+        return _lazy_import_litellm_logging(name)
+    
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

litellm/images/main.py

@@ -6,11 +6,12 @@ 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
+# Logging is imported at module level since litellm_logging is already loaded via main.py imports
+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

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

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 (
@@ -247,6 +246,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
@@ -3457,6 +3473,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
@@ -3934,7 +3952,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

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,
@@ -77,7 +79,7 @@ from litellm.litellm_core_utils.health_check_utils import (
     _create_health_check_response,
     _filter_model_params,
 )
-from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj, Logging
 from litellm.litellm_core_utils.mock_functions import (
     mock_embedding,
     mock_image_generation,
@@ -6295,7 +6297,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: