diff --git a/dspy/clients/__init__.py b/dspy/clients/__init__.py index 8f314dc4ff..c371f3e06b 100644 --- a/dspy/clients/__init__.py +++ b/dspy/clients/__init__.py @@ -3,8 +3,7 @@ from pathlib import Path from typing import Any -import litellm - +from dspy.clients._litellm import get_litellm from dspy.clients.base_lm import BaseLM, inspect_history from dspy.clients.cache import Cache from dspy.clients.embedding import Embedder @@ -56,9 +55,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 +87,8 @@ 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 + litellm = get_litellm(feature="LiteLLM logging") + verbose_logger = litellm._logging.verbose_logger numeric_logging_level = getattr(logging, level) @@ -101,18 +98,18 @@ def configure_litellm_logging(level: str = "ERROR"): def enable_litellm_logging(): + litellm = get_litellm(feature="LiteLLM logging") litellm.suppress_debug_info = False + litellm._dspy_logging_configured = True configure_litellm_logging("DEBUG") def disable_litellm_logging(): + litellm = get_litellm(feature="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", diff --git a/dspy/clients/_litellm.py b/dspy/clients/_litellm.py new file mode 100644 index 0000000000..06f30f042d --- /dev/null +++ b/dspy/clients/_litellm.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +import functools +import sys +import types +from typing import Any + +from dspy.utils.lazy_import import require + + +@functools.cache +def _configure_litellm_defaults(litellm: types.ModuleType) -> None: + """Apply DSPy's global LiteLLM defaults once when LiteLLM is first imported.""" + 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 + + +def _materialize_litellm(litellm: types.ModuleType) -> None: + """Force LiteLLM's lazy module to execute, or raise the missing dependency error.""" + # `require()` returns either an importlib LazyLoader-backed module or a _MissingModule. + # Accessing a real LiteLLM attribute forces LazyLoader execution; on _MissingModule it raises + # the helpful install-hint ImportError immediately at the DSPy call site. + _completion = litellm.completion + + +@functools.cache +def get_litellm(*, feature: str) -> Any: + """Import LiteLLM, apply DSPy's defaults once, and return the module.""" + litellm = require("litellm", extra="litellm", feature=feature) + _materialize_litellm(litellm) + _configure_litellm_defaults(litellm) + return litellm + + +def is_litellm_context_window_error(error: Exception) -> bool: + """Return whether an exception is LiteLLM's context-window error, if LiteLLM is loaded.""" + litellm_module = sys.modules.get("litellm") + context_window_error = getattr(litellm_module, "ContextWindowExceededError", None) + return context_window_error is not None and isinstance(error, context_window_error) diff --git a/dspy/clients/embedding.py b/dspy/clients/embedding.py index 0f74407c44..01bfc1e641 100644 --- a/dspy/clients/embedding.py +++ b/dspy/clients/embedding.py @@ -2,13 +2,17 @@ from typing import Any, Callable -import litellm - +from dspy.clients._litellm import get_litellm from dspy.clients.cache import request_cache from dspy.utils.lazy_import import require np = require("numpy") + +def _get_litellm(): + return get_litellm(feature="dspy.Embedder") + + class Embedder: """DSPy embedding class. @@ -152,8 +156,8 @@ async def acall(self, inputs, batch_size=None, caching=None, **kwargs): def _compute_embeddings(model, batch_inputs, caching=False, **kwargs): if isinstance(model, str): - caching = caching and litellm.cache is not None - embedding_response = litellm.embedding(model=model, input=batch_inputs, caching=caching, **kwargs) + caching = caching and _get_litellm().cache is not None + embedding_response = _get_litellm().embedding(model=model, input=batch_inputs, caching=caching, **kwargs) return [data["embedding"] for data in embedding_response.data] elif callable(model): return model(batch_inputs, **kwargs) @@ -168,8 +172,8 @@ def _cached_compute_embeddings(model, batch_inputs, caching=True, **kwargs): async def _acompute_embeddings(model, batch_inputs, caching=False, **kwargs): if isinstance(model, str): - caching = caching and litellm.cache is not None - embedding_response = await litellm.aembedding(model=model, input=batch_inputs, caching=caching, **kwargs) + caching = caching and _get_litellm().cache is not None + embedding_response = await _get_litellm().aembedding(model=model, input=batch_inputs, caching=caching, **kwargs) return [data["embedding"] for data in embedding_response.data] elif callable(model): return model(batch_inputs, **kwargs) diff --git a/dspy/clients/lm.py b/dspy/clients/lm.py index 89c25ece3b..51ad677e0e 100644 --- a/dspy/clients/lm.py +++ b/dspy/clients/lm.py @@ -7,12 +7,11 @@ 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._litellm import get_litellm, is_litellm_context_window_error from dspy.clients.cache import request_cache from dspy.clients.openai import OpenAIProvider from dspy.clients.provider import Provider, ReinforceJob, TrainingJob @@ -26,6 +25,10 @@ logger = logging.getLogger(__name__) +def _get_litellm(): + return get_litellm(feature="dspy.LM") + + class LM(BaseLM): """ A language model supporting chat or text completion requests for use with DSPy modules. @@ -122,19 +125,19 @@ def _provider_name(self) -> str: @property def supports_function_calling(self) -> bool: - return litellm.supports_function_calling(model=self.model) + return _get_litellm().supports_function_calling(model=self.model) @property def supports_reasoning(self) -> bool: - return litellm.supports_reasoning(self.model) + return _get_litellm().supports_reasoning(self.model) @property def supports_response_schema(self) -> bool: - return litellm.supports_response_schema(model=self.model, custom_llm_provider=self._provider_name) + return _get_litellm().supports_response_schema(model=self.model, custom_llm_provider=self._provider_name) @property def supported_params(self) -> set[str]: - params = litellm.get_supported_openai_params(model=self.model, custom_llm_provider=self._provider_name) + params = _get_litellm().get_supported_openai_params(model=self.model, custom_llm_provider=self._provider_name) return set(params) if params else set() def _warn_zero_temp_rollout(self, temperature: float | None, rollout_id): @@ -189,8 +192,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 +235,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) @@ -355,7 +362,7 @@ def _get_stream_completion_fn( request["stream_options"] = {"include_usage": True} async def stream_completion(request: dict[str, Any], cache_kwargs: dict[str, Any]): - response = await litellm.acompletion( + response = await _get_litellm().acompletion( cache=cache_kwargs, stream=True, headers=headers, @@ -368,7 +375,7 @@ async def stream_completion(request: dict[str, Any], cache_kwargs: dict[str, Any chunk.predict_id = caller_predict_id chunks.append(chunk) await stream.send(chunk) - return litellm.stream_chunk_builder(chunks) + return _get_litellm().stream_chunk_builder(chunks) def sync_stream_completion(): return anyio.from_thread.run(functools.partial(stream_completion, request, cache_kwargs)) @@ -389,7 +396,7 @@ def litellm_completion(request: dict[str, Any], num_retries: int, cache: dict[st headers = _add_dspy_identifier_to_headers(request.pop("headers", None)) stream_completion = _get_stream_completion_fn(request, cache, sync=True, headers=headers) if stream_completion is None: - return litellm.completion( + return _get_litellm().completion( cache=cache, num_retries=num_retries, retry_strategy="exponential_backoff_retry", @@ -417,7 +424,7 @@ def litellm_text_completion(request: dict[str, Any], num_retries: int, cache: di # Build the prompt from the messages. prompt = "\n\n".join([x["content"] for x in request.pop("messages")] + ["BEGIN RESPONSE:"]) - return litellm.text_completion( + return _get_litellm().text_completion( cache=cache, model=f"text-completion-openai/{model}", api_key=api_key, @@ -437,7 +444,7 @@ async def alitellm_completion(request: dict[str, Any], num_retries: int, cache: headers = _add_dspy_identifier_to_headers(request.pop("headers", None)) stream_completion = _get_stream_completion_fn(request, cache, sync=False, headers=headers) if stream_completion is None: - return await litellm.acompletion( + return await _get_litellm().acompletion( cache=cache, num_retries=num_retries, retry_strategy="exponential_backoff_retry", @@ -463,7 +470,7 @@ async def alitellm_text_completion(request: dict[str, Any], num_retries: int, ca # Build the prompt from the messages. prompt = "\n\n".join([x["content"] for x in request.pop("messages")] + ["BEGIN RESPONSE:"]) - return await litellm.atext_completion( + return await _get_litellm().atext_completion( cache=cache, model=f"text-completion-openai/{model}", api_key=api_key, @@ -483,7 +490,7 @@ def litellm_responses_completion(request: dict[str, Any], num_retries: int, cach headers = request.pop("headers", None) request = _convert_chat_request_to_responses_request(request) - return litellm.responses( + return _get_litellm().responses( cache=cache, num_retries=num_retries, retry_strategy="exponential_backoff_retry", @@ -499,7 +506,7 @@ async def alitellm_responses_completion(request: dict[str, Any], num_retries: in headers = request.pop("headers", None) request = _convert_chat_request_to_responses_request(request) - return await litellm.aresponses( + return await _get_litellm().aresponses( cache=cache, num_retries=num_retries, retry_strategy="exponential_backoff_retry", diff --git a/dspy/streaming/streamify.py b/dspy/streaming/streamify.py index 730223f4c8..90a055baab 100644 --- a/dspy/streaming/streamify.py +++ b/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:"): diff --git a/dspy/streaming/streaming_listener.py b/dspy/streaming/streaming_listener.py index 4080f26cdd..754cbf54de 100644 --- a/dspy/streaming/streaming_listener.py +++ b/dspy/streaming/streaming_listener.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import inspect import re from collections import defaultdict @@ -5,7 +7,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 +16,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] @@ -368,7 +371,7 @@ def _output_type(self) -> type | None: def find_predictor_for_stream_listeners( - program: "Module", stream_listeners: list[StreamListener] + program: Module, stream_listeners: list[StreamListener] ) -> dict[int, list[StreamListener]]: """Find the predictor for each stream listener. @@ -385,7 +388,7 @@ def find_predictor_for_stream_listeners( field_name_to_named_predictor[listener.signature_field_name] = None for name, predictor in predictors: - for field_name, field_info in predictor.signature.output_fields.items(): + for field_name in predictor.signature.output_fields: if field_name not in field_name_to_named_predictor: continue diff --git a/dspy/utils/lazy_import.py b/dspy/utils/lazy_import.py index 1090011b84..1cebf13a98 100644 --- a/dspy/utils/lazy_import.py +++ b/dspy/utils/lazy_import.py @@ -42,6 +42,7 @@ def _detect_dspy_dist() -> str: "weaviate": "weaviate", "anthropic": "anthropic", "numpy": "numpy", + "litellm": "litellm", } diff --git a/pyproject.toml b/pyproject.toml index 0874bea76f..40ea542087 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/clients/test_lazy_litellm_import.py b/tests/clients/test_lazy_litellm_import.py new file mode 100644 index 0000000000..5eeec4092f --- /dev/null +++ b/tests/clients/test_lazy_litellm_import.py @@ -0,0 +1,58 @@ +import importlib.util +import sys + +import pytest + + +def _hide_litellm(monkeypatch): + real_find_spec = importlib.util.find_spec + + def find_spec(name, *args, **kwargs): + if name == "litellm" or name.startswith("litellm."): + return None + return real_find_spec(name, *args, **kwargs) + + monkeypatch.setattr(importlib.util, "find_spec", find_spec) + monkeypatch.delitem(sys.modules, "litellm", raising=False) + + from dspy.clients._litellm import get_litellm + + get_litellm.cache_clear() + + +def test_import_dspy_does_not_import_litellm(monkeypatch): + monkeypatch.delitem(sys.modules, "litellm", raising=False) + + import dspy + + _ = dspy.LM + _ = dspy.Embedder + _ = dspy.streamify + + assert "litellm" not in sys.modules + + +def test_lm_litellm_use_raises_helpful_error_without_litellm(monkeypatch): + import dspy + + _hide_litellm(monkeypatch) + + with pytest.raises(ImportError) as exc_info: + _ = dspy.LM("openai/gpt-4o-mini").supports_function_calling + + msg = str(exc_info.value) + assert "[litellm]" in msg + assert "dspy.LM" in msg + + +def test_embedder_litellm_use_raises_helpful_error_without_litellm(monkeypatch): + import dspy + + _hide_litellm(monkeypatch) + + with pytest.raises(ImportError) as exc_info: + dspy.Embedder("openai/text-embedding-3-small")(["hello"]) + + msg = str(exc_info.value) + assert "[litellm]" in msg + assert "dspy.Embedder" in msg diff --git a/tests/utils/test_lazy_import.py b/tests/utils/test_lazy_import.py index 49a718058c..85495641b9 100644 --- a/tests/utils/test_lazy_import.py +++ b/tests/utils/test_lazy_import.py @@ -46,17 +46,32 @@ def test_require_stub_raises_on_access_with_install_hint(): dist = _detect_dspy_dist() stub = require("nonexistent_abc", feature="dspy.Test") with pytest.raises(ImportError) as exc_info: - stub.something + _ = stub.something msg = str(exc_info.value) assert f"{dist}[nonexistent_abc]" in msg, msg assert "dspy.Test" in msg +def test_require_stub_uses_install_hint_for_litellm(monkeypatch): + import importlib.util + import sys + + dist = _detect_dspy_dist() + find_spec = importlib.util.find_spec + monkeypatch.delitem(sys.modules, "litellm", raising=False) + monkeypatch.setattr(importlib.util, "find_spec", lambda module: None if module == "litellm" else find_spec(module)) + + stub = require("litellm", feature="dspy.LM") + with pytest.raises(ImportError) as exc_info: + _ = stub.something + assert f"{dist}[litellm]" in str(exc_info.value) + + def test_require_stub_uses_explicit_extra(): dist = _detect_dspy_dist() stub = require("nonexistent_xyz", extra="custom", feature="dspy.X") with pytest.raises(ImportError) as exc_info: - stub.something + _ = stub.something assert f"{dist}[custom]" in str(exc_info.value) @@ -64,7 +79,7 @@ def test_require_stub_falls_back_to_module_name(): dist = _detect_dspy_dist() stub = require("nonexistent_xyz", feature="dspy.X") with pytest.raises(ImportError) as exc_info: - stub.something + _ = stub.something assert f"{dist}[nonexistent_xyz]" in str(exc_info.value) diff --git a/uv.lock b/uv.lock index c9d87f4ba4..a8d35588d4 100644 --- a/uv.lock +++ b/uv.lock @@ -778,6 +778,10 @@ dev = [ langchain = [ { name = "langchain-core" }, ] +litellm = [ + { name = "litellm", version = "1.68.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.14' and sys_platform != 'win32'" }, + { name = "litellm", version = "1.72.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14' or sys_platform == 'win32'" }, +] mcp = [ { name = "mcp", version = "1.5.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.14' and sys_platform != 'win32'" }, { name = "mcp", version = "1.9.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14' or sys_platform == 'win32'" }, @@ -819,6 +823,7 @@ requires-dist = [ { name = "langchain-core", marker = "extra == 'test-extras'" }, { name = "litellm", specifier = ">=1.64.0" }, { name = "litellm", marker = "(python_full_version == '3.14.*' and extra == 'dev') or (sys_platform == 'win32' and extra == 'dev')", specifier = ">=1.64.0" }, + { name = "litellm", marker = "extra == 'litellm'", specifier = ">=1.64.0" }, { name = "litellm", extras = ["proxy"], marker = "python_full_version < '3.14' and sys_platform != 'win32' and extra == 'dev'", specifier = ">=1.64.0" }, { name = "mcp", marker = "python_full_version >= '3.10' and extra == 'mcp'" }, { name = "mcp", marker = "python_full_version >= '3.10' and extra == 'test-extras'" }, @@ -843,7 +848,7 @@ requires-dist = [ { name = "tqdm", specifier = ">=4.66.1" }, { name = "weaviate-client", marker = "extra == 'weaviate'", specifier = "~=4.5.4" }, ] -provides-extras = ["anthropic", "weaviate", "mcp", "langchain", "optuna", "numpy", "dev", "test-extras"] +provides-extras = ["anthropic", "weaviate", "mcp", "langchain", "optuna", "numpy", "litellm", "dev", "test-extras"] [[package]] name = "email-validator"