From b52feaa965d326fbf1272f7e46694982eaf3a8d2 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Sat, 3 Jan 2026 13:44:41 -0800 Subject: [PATCH] refactor(utils): lazy load CustomStreamWrapper to improve import time - Move CustomStreamWrapper from streaming_handler to lazy loading via __getattr__ - Add type stub in TYPE_CHECKING block for mypy type checking - CustomStreamWrapper is not used internally in utils.py, only exported for other modules This reduces import time by deferring the streaming_handler module import until CustomStreamWrapper is actually accessed. --- litellm/utils.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index cb4471ca791..fc79384e0ca 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -99,7 +99,6 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) from litellm.litellm_core_utils.rules import Rules -from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.base_llm.google_genai.transformation import ( BaseGoogleGenAIGenerateContentConfig, ) @@ -340,6 +339,7 @@ if TYPE_CHECKING: LiteLLMLoggingObject, redact_message_input_output_from_logging, ) + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseConfig @@ -8905,4 +8905,14 @@ def __getattr__(name: str) -> Any: # noqa: PLR0915 _globals["redact_message_input_output_from_logging"] = _redact_message_input_output_from_logging return _globals["redact_message_input_output_from_logging"] + # Lazy load CustomStreamWrapper to avoid loading at module import time + if name == "CustomStreamWrapper": + # Check if already cached + if "CustomStreamWrapper" not in _globals: + from litellm.litellm_core_utils.streaming_handler import ( + CustomStreamWrapper as _CustomStreamWrapper, + ) + _globals["CustomStreamWrapper"] = _CustomStreamWrapper + return _globals["CustomStreamWrapper"] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}")