From b4b8133d47e25e8d2557de10e4a0f5a881803f9a Mon Sep 17 00:00:00 2001 From: AlexsanderHamir Date: Tue, 18 Nov 2025 14:47:08 -0800 Subject: [PATCH] refactor: litellm init file #1 --- litellm/__init__.py | 60 ++++++++++++++++++- litellm/images/main.py | 5 +- litellm/integrations/prometheus.py | 24 +++++++- litellm/litellm_core_utils/litellm_logging.py | 23 ++++++- litellm/main.py | 8 ++- 5 files changed, 109 insertions(+), 11 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index d93f44c37e0..739cef04e2d 100644 --- a/litellm/__init__.py +++ b/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}") diff --git a/litellm/images/main.py b/litellm/images/main.py index 333a751b045..2e93765c8fb 100644 --- a/litellm/images/main.py +++ b/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 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 9c4a7e38768..67407ff4c7a 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 ( @@ -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 diff --git a/litellm/main.py b/litellm/main.py index b082b491f24..57f8256ec35 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, @@ -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: