Fix circular import and serialization issues.

Author: FireMasterK

dspy/predict/chain_of_thought.py

@@ -1,9 +1,9 @@
+from functools import cached_property
 from typing import Any, Optional, Type, Union
 
 from pydantic.fields import FieldInfo
 
 import dspy
-from dspy.clients.base_lm import BaseLM
 from dspy.primitives.module import Module
 from dspy.signatures.field import OutputField
 from dspy.signatures.signature import Signature, ensure_signature
@@ -27,43 +27,60 @@ def __init__(
             **config: The configuration for the module.
         """
         super().__init__()
-        signature = ensure_signature(signature)
-        self.predict_reasoning = dspy.Predict(signature, **config)
-        prefix = "Reasoning: Let's think step by step in order to"
-        desc = "${reasoning}"
-        rationale_field_type = rationale_field.annotation if rationale_field else rationale_field_type
-        rationale_field = rationale_field if rationale_field else dspy.OutputField(prefix=prefix, desc=desc)
-        extended_signature = signature.prepend(name="reasoning", field=rationale_field, type_=rationale_field_type)
-        self.predict = dspy.Predict(extended_signature, **config)
+        self._signature = ensure_signature(signature)
+        self._config = config
+        self._rationale_field = rationale_field
+        self._rationale_field_type = rationale_field_type
 
-    def _validate_lm(self, **kwargs):
-        """Helper method to validate the LM configuration."""
-        lm = kwargs.get("lm", dspy.settings.lm)
-
-        if lm is None:
-            raise ValueError(
-                "No LM is loaded. Please configure the LM using `dspy.configure(lm=dspy.LM(...))`. e.g, "
-                "`dspy.configure(lm=dspy.LM('openai/gpt-4o-mini'))`"
+    @cached_property
+    def predict(self):
+        """Returns the appropriate predict instance based on the LM's reasoning model capability."""
+        lm = dspy.settings.lm
+        if lm and getattr(lm, "reasoning_model", False):
+            return dspy.Predict(self._signature, **self._config)
+        else:
+            prefix = "Reasoning: Let's think step by step in order to"
+            desc = "${reasoning}"
+            rationale_field_type = (
+                self._rationale_field.annotation if self._rationale_field else self._rationale_field_type
             )
-
-        if isinstance(lm, str):
-            # Many users mistakenly use `dspy.configure(lm="openai/gpt-4o-mini")` instead of
-            # `dspy.configure(lm=dspy.LM("openai/gpt-4o-mini"))`, so we are providing a specific error message.
-            raise ValueError(
-                f"LM must be an instance of `dspy.BaseLM`, not a string. Instead of using a string like "
-                f"'dspy.configure(lm=\"{lm}\")', please configure the LM like 'dspy.configure(lm=dspy.LM(\"{lm}\"))'"
+            rationale_field = (
+                self._rationale_field if self._rationale_field else dspy.OutputField(prefix=prefix, desc=desc)
             )
-        elif not isinstance(lm, BaseLM):
-            raise ValueError(f"LM must be an instance of `dspy.BaseLM`, not {type(lm)}. Received `lm={lm}`.")
-
-        return lm
+            extended_signature = self._signature.prepend(
+                name="reasoning", field=rationale_field, type_=rationale_field_type
+            )
+            return dspy.Predict(extended_signature, **self._config)
 
     def forward(self, **kwargs):
-        lm = self._validate_lm(**kwargs)
-        return self.predict(**kwargs) if not lm.reasoning_model else self.predict_reasoning(**kwargs)
+        return self.predict(**kwargs)
 
     async def aforward(self, **kwargs):
-        lm = self._validate_lm(**kwargs)
-        return await (
-            self.predict.acall(**kwargs) if not lm.reasoning_model else self.predict_reasoning.acall(**kwargs)
-        )
+        return await self.predict.acall(**kwargs)
+
+    def load_state(self, state):
+        """Override to ensure predict parameter is created before loading state."""
+        # If predict state exists but predict hasn't been accessed yet, access it first
+        if "predict" in state and "predict" not in self.__dict__:
+            _ = self.predict  # This creates the predict instance
+
+        # Now call the base load_state which will load into all named_parameters
+        return super().load_state(state)
+
+    def __setstate__(self, state):
+        """Custom deserialization for cloudpickle to preserve predict instance."""
+        # Restore the state normally
+        self.__dict__.update(state)
+
+        # If predict was cached and serialized, we don't need to do anything special
+        # since cloudpickle should have preserved it correctly
+
+    def __getstate__(self):
+        """Custom serialization for cloudpickle to ensure predict instance is preserved."""
+        state = self.__dict__.copy()
+        # Force evaluation of cached property if not already done
+        if "predict" not in state:
+            # Access the predict property to cache it before serialization
+            _ = self.predict
+            state = self.__dict__.copy()
+        return state

tests/predict/test_chain_of_thought.py

@@ -31,7 +31,7 @@ def test_cot_skips_with_reasoning_model():
     dspy.settings.configure(lm=lm)
     signature = dspy.Signature("question -> answer")
     predict = ChainOfThought(signature)
-    assert list(predict.predict_reasoning.signature.output_fields.keys()) == [
+    assert list(predict.predict.signature.output_fields.keys()) == [
         "answer",
     ]
     assert predict(question="What is 1+1?").answer == "2"