From 28a56068a6a7975f446dcb8018a21f8d5377693d Mon Sep 17 00:00:00 2001 From: Maxime Rivest Date: Sun, 24 May 2026 07:58:11 -0400 Subject: [PATCH] Move LM capabilities into BaseLM --- dspy/__init__.py | 18 +++- dspy/adapters/_type_runtime.py | 13 +-- dspy/adapters/base.py | 32 ++++--- dspy/clients/__init__.py | 7 +- dspy/clients/base_lm.py | 88 +++++++++++++++++--- dspy/clients/lm.py | 12 ++- dspy/predict/predict.py | 57 ++++++++++++- dspy/utils/exceptions.py | 147 +++++++++++++++++++++++++++++---- 8 files changed, 314 insertions(+), 60 deletions(-) diff --git a/dspy/__init__.py b/dspy/__init__.py index cc9d1fe5e8..56fc761cf6 100644 --- a/dspy/__init__.py +++ b/dspy/__init__.py @@ -8,7 +8,23 @@ from dspy.clients import * # isort: skip from dspy.adapters import Adapter, ChatAdapter, JSONAdapter, XMLAdapter, TwoStepAdapter, Image, Audio, File, History, Type, Tool, ToolCalls, ToolCallResults, Code, Reasoning # isort: skip from dspy.primitives.sandbox_serializable import SandboxSerializable # isort: skip -from dspy.utils.exceptions import ContextWindowExceededError +from dspy.utils.exceptions import ( + ContextWindowExceededError, + DSPyError, + LMAuthError, + LMBillingError, + LMConfigurationError, + LMError, + LMInvalidRequestError, + LMNotConfiguredError, + LMProviderError, + LMRateLimitError, + LMServerError, + LMTimeoutError, + LMTransportError, + LMUnsupportedFeatureError, + LMUnsupportedModelError, +) from dspy.utils.logging_utils import configure_dspy_loggers, disable_logging, enable_logging from dspy.utils.asyncify import asyncify from dspy.utils.syncify import syncify diff --git a/dspy/adapters/_type_runtime.py b/dspy/adapters/_type_runtime.py index 757aa3ae22..9c5e5d1985 100644 --- a/dspy/adapters/_type_runtime.py +++ b/dspy/adapters/_type_runtime.py @@ -10,26 +10,17 @@ if TYPE_CHECKING: from dspy.adapters.base import Adapter + from dspy.clients.base_lm import BaseLM from dspy.signatures.signature import Signature -@dataclass(frozen=True) -class _LMCapabilities: - model: str - model_type: str - supported_params: frozenset[str] - supports_function_calling: bool - supports_response_schema: bool - supports_reasoning: bool - - @dataclass(frozen=True) class _CallContext: adapter: Adapter use_native_function_calling: bool allow_parallel_tool_calls: bool | None native_response_types: tuple[type[object], ...] - lm: _LMCapabilities + lm: BaseLM lm_kwargs: dict[str, object] = field(default_factory=dict) lm_default_kwargs: dict[str, object] = field(default_factory=dict) diff --git a/dspy/adapters/base.py b/dspy/adapters/base.py index 599067981d..86d3eeae01 100644 --- a/dspy/adapters/base.py +++ b/dspy/adapters/base.py @@ -19,7 +19,6 @@ from dspy.adapters._type_runtime import ( _AdapterCallPlan, _CallContext, - _LMCapabilities, _merge_lm_config, _OutputParser, _TypeFeatureHandler, @@ -30,7 +29,7 @@ from dspy.adapters.types.reasoning import Reasoning from dspy.adapters.types.tool import ToolCallResults, ToolCalls from dspy.adapters.utils import format_field_value, parse_value -from dspy.clients.base_lm import BaseLM +from dspy.clients.base_lm import BaseLM, LMCapabilities from dspy.clients.openai_format import ( lm_response_from_legacy_outputs, message_to_openai_chat, @@ -59,6 +58,17 @@ def signature(self) -> type[Signature]: return self.call_plan.render_signature +class _AdapterContextLM(BaseLM): + """Minimal BaseLM used when rendering adapter messages without a real LM.""" + + def __init__(self, *, use_native_tool_calls: bool = False): + super().__init__(model="", model_type="chat", temperature=None, max_tokens=None) + self._capabilities = LMCapabilities(function_calling=use_native_tool_calls) + + def get_capabilities(self) -> LMCapabilities: + return self._capabilities + + class Adapter: """Base Adapter class. @@ -116,14 +126,7 @@ def build_call_context(self, lm: BaseLM, lm_kwargs: dict[str, Any] | None = None use_native_function_calling=self.use_native_function_calling, allow_parallel_tool_calls=self.allow_parallel_tool_calls, native_response_types=tuple(self.native_response_types), - lm=_LMCapabilities( - model=lm.model, - model_type=getattr(lm, "model_type", "chat"), - supported_params=frozenset(getattr(lm, "supported_params", set())), - supports_function_calling=bool(getattr(lm, "supports_function_calling", False)), - supports_response_schema=bool(getattr(lm, "supports_response_schema", False)), - supports_reasoning=bool(getattr(lm, "supports_reasoning", False)), - ), + lm=lm, lm_kwargs=dict(lm_kwargs or {}), lm_default_kwargs=dict(getattr(lm, "kwargs", {}) or {}), ) @@ -134,14 +137,7 @@ def _default_call_context(self, *, use_native_tool_calls: bool = False) -> _Call use_native_function_calling=use_native_tool_calls, allow_parallel_tool_calls=self.allow_parallel_tool_calls, native_response_types=tuple(self.native_response_types), - lm=_LMCapabilities( - model="", - model_type="chat", - supported_params=frozenset(), - supports_function_calling=use_native_tool_calls, - supports_response_schema=False, - supports_reasoning=False, - ), + lm=_AdapterContextLM(use_native_tool_calls=use_native_tool_calls), ) def value_to_lm_parts(self, value: object, field_info: FieldInfo) -> list[LMPart]: diff --git a/dspy/clients/__init__.py b/dspy/clients/__init__.py index c371f3e06b..be5240e20f 100644 --- a/dspy/clients/__init__.py +++ b/dspy/clients/__init__.py @@ -4,7 +4,7 @@ from typing import Any from dspy.clients._litellm import get_litellm -from dspy.clients.base_lm import BaseLM, inspect_history +from dspy.clients.base_lm import BaseLM, LMCapabilities, inspect_history from dspy.clients.cache import Cache from dspy.clients.embedding import Embedder from dspy.clients.lm import LM @@ -39,7 +39,7 @@ def configure_cache( safe_types: Additional types to allow when restrict_pickle is True. """ - DSPY_CACHE = Cache( + dspy_cache = Cache( enable_disk_cache, enable_memory_cache, disk_cache_dir, @@ -52,7 +52,7 @@ def configure_cache( import dspy # Update the reference to point to the new cache - dspy.cache = DSPY_CACHE + dspy.cache = dspy_cache @@ -112,6 +112,7 @@ def disable_litellm_logging(): __all__ = [ "BaseLM", + "LMCapabilities", "LM", "Provider", "TrainingJob", diff --git a/dspy/clients/base_lm.py b/dspy/clients/base_lm.py index 760a2679b5..2acb49e3b7 100644 --- a/dspy/clients/base_lm.py +++ b/dspy/clients/base_lm.py @@ -1,7 +1,11 @@ +import copy as copy_module import datetime import uuid +from dataclasses import dataclass, field from typing import Any, TextIO +from typing_extensions import Self + from dspy.dsp.utils import settings from dspy.utils.callback import with_callbacks from dspy.utils.inspect_history import pretty_print_history @@ -10,6 +14,27 @@ GLOBAL_HISTORY = [] +@dataclass(frozen=True) +class LMCapabilities: + """Optional model and deployment metadata for an LM backend. + + Capabilities are descriptive hints. Adapters can use them to select native + paths, but concrete LM implementations still decide how to handle requests. + """ + + function_calling: bool = False + reasoning: bool = False + response_schema: bool = False + streaming: bool = False + input_image: bool = False + input_audio: bool = False + input_file: bool = False + output_image: bool = False + output_audio: bool = False + tool_results: bool = False + extensions: dict[str, Any] = field(default_factory=dict) + + class BaseLM: """Base class for handling LLM calls. @@ -60,32 +85,55 @@ def forward(self, prompt, messages=None, **kwargs): ``` """ - def __init__(self, model, model_type="chat", temperature=0.0, max_tokens=1000, cache=True, **kwargs): + def __init__( + self, + model, + model_type="chat", + temperature=0.0, + max_tokens=1000, + cache=True, + callbacks=None, + num_retries=0, + **kwargs, + ): self.model = model self.model_type = model_type self.cache = cache + self.callbacks = callbacks or [] + self.num_retries = num_retries self.kwargs = dict(temperature=temperature, max_tokens=max_tokens, **kwargs) self.history = [] + self._warned_zero_temp_rollout = False + + @property + def capabilities(self) -> LMCapabilities: + """Native metadata available for this model instance.""" + return self.get_capabilities() + + def get_capabilities(self) -> LMCapabilities: + """Return optional native model and deployment hints.""" + return LMCapabilities() @property def supports_function_calling(self) -> bool: """Whether the model supports function calling (tool use).""" - return False + return self.capabilities.function_calling @property def supports_reasoning(self) -> bool: """Whether the model supports native reasoning (extended thinking).""" - return False + return self.capabilities.reasoning @property def supports_response_schema(self) -> bool: """Whether the model supports structured output via response schema.""" - return False + return self.capabilities.response_schema @property def supported_params(self) -> set[str]: """Set of supported OpenAI-style parameter names for the model.""" - return set() + supported = self.capabilities.extensions.get("supported_params", set()) + return set(supported) if supported else set() def _process_lm_response(self, response, prompt, messages, **kwargs): merged_kwargs = {**self.kwargs, **kwargs} @@ -190,23 +238,43 @@ async def aforward( """ raise NotImplementedError("Subclasses must implement this method.") + def dump_state(self) -> dict[str, Any]: + """Return a sanitized reconstruction state for this LM.""" + filtered_kwargs = {key: value for key, value in self.kwargs.items() if key != "api_key"} + return { + "model": self.model, + "model_type": self.model_type, + "cache": self.cache, + "num_retries": self.num_retries, + **filtered_kwargs, + } + + @classmethod + def load_state(cls, state: dict[str, Any]) -> Self: + """Reconstruct this LM from `dump_state()` output.""" + return cls(**state) + def copy(self, **kwargs): """Returns a copy of the language model with possibly updated parameters. Any provided keyword arguments update the corresponding attributes or LM kwargs of the copy. For example, ``lm.copy(rollout_id=1, temperature=1.0)`` returns an LM whose requests use a different rollout ID at non-zero temperature to bypass cache collisions. - """ - import copy + The default implementation uses a shallow runtime copy so provider clients, sessions, + and local model handles are preserved by reference while DSPy-owned mutable state is + isolated on the copy. + """ - new_instance = copy.deepcopy(self) + new_instance = copy_module.copy(self) new_instance.history = [] + new_instance.callbacks = list(getattr(self, "callbacks", []) or []) + new_instance.kwargs = dict(getattr(self, "kwargs", {}) or {}) for key, value in kwargs.items(): - if hasattr(self, key): + if hasattr(new_instance, key): setattr(new_instance, key, value) - if (key in self.kwargs) or (not hasattr(self, key)): + if (key in new_instance.kwargs) or (not hasattr(self, key)): if value is None: new_instance.kwargs.pop(key, None) else: diff --git a/dspy/clients/lm.py b/dspy/clients/lm.py index 51ad677e0e..1405a745c9 100644 --- a/dspy/clients/lm.py +++ b/dspy/clients/lm.py @@ -20,7 +20,7 @@ from dspy.utils.callback import BaseCallback from dspy.utils.exceptions import ContextWindowExceededError -from .base_lm import BaseLM +from .base_lm import BaseLM, LMCapabilities logger = logging.getLogger(__name__) @@ -123,6 +123,16 @@ def _provider_name(self) -> str: return self.model.split("/", 1)[0] return "openai" + def get_capabilities(self) -> LMCapabilities: + params = self.supported_params + return LMCapabilities( + function_calling=self.supports_function_calling, + reasoning=self.supports_reasoning, + response_schema=self.supports_response_schema, + streaming=True, + extensions={"supported_params": params}, + ) + @property def supports_function_calling(self) -> bool: return _get_litellm().supports_function_calling(model=self.model) diff --git a/dspy/predict/predict.py b/dspy/predict/predict.py index 8b10493156..6212806bf7 100644 --- a/dspy/predict/predict.py +++ b/dspy/predict/predict.py @@ -1,3 +1,4 @@ +import importlib import logging import random import types @@ -21,6 +22,7 @@ logger = logging.getLogger(__name__) UNSAFE_LM_STATE_KEYS = {"api_base", "base_url", "model_list"} +_LM_CLASS_STATE_KEY = "_dspy_lm_class" def _sanitize_lm_state(lm_state: dict, allow_unsafe_lm_state: bool) -> dict: @@ -99,7 +101,7 @@ def dump_state(self, json_mode=True): state["demos"].append(demo.toDict()) state["signature"] = self.signature.dump_state() - state["lm"] = self.lm.dump_state() if self.lm else None + state["lm"] = _dump_lm_state(self.lm) if self.lm else None return state def load_state(self, state: dict, *, allow_unsafe_lm_state: bool = False) -> "Predict": @@ -121,7 +123,7 @@ def load_state(self, state: dict, *, allow_unsafe_lm_state: bool = False) -> "Pr self.signature = self.signature.load_state(state["signature"]) sanitized_lm_state = _sanitize_lm_state(state["lm"], allow_unsafe_lm_state) if state["lm"] else None - self.lm = LM(**sanitized_lm_state) if sanitized_lm_state else None + self.lm = _load_lm_from_state(sanitized_lm_state) if sanitized_lm_state else None if "extended_signature" in state: # legacy, up to and including 2.5, for CoT. raise NotImplementedError("Loading extended_signature is no longer supported in DSPy 2.6+") @@ -291,6 +293,57 @@ def get_config(self): def __repr__(self): return f"{self.__class__.__name__}({self.signature})" + +def _dump_lm_state(lm: BaseLM) -> dict[str, Any]: + """Return LM state with enough type information to reconstruct it.""" + lm_state = dict(lm.dump_state()) + lm_state.setdefault(_LM_CLASS_STATE_KEY, f"{type(lm).__module__}.{type(lm).__qualname__}") + return lm_state + + +def _load_lm_from_state(lm_state: dict[str, Any]) -> BaseLM: + """Reconstruct a legacy or custom BaseLM from serialized state. + + Older saved programs did not record the concrete LM class, so they are + loaded through the legacy LiteLLM-backed `dspy.clients.lm.LM` constructor. + Newer states record the concrete class so `Predict` can round-trip any + importable `BaseLM` subclass. + """ + lm_state = dict(lm_state) + class_path = lm_state.pop(_LM_CLASS_STATE_KEY, None) + if class_path is None: + return LM(**lm_state) + + lm_cls = _import_lm_class(class_path) + if not issubclass(lm_cls, BaseLM): + raise TypeError(f"Serialized LM class `{class_path}` must be a subclass of dspy.BaseLM.") + + load_state = getattr(lm_cls, "load_state", None) + if callable(load_state): + return load_state(lm_state) + return lm_cls(**lm_state) + + +def _import_lm_class(class_path: str) -> type: + try: + module_name, class_name = class_path.rsplit(".", 1) + except ValueError as exc: + raise ValueError(f"Invalid serialized LM class path: {class_path!r}") from exc + + try: + module = importlib.import_module(module_name) + lm_cls = getattr(module, class_name) + except (ImportError, AttributeError) as exc: + raise ImportError( + f"Could not import serialized LM class `{class_path}`. Ensure the class is importable, " + "or load an older state without the concrete LM class marker." + ) from exc + + if not isinstance(lm_cls, type): + raise TypeError(f"Serialized LM class `{class_path}` did not resolve to a class.") + return lm_cls + + def _get_type_name(type_annotation) -> str: """Helper method to get the name for a type annotation.""" diff --git a/dspy/utils/exceptions.py b/dspy/utils/exceptions.py index 928bbe351b..dbfd79ae1c 100644 --- a/dspy/utils/exceptions.py +++ b/dspy/utils/exceptions.py @@ -1,25 +1,144 @@ +from __future__ import annotations from dspy.signatures.signature import Signature -class ContextWindowExceededError(Exception): - """Raised when the prompt exceeds the model's context window. +class DSPyError(Exception): + """Base class for DSPy errors with structured metadata.""" - Any `BaseLM` subclass should raise this error (or a subclass of it) when the - request fails because the input is too long for the model. Adapters and some - modules rely on catching this specific type to decide whether a fallback - retry is appropriate. + default_code: str | None = None - Args: - model: The model identifier that rejected the request. - message: Description of the error. Defaults to `"Context window exceeded"`. - """ - - def __init__(self, *, model: str | None = None, message: str = "Context window exceeded"): + def __init__( + self, + message: str = "", + *, + code: str | None = None, + model: str | None = None, + provider: str | None = None, + provider_code: str | None = None, + status: int | None = None, + request_id: str | None = None, + retry_after: float | None = None, + ): + self.message = message + self.code = code or self.default_code self.model = model - msg = message + self.provider = provider + self.provider_code = provider_code + self.status = status + self.request_id = request_id + self.retry_after = retry_after + prefix = f"[{model}] " if model else "" - super().__init__(f"{prefix}{msg}") + super().__init__(f"{prefix}{message}" if message else prefix.rstrip()) + + +class LMError(DSPyError): + """Base class for language model errors.""" + + default_code = "lm_error" + + +class LMTransportError(LMError): + """The LM request failed before the provider returned a response.""" + + default_code = "transport" + + +class LMConfigurationError(LMError): + """The LM or provider client is not configured correctly.""" + + default_code = "configuration" + + +class LMNotConfiguredError(LMConfigurationError): + """The LM is missing required provider configuration or credentials.""" + + default_code = "not_configured" + + +class LMUnsupportedFeatureError(LMError): + """The LM does not support a requested feature.""" + + default_code = "unsupported_feature" + + def __init__( + self, + message: str = "", + *, + features: list[str] | None = None, + issues: list[str] | None = None, + **kwargs, + ): + self.features = list(features or []) + self.issues = list(issues or []) + super().__init__(message, **kwargs) + + +class LMProviderError(LMError): + """The provider returned an error response.""" + + default_code = "provider" + + +class LMAuthError(LMProviderError): + """The provider rejected the request because authentication failed.""" + + default_code = "auth" + + +class LMBillingError(LMProviderError): + """The provider rejected the request because billing or quota failed.""" + + default_code = "billing" + + +class LMRateLimitError(LMProviderError): + """The provider rate-limited the request.""" + + default_code = "rate_limit" + + +class LMInvalidRequestError(LMProviderError): + """The provider rejected the request shape or resource.""" + + default_code = "invalid_request" + + +class ContextWindowExceededError(LMInvalidRequestError): + """Raised when the prompt exceeds the model's context window.""" + + default_code = "context_window_exceeded" + + def __init__( + self, + *, + model: str | None = None, + message: str = "Context window exceeded", + **kwargs, + ): + super().__init__(message, model=model, **kwargs) + + +class LMUnsupportedModelError(LMInvalidRequestError): + """The requested model is unavailable or unsupported by the provider.""" + + default_code = "unsupported_model" + + +class LMTimeoutError(LMProviderError): + """The provider request timed out.""" + + default_code = "timeout" + + +class LMServerError(LMProviderError): + """The provider failed while handling the request.""" + + default_code = "server" + + +RETRYABLE_LM_ERRORS = (LMRateLimitError, LMTimeoutError, LMServerError, LMTransportError) class AdapterParseError(Exception):