Make LiteLLM imports lazy

dspy/clients/__init__.py

@@ -3,13 +3,14 @@
 from pathlib import Path
 from typing import Any
 
-import litellm
-
 from dspy.clients.base_lm import BaseLM, inspect_history
 from dspy.clients.cache import Cache
 from dspy.clients.embedding import Embedder
 from dspy.clients.lm import LM
 from dspy.clients.provider import Provider, TrainingJob
+from dspy.utils.lazy_import import require
+
+litellm = require("litellm", extra="litellm", feature="LiteLLM logging")
 
 logger = logging.getLogger(__name__)
 
@@ -56,9 +57,6 @@ def configure_cache(
     dspy.cache = DSPY_CACHE
 
 
-litellm.telemetry = False
-litellm.cache = None  # By default we disable LiteLLM cache and use DSPy on-disk cache.
-
 
 def _get_dspy_cache():
     disk_cache_dir = os.environ.get("DSPY_CACHEDIR") or os.path.join(Path.home(), ".dspy_cache")
@@ -91,7 +89,7 @@ def _get_dspy_cache():
 def configure_litellm_logging(level: str = "ERROR"):
     """Configure LiteLLM logging to the specified level."""
     # Litellm uses a global logger called `verbose_logger` to control all loggings.
-    from litellm._logging import verbose_logger
+    verbose_logger = litellm._logging.verbose_logger
 
     numeric_logging_level = getattr(logging, level)
 
@@ -102,17 +100,15 @@ def configure_litellm_logging(level: str = "ERROR"):
 
 def enable_litellm_logging():
     litellm.suppress_debug_info = False
+    litellm._dspy_logging_configured = True
     configure_litellm_logging("DEBUG")
 
 
 def disable_litellm_logging():
     litellm.suppress_debug_info = True
+    litellm._dspy_logging_configured = True
     configure_litellm_logging("ERROR")
 
-
-# By default, we disable LiteLLM logging for clean logging
-disable_litellm_logging()
-
 __all__ = [
     "BaseLM",
     "LM",

dspy/clients/embedding.py

@@ -2,13 +2,23 @@
 
 from typing import Any, Callable
 
-import litellm
-
 from dspy.clients.cache import request_cache
 from dspy.utils.lazy_import import require
 
 np = require("numpy")
 
+
+def _configure_litellm_defaults(litellm):
+    litellm.telemetry = False
+    litellm.cache = None  # By default we disable LiteLLM cache and use DSPy on-disk cache.
+    if not getattr(litellm, "_dspy_logging_configured", False):
+        litellm.suppress_debug_info = True
+        litellm._dspy_logging_configured = True
+
+
+litellm = require("litellm", extra="litellm", feature="dspy.Embedder", on_load=_configure_litellm_defaults)
+
+
 class Embedder:
     """DSPy embedding class.
 

dspy/clients/lm.py

@@ -7,10 +7,8 @@
 from typing import Any, Literal, cast
 
 import anyio.from_thread
-import litellm
 import pydantic
 from anyio.streams.memory import MemoryObjectSendStream
-from litellm import ContextWindowExceededError as LitellmContextWindowExceededError
 
 import dspy
 from dspy.clients.cache import request_cache
@@ -20,12 +18,29 @@
 from dspy.dsp.utils.settings import settings
 from dspy.utils.callback import BaseCallback
 from dspy.utils.exceptions import ContextWindowExceededError
+from dspy.utils.lazy_import import require
 
 from .base_lm import BaseLM
 
 logger = logging.getLogger(__name__)
 
 
+def _configure_litellm_defaults(litellm):
+    litellm.telemetry = False
+    litellm.cache = None  # By default we disable LiteLLM cache and use DSPy on-disk cache.
+    if not getattr(litellm, "_dspy_logging_configured", False):
+        litellm.suppress_debug_info = True
+        litellm._dspy_logging_configured = True
+
+
+litellm = require("litellm", extra="litellm", feature="dspy.LM", on_load=_configure_litellm_defaults)
+
+
+def _is_litellm_context_window_error(error: Exception) -> bool:
+    error_type = type(error)
+    return error_type.__module__.startswith("litellm") and "ContextWindowExceeded" in error_type.__name__
+
+
 class LM(BaseLM):
     """
     A language model supporting chat or text completion requests for use with DSPy modules.
@@ -189,8 +204,10 @@ def forward(
                 num_retries=self.num_retries,
                 cache=litellm_cache_args,
             )
-        except LitellmContextWindowExceededError as e:
-            raise ContextWindowExceededError(model=self.model) from e
+        except Exception as e:
+            if _is_litellm_context_window_error(e):
+                raise ContextWindowExceededError(model=self.model) from e
+            raise
 
         self._check_truncation(results)
 
@@ -230,8 +247,10 @@ async def aforward(
                 num_retries=self.num_retries,
                 cache=litellm_cache_args,
             )
-        except LitellmContextWindowExceededError as e:
-            raise ContextWindowExceededError(model=self.model) from e
+        except Exception as e:
+            if _is_litellm_context_window_error(e):
+                raise ContextWindowExceededError(model=self.model) from e
+            raise
 
         self._check_truncation(results)
 

dspy/streaming/streamify.py

@@ -6,11 +6,9 @@
 from queue import Queue
 from typing import TYPE_CHECKING, Any, AsyncGenerator, Awaitable, Callable, Generator
 
-import litellm
 import orjson
 from anyio import create_memory_object_stream, create_task_group
 from anyio.streams.memory import MemoryObjectSendStream
-from litellm import ModelResponseStream
 
 from dspy.dsp.utils.settings import settings
 from dspy.primitives.prediction import Prediction
@@ -20,6 +18,12 @@
 
 logger = logging.getLogger(__name__)
 
+
+def _is_litellm_model_response_stream(value: Any) -> bool:
+    cls = type(value)
+    return cls.__name__ == "ModelResponseStream" and cls.__module__.startswith("litellm")
+
+
 if TYPE_CHECKING:
     from dspy.primitives.module import Module
 
@@ -178,7 +182,7 @@ async def async_streamer(*args, **kwargs):
             tg.start_soon(generator, args, kwargs, send_stream)
 
             async for value in receive_stream:
-                if isinstance(value, ModelResponseStream):
+                if _is_litellm_model_response_stream(value):
                     if len(predict_id_to_listener) == 0:
                         # No listeners are configured, yield the chunk directly for backwards compatibility.
                         yield value
@@ -271,7 +275,7 @@ async def streaming_response(streamer: AsyncGenerator) -> AsyncGenerator:
         if isinstance(value, Prediction):
             data = {"prediction": dict(value.items(include_dspy=False))}
             yield f"data: {orjson.dumps(data).decode()}\n\n"
-        elif isinstance(value, litellm.ModelResponseStream):
+        elif _is_litellm_model_response_stream(value):
             data = {"chunk": value.json()}
             yield f"data: {orjson.dumps(data).decode()}\n\n"
         elif isinstance(value, str) and value.startswith("data:"):

dspy/streaming/streaming_listener.py

@@ -5,7 +5,6 @@
 from typing import TYPE_CHECKING, Any
 
 import jiter
-from litellm import ModelResponseStream
 
 from dspy.adapters.chat_adapter import ChatAdapter
 from dspy.adapters.json_adapter import JSONAdapter
@@ -15,6 +14,8 @@
 from dspy.streaming.messages import StreamResponse
 
 if TYPE_CHECKING:
+    from litellm import ModelResponseStream
+
     from dspy.primitives.module import Module
 
 ADAPTER_SUPPORT_STREAMING = [ChatAdapter, XMLAdapter, JSONAdapter]
@@ -112,7 +113,7 @@ def _could_form_end_identifier(self, concat_message: str, adapter_name: str) ->
 
         return False
 
-    def receive(self, chunk: ModelResponseStream):
+    def receive(self, chunk: "ModelResponseStream"):
         adapter_name = settings.adapter.__class__.__name__ if settings.adapter else "ChatAdapter"
         if adapter_name not in self.adapter_identifiers:
             raise ValueError(

dspy/utils/lazy_import.py

@@ -23,6 +23,7 @@
 import inspect
 import sys
 import types
+from collections.abc import Callable
 from typing import Any
 
 
@@ -42,7 +43,30 @@ def _detect_dspy_dist() -> str:
     "weaviate": "weaviate",
     "anthropic": "anthropic",
     "numpy": "numpy",
+    "litellm": "litellm",
 }
+_ON_LOAD_CALLBACKS: dict[str, list[Callable[[types.ModuleType], None]]] = {}
+_LAZY_MODULES: set[str] = set()
+
+
+class _OnLoadLoader:
+    def __init__(self, module: str, loader):
+        self.module = module
+        self.loader = loader
+
+    def create_module(self, spec):
+        if hasattr(self.loader, "create_module"):
+            return self.loader.create_module(spec)
+        return None
+
+    def exec_module(self, module: types.ModuleType):
+        self.loader.exec_module(module)
+        _LAZY_MODULES.discard(self.module)
+        for callback in _ON_LOAD_CALLBACKS.pop(self.module, []):
+            callback(module)
+
+    def __getattr__(self, name: str):
+        return getattr(self.loader, name)
 
 
 class _MissingModule(types.ModuleType):
@@ -76,7 +100,13 @@ def is_available(module: str) -> bool:
         return False
 
 
-def require(module: str, *, extra: str | None = None, feature: str | None = None) -> Any:
+def require(
+    module: str,
+    *,
+    extra: str | None = None,
+    feature: str | None = None,
+    on_load: Callable[[types.ModuleType], None] | None = None,
+) -> Any:
     """Return a lazily-loaded module, or a stub that raises on access.
 
     Safe to call at module level:
@@ -94,9 +124,16 @@ def require(module: str, *, extra: str | None = None, feature: str | None = None
         module: Dotted module path (e.g. `"numpy"`).
         extra: Name of the dspy extra that provides this dep.
         feature: Label shown in the error (e.g. `"dspy.Embeddings"`).
+        on_load: Optional callback run once with the real module after it is loaded.
     """
     if module in sys.modules:
-        return sys.modules[module]
+        mod = sys.modules[module]
+        if on_load is not None:
+            if module in _LAZY_MODULES:
+                _ON_LOAD_CALLBACKS.setdefault(module, []).append(on_load)
+            else:
+                on_load(mod)
+        return mod
 
     spec = importlib.util.find_spec(module)
     if spec is None or spec.loader is None:
@@ -118,9 +155,14 @@ def require(module: str, *, extra: str | None = None, feature: str | None = None
         del parent
         return _MissingModule(module, message, frame_data)
 
+    if on_load is not None:
+        _ON_LOAD_CALLBACKS.setdefault(module, []).append(on_load)
+    spec.loader = _OnLoadLoader(module, spec.loader)
+
     loader = importlib.util.LazyLoader(spec.loader)
     spec.loader = loader
     mod = importlib.util.module_from_spec(spec)
     sys.modules[module] = mod
+    _LAZY_MODULES.add(module)
     loader.exec_module(mod)
     return mod

pyproject.toml

@@ -44,6 +44,7 @@ mcp = ["mcp; python_version >= '3.10'"]
 langchain = ["langchain_core"]
 optuna = ["optuna>=3.4.0"]
 numpy = ["numpy>=1.26.0"]
+litellm = ["litellm>=1.64.0"]
 dev = [
     "pytest>=6.2.5",
     "pytest-mock>=3.12.0",