Move LM capabilities into BaseLM - #60
Conversation
Greptile SummaryThis PR consolidates LM capability metadata into a structured
Confidence Score: 3/5The core capability-metadata and error-hierarchy changes are solid, but the generic LM deserialization path introduces an unguarded importlib.import_module call driven by attacker-controllable state-file content, and the new capabilities property is inconsistent with the documented subclassing pattern. The generic state-loading path runs importlib.import_module on a string read directly from a saved JSON file without any sanitization by sanitize_lm_state. Loading a crafted state file executes the imported module's top-level code before the issubclass guard fires. The stale BaseLM docstring example teaches subclass authors to override supports* properties, but the new capabilities property will silently return incorrect data for those subclasses, making lm.capabilities.function_calling and lm.supports_function_calling disagree on the same object. dspy/predict/predict.py (the _import_lm_class / _load_lm_from_state block) and dspy/clients/base_lm.py (the class-level docstring example and the LMCapabilities.extensions field).
|
| Filename | Overview |
|---|---|
| dspy/clients/base_lm.py | Introduces LMCapabilities dataclass and promotes LM capability metadata from ad-hoc attributes to a proper interface; docstring example remains on the old override pattern causing capabilities vs supports_* inconsistency for custom subclasses. |
| dspy/predict/predict.py | Adds _dump_lm_state/_load_lm_from_state to round-trip any BaseLM subclass; the _LM_CLASS_STATE_KEY value is not screened by _sanitize_lm_state, allowing arbitrary module imports from crafted state files. |
| dspy/clients/lm.py | Adds get_capabilities() override that delegates to existing supports_* litellm properties; straightforward and correct. |
| dspy/adapters/base.py | Removes the intermediate _LMCapabilities struct and passes the real BaseLM directly into _CallContext; cleaner and functionally equivalent. |
| dspy/utils/exceptions.py | Introduces a full LM error hierarchy with structured metadata fields; ContextWindowExceededError is correctly integrated as a leaf. |
Class Diagram
%%{init: {'theme': 'neutral'}}%%
classDiagram
class BaseLM {
+model: str
+capabilities: LMCapabilities
+get_capabilities() LMCapabilities
+supports_function_calling: bool
+dump_state() dict
+load_state(state) Self
}
class LM {
+get_capabilities() LMCapabilities
+supports_function_calling: bool
}
class _AdapterContextLM {
+get_capabilities() LMCapabilities
}
class LMCapabilities {
+function_calling: bool
+reasoning: bool
+response_schema: bool
+extensions: dict
}
class _CallContext {
+lm: BaseLM
}
class DSPyError
class LMError
class LMProviderError
class ContextWindowExceededError
BaseLM <|-- LM
BaseLM <|-- _AdapterContextLM
BaseLM --> LMCapabilities : get_capabilities()
_CallContext --> BaseLM : lm
DSPyError <|-- LMError
LMError <|-- LMProviderError
LMProviderError <|-- ContextWindowExceededError
Comments Outside Diff (1)
-
dspy/clients/base_lm.py, line 54-84 (link)Stale docstring creates inconsistency between
capabilitiesandsupports_*propertiesThe docstring example teaches users to override
supports_function_calling,supports_reasoning, etc. as individual properties. With the new design,BaseLM.supports_function_callingnow delegates toself.capabilities.function_calling(line 120). A user subclass that follows the docstring — overridingsupports_function_callingwithout overridingget_capabilities()— will silently produce inconsistent state:lm.supports_function_callingreturnsTruewhilelm.capabilities.function_callingreturnsFalse. Any code path (including user code) that readslm.capabilities.function_callingdirectly will miss the override. The docstring example should be updated to showget_capabilities()as the override point instead.
Reviews (1): Last reviewed commit: "Move LM capabilities into BaseLM" | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| 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 |
There was a problem hiding this comment.
_LM_CLASS_STATE_KEY is not screened by _sanitize_lm_state, enabling arbitrary module import
_sanitize_lm_state only removes keys from UNSAFE_LM_STATE_KEYS (api_base, base_url, model_list). The _dspy_lm_class field is left untouched and is fed directly into importlib.import_module(module_name) in _import_lm_class. Module import executes module-level code, so a crafted state file with _dspy_lm_class: "malicious_package.EvilLM" would run malicious_package's __init__.py even before the issubclass(lm_cls, BaseLM) guard is reached. Users relying on allow_unsafe_lm_state=False (the default) would reasonably expect protection here — _dspy_lm_class should either be added to UNSAFE_LM_STATE_KEYS or validated against an allowlist before import.
| output_image: bool = False | ||
| output_audio: bool = False | ||
| tool_results: bool = False | ||
| extensions: dict[str, Any] = field(default_factory=dict) |
There was a problem hiding this comment.
extensions dict inside a frozen=True dataclass
LMCapabilities is declared frozen=True, which prevents re-assigning the extensions field reference, but the dict itself remains mutable. Code that holds a reference to a LMCapabilities object and later calls capabilities.extensions["supported_params"] = new_set would silently mutate the shared state. Since get_capabilities() in LM constructs a fresh LMCapabilities on every call there is no shared state today, but any subclass that caches the result (the natural optimisation) could be affected. Switching extensions to types.MappingProxyType would make the immutability guarantee consistent.
1 participant
No description provided.