Repository: stanfordnlp/dspy
Pull request #1943 (2024-12-16 19:10:09)

Simplify built-in modules (remove new_* & extended_signature). Prepare for assertions v2.


Patch:
diff --git a/dspy/predict/__init__.py b/dspy/predict/__init__.py
index 0260ead3..e273a3dc 100644
--- a/dspy/predict/__init__.py
+++ b/dspy/predict/__init__.py
@@ -6,5 +6,5 @@ from .multi_chain_comparison import MultiChainComparison
 from .predict import Predict
 from .program_of_thought import ProgramOfThought
 from .react import ReAct, Tool
-from .retry import Retry
-from .parallel import Parallel
\ No newline at end of file
+from .parallel import Parallel
+# from .retry import Retry
\ No newline at end of file
diff --git a/dspy/predict/chain_of_thought.py b/dspy/predict/chain_of_thought.py
index b6b9cb57..4a56c0ce 100644
--- a/dspy/predict/chain_of_thought.py
+++ b/dspy/predict/chain_of_thought.py
@@ -2,58 +2,19 @@ import dspy
 from dspy.primitives.program import Module
 from dspy.signatures.signature import ensure_signature
 
-# TODO: This shouldn't inherit from Predict. It should be a module that has one or two predictors.
-# Let's focus on the activated case. It's a predictor with the expanded signature.
-# Now, when deactivated, it's a predictor with the original signature.
-# When activate is None, though, we need the expanded one but during forward we need to pass the right signature.
-
 
 class ChainOfThought(Module):
-    def __init__(self, signature, rationale_type=None, activated=True, **config):
+    def __init__(self, signature, rationale_type=None, **config):
         super().__init__()
 
-        self.activated = activated
-
-        self.signature = signature = ensure_signature(signature)
-        *_keys, last_key = signature.output_fields.keys()
+        signature = ensure_signature(signature)
 
         prefix = "Reasoning: Let's think step by step in order to"
-
-        if isinstance(dspy.settings.lm, dspy.LM):
-            desc = "${reasoning}"
-        elif dspy.settings.experimental:
-            desc = "${produce the output fields}. We ..."
-        else:
-            desc = f"${{produce the {last_key}}}. We ..."
-
+        desc = "${reasoning}"
         rationale_type = rationale_type or dspy.OutputField(prefix=prefix, desc=desc)
-
-        # Add "rationale" field to the output signature.
-        if isinstance(dspy.settings.lm, dspy.LM) or dspy.settings.experimental:
-            extended_signature = signature.prepend("reasoning", rationale_type, type_=str)
-        else:
-            extended_signature = signature.prepend("rationale", rationale_type, type_=str)
+        extended_signature = signature.prepend("reasoning", rationale_type, type_=str)
         
-        self._predict = dspy.Predict(extended_signature, **config)
-        self._predict.extended_signature = extended_signature
+        self.predict = dspy.Predict(extended_signature, **config)
 
     def forward(self, **kwargs):
-        assert self.activated in [True, False]
-
-        signature = kwargs.pop("new_signature", self._predict.extended_signature if self.activated else self.signature)
-        return self._predict(signature=signature, **kwargs)
-
-    @property
-    def demos(self):
-        return self._predict.demos
-
-    @property
-    def extended_signature(self):
-        return self._predict.extended_signature
-
-
-"""
-TODO: In principle, we can update the field's prefix during forward too to fill any thing based on the input args.
-
-IF the user didn't overwrite our default rationale_type.
-"""
+        return self.predict(**kwargs)
diff --git a/dspy/predict/predict.py b/dspy/predict/predict.py
index 3f851e7d..05d66bb4 100644
--- a/dspy/predict/predict.py
+++ b/dspy/predict/predict.py
@@ -51,9 +51,6 @@ class Predict(Module, Parameter):
             state["demos"].append(demo)
 
         state["signature"] = self.signature.dump_state()
-        # `extended_signature` is a special field for `Predict`s like CoT.
-        if hasattr(self, "extended_signature"):
-            state["extended_signature"] = self.extended_signature.dump_state()
         return state
 
     def load_state(self, state):
@@ -82,8 +79,8 @@ class Predict(Module, Parameter):
 
         self.signature = self.signature.load_state(state["signature"])
 
-        if "extended_signature" in state:
-            self.extended_signature = self.extended_signature.load_state(state["extended_signature"])
+        if "extended_signature" in state: # legacy, up to and including 2.5, for CoT.
+            self.signature = self.signature.load_state(state["extended_signature"])
 
         return self
 
@@ -96,14 +93,14 @@ class Predict(Module, Parameter):
         assert not dspy.settings.compiling, "It's no longer ever the case that .compiling is True"
 
         # Extract the three privileged keyword arguments.
-        new_signature = ensure_signature(kwargs.pop("new_signature", None))
+        assert "new_signature" not in kwargs, "new_signature is no longer a valid keyword argument."
         signature = ensure_signature(kwargs.pop("signature", self.signature))
         demos = kwargs.pop("demos", self.demos)
         config = dict(**self.config, **kwargs.pop("config", {}))
 
         # Get the right LM to use.
         lm = kwargs.pop("lm", self.lm) or dspy.settings.lm
-        assert lm is not None, "No LM is loaded."
+        assert isinstance(lm, dspy.LM), "No LM is loaded."
 
         # If temperature is 0.0 but its n > 1, set temperature to 0.7.
         temperature = config.get("temperature")
@@ -113,15 +110,11 @@ class Predict(Module, Parameter):
         if (temperature is None or temperature <= 0.15) and num_generations > 1:
             config["temperature"] = 0.7
 
-        if new_signature is not None:
-            signature = new_signature
-
         if not all(k in kwargs for k in signature.input_fields):
             present = [k for k in signature.input_fields if k in kwargs]
             missing = [k for k in signature.input_fields if k not in kwargs]
             print(f"WARNING: Not all input fields were provided to module. Present: {present}. Missing: {missing}.")
 
-        assert isinstance(lm, dspy.LM)
         completions = v2_5_generate(lm, config, signature, demos, kwargs, _parse_values=self._parse_values)
 
         pred = Prediction.from_completions(completions, signature=signature)
diff --git a/dspy/predict/react.py b/dspy/predict/react.py
index b626f1dc..2690d066 100644
--- a/dspy/predict/react.py
+++ b/dspy/predict/react.py
@@ -124,6 +124,9 @@ means that care must be taken that the adapter is accessed at `forward` runtime,
 Another potential fix is to more natively support a "variadic" input field, where the input is a list of dictionaries,
 or a big dictionary, and have each adatper format it accordingly.
 
+Trajectories also affect meta-programming modules that view the trace later. It's inefficient O(n^2) to view the
+trace of every module repeating the prefix.
+
 
 TOPIC 02: Handling default arguments in the Tool class.
 
@@ -140,4 +143,13 @@ TOPIC 04: Default behavior when the trajectory gets too long.
 TOPIC 05: Adding more structure around how the instruction is formatted.
     * Concretely, it's now a string, so an optimizer can and does rewrite it freely.
     * An alternative would be to add more structure, such that a certain template is fixed but values are variable?
+
+
+TOPIC 06: Idiomatically allowing tools that maintain state across iterations, but not across different `forward` calls.
+    * So the tool would be newly initialized at the start of each `forward` call, but maintain state across iterations.
+    * This is pretty useful for allowing the agent to keep notes or count certain things, etc.
+
+TOPIC 07: Make max_iters a bit more expressive.
+    * Allow passing `max_iters` in forward to overwrite the default.
+    * Get rid of `last_iteration: bool` in the format function. It's not necessary now.
 """
diff --git a/dspy/predict/retry.py b/dspy/predict/retry.py
index b515dfba..66542ba4 100644
--- a/dspy/predict/retry.py
+++ b/dspy/predict/retry.py
@@ -1,74 +1,74 @@
-import copy
+# import copy
 
-import dspy
+# import dspy
 
-from .predict import Predict
+# from .predict import Predict
 
 
-class Retry(Predict):
-    def __init__(self, module):
-        super().__init__(module.signature)
-        self.module = module
-        self.original_signature = module.extended_signature if isinstance(module, dspy.ChainOfThought) else module.signature
-        self.original_forward = module.forward
-        self.new_signature = self._create_new_signature(self.original_signature)
+# class Retry(Predict):
+#     def __init__(self, module):
+#         super().__init__(module.signature)
+#         self.module = module
+#         self.original_signature = module.signature
+#         self.original_forward = module.forward
+#         self.new_signature = self._create_new_signature(self.original_signature)
 
-    def _create_new_signature(self, signature):
-        # Add "Past" input fields for each output field
-        for key, value in signature.output_fields.items():
-            actual_prefix = value.json_schema_extra["prefix"].split(":")[0] + ":"
-            signature = signature.append(f"past_{key}", dspy.InputField(
-                prefix="Previous " + actual_prefix,
-                desc=f"past {actual_prefix[:-1]} with errors",
-                format=value.json_schema_extra.get("format"),
-            ))
+#     def _create_new_signature(self, signature):
+#         # Add "Past" input fields for each output field
+#         for key, value in signature.output_fields.items():
+#             actual_prefix = value.json_schema_extra["prefix"].split(":")[0] + ":"
+#             signature = signature.append(f"past_{key}", dspy.InputField(
+#                 prefix="Previous " + actual_prefix,
+#                 desc=f"past {actual_prefix[:-1]} with errors",
+#                 format=value.json_schema_extra.get("format"),
+#             ))
 
-        signature = signature.append("feedback", dspy.InputField(
-            prefix="Instructions:",
-            desc="Some instructions you must satisfy",
-            format=str,
-        ))
+#         signature = signature.append("feedback", dspy.InputField(
+#             prefix="Instructions:",
+#             desc="Some instructions you must satisfy",
+#             format=str,
+#         ))
 
-        return signature
+#         return signature
 
-    def forward(self, *, past_outputs, **kwargs):
-        # Take into account the possible new signature, as in TypedPredictor
-        new_signature = kwargs.pop("new_signature", None)
-        if new_signature:
-            self.original_signature = new_signature
-            self.new_signature = self._create_new_signature(self.original_signature)
+#     def forward(self, *, past_outputs, **kwargs):
+#         # Take into account the possible new signature, as in TypedPredictor
+#         new_signature = kwargs.pop("new_signature", None)
+#         if new_signature:
+#             self.original_signature = new_signature
+#             self.new_signature = self._create_new_signature(self.original_signature)
 
-        # Convert the dict past_outputs={"answer": ...} to kwargs
-        # {past_answer=..., ...}
-        for key, value in past_outputs.items():
-            past_key = f"past_{key}"
-            if past_key in self.new_signature.input_fields:
-                kwargs[past_key] = value
-        # Tell the wrapped module to use the new signature.
-        # Note: This only works if the wrapped module is a Predict or ChainOfThought.
-        kwargs["new_signature"] = self.new_signature
-        return self.original_forward(**kwargs)
+#         # Convert the dict past_outputs={"answer": ...} to kwargs
+#         # {past_answer=..., ...}
+#         for key, value in past_outputs.items():
+#             past_key = f"past_{key}"
+#             if past_key in self.new_signature.input_fields:
+#                 kwargs[past_key] = value
+#         # Tell the wrapped module to use the new signature.
+#         # Note: This only works if the wrapped module is a Predict or ChainOfThought.
+#         kwargs["new_signature"] = self.new_signature
+#         return self.original_forward(**kwargs)
 
-    def __call__(self, **kwargs):
-        copy.deepcopy(kwargs)
-        kwargs["_trace"] = False
-        kwargs.setdefault("demos", self.demos if self.demos is not None else [])
+#     def __call__(self, **kwargs):
+#         copy.deepcopy(kwargs)
+#         kwargs["_trace"] = False
+#         kwargs.setdefault("demos", self.demos if self.demos is not None else [])
 
-        # perform backtracking
-        if dspy.settings.backtrack_to == self:
-            for key, value in dspy.settings.backtrack_to_args.items():
-                kwargs.setdefault(key, value)
-            pred = self.forward(**kwargs)
-        else:
-            pred = self.module(**kwargs)
+#         # perform backtracking
+#         if dspy.settings.backtrack_to == self:
+#             for key, value in dspy.settings.backtrack_to_args.items():
+#                 kwargs.setdefault(key, value)
+#             pred = self.forward(**kwargs)
+#         else:
+#             pred = self.module(**kwargs)
 
-        # now pop multiple reserved keys
-        # NOTE(shangyin) past_outputs seems not useful to include in demos,
-        # therefore dropped
-        for key in ["_trace", "demos", "signature", "new_signature", "config", "lm", "past_outputs"]:
-            kwargs.pop(key, None)
+#         # now pop multiple reserved keys
+#         # NOTE(shangyin) past_outputs seems not useful to include in demos,
+#         # therefore dropped
+#         for key in ["_trace", "demos", "signature", "new_signature", "config", "lm", "past_outputs"]:
+#             kwargs.pop(key, None)
 
-        if dspy.settings.trace is not None:
-            trace = dspy.settings.trace
-            trace.append((self, {**kwargs}, pred))
-        return pred
+#         if dspy.settings.trace is not None:
+#             trace = dspy.settings.trace
+#             trace.append((self, {**kwargs}, pred))
+#         return pred
diff --git a/dspy/primitives/assertions.py b/dspy/primitives/assertions.py
index a612d02a..8e456cf5 100644
--- a/dspy/primitives/assertions.py
+++ b/dspy/primitives/assertions.py
@@ -1,344 +1,344 @@
-import inspect
-import logging
-import uuid
-from typing import Any
-
-# import dspy.dsp as dsp
-import dspy
-
-logger = logging.getLogger(__name__)
-#################### Assertion Helpers ####################
-
-
-def _build_error_msg(feedback_msgs):
-    """Build an error message from a list of feedback messages."""
-    return "\n".join([msg for msg in feedback_msgs])
-
-
-#################### Assertion Exceptions ####################
-
-
-class DSPyAssertionError(AssertionError):
-    """Custom exception raised when a DSPy `Assert` fails."""
-
-    def __init__(
-        self,
-        id: str,
-        msg: str,
-        target_module: Any = None,
-        state: Any = None,
-        is_metric: bool = False,
-    ) -> None:
-        super().__init__(msg)
-        self.id = id
-        self.msg = msg
-        self.target_module = target_module
-        self.state = state
-        self.is_metric = is_metric
-
-
-class DSPySuggestionError(AssertionError):
-    """Custom exception raised when a DSPy `Suggest` fails."""
-
-    def __init__(
-        self,
-        id: str,
-        msg: str,
-        target_module: Any = None,
-        state: Any = None,
-        is_metric: bool = False,
-    ) -> None:
-        super().__init__(msg)
-        self.id = id
-        self.msg = msg
-        self.target_module = target_module
-        self.state = state
-        self.is_metric = is_metric
-
-
-#################### Assertion Primitives ####################
-
-
-class Constraint:
-    def __init__(
-        self,
-        result: bool,
-        msg: str = "",
-        target_module=None,
-        is_metric: bool = False,
-    ):
-        self.id = str(uuid.uuid4())
-        self.result = result
-        self.msg = msg
-        self.target_module = target_module
-        self.is_metric = is_metric
+# import inspect
+# import logging
+# import uuid
+# from typing import Any
+
+# # import dspy.dsp as dsp
+# import dspy
+
+# logger = logging.getLogger(__name__)
+# #################### Assertion Helpers ####################
+
+
+# def _build_error_msg(feedback_msgs):
+#     """Build an error message from a list of feedback messages."""
+#     return "\n".join([msg for msg in feedback_msgs])
+
+
+# #################### Assertion Exceptions ####################
+
+
+# class DSPyAssertionError(AssertionError):
+#     """Custom exception raised when a DSPy `Assert` fails."""
+
+#     def __init__(
+#         self,
+#         id: str,
+#         msg: str,
+#         target_module: Any = None,
+#         state: Any = None,
+#         is_metric: bool = False,
+#     ) -> None:
+#         super().__init__(msg)
+#         self.id = id
+#         self.msg = msg
+#         self.target_module = target_module
+#         self.state = state
+#         self.is_metric = is_metric
+
+
+# class DSPySuggestionError(AssertionError):
+#     """Custom exception raised when a DSPy `Suggest` fails."""
+
+#     def __init__(
+#         self,
+#         id: str,
+#         msg: str,
+#         target_module: Any = None,
+#         state: Any = None,
+#         is_metric: bool = False,
+#     ) -> None:
+#         super().__init__(msg)
+#         self.id = id
+#         self.msg = msg
+#         self.target_module = target_module
+#         self.state = state
+#         self.is_metric = is_metric
+
+
+# #################### Assertion Primitives ####################
+
+
+# class Constraint:
+#     def __init__(
+#         self,
+#         result: bool,
+#         msg: str = "",
+#         target_module=None,
+#         is_metric: bool = False,
+#     ):
+#         self.id = str(uuid.uuid4())
+#         self.result = result
+#         self.msg = msg
+#         self.target_module = target_module
+#         self.is_metric = is_metric
 
-        self.__call__()
+#         self.__call__()
 
 
-class Assert(Constraint):
-    """DSPy Assertion"""
+# class Assert(Constraint):
+#     """DSPy Assertion"""
 
-    def __call__(self) -> bool:
-        if isinstance(self.result, bool):
-            if self.result:
-                return True
-            elif dspy.settings.bypass_assert:
-                logger.error(f"AssertionError: {self.msg}")
-                return True
-            else:
-                logger.error(f"AssertionError: {self.msg}")
-                raise DSPyAssertionError(
-                    id=self.id,
-                    msg=self.msg,
-                    target_module=self.target_module,
-                    state=dspy.settings.trace,
-                    is_metric=self.is_metric,
-                )
-        else:
-            raise ValueError("Assertion function should always return [bool]")
+#     def __call__(self) -> bool:
+#         if isinstance(self.result, bool):
+#             if self.result:
+#                 return True
+#             elif dspy.settings.bypass_assert:
+#                 logger.error(f"AssertionError: {self.msg}")
+#                 return True
+#             else:
+#                 logger.error(f"AssertionError: {self.msg}")
+#                 raise DSPyAssertionError(
+#                     id=self.id,
+#                     msg=self.msg,
+#                     target_module=self.target_module,
+#                     state=dspy.settings.trace,
+#                     is_metric=self.is_metric,
+#                 )
+#         else:
+#             raise ValueError("Assertion function should always return [bool]")
 
 
-class Suggest(Constraint):
-    """DSPy Suggestion"""
+# class Suggest(Constraint):
+#     """DSPy Suggestion"""
 
-    def __call__(self) -> Any:
-        if isinstance(self.result, bool):
-            if self.result:
-                return True
-            elif dspy.settings.bypass_suggest:
-                logger.info(f"SuggestionFailed: {self.msg}")
-                return True
-            else:
-                logger.info(f"SuggestionFailed: {self.msg}")
-                raise DSPySuggestionError(
-                    id=self.id,
-                    msg=self.msg,
-                    target_module=self.target_module,
-                    state=dspy.settings.trace,
-                    is_metric=self.is_metric,
-                )
-        else:
-            raise ValueError("Suggestion function should always return [bool]")
-
-
-#################### Assertion Handlers ####################
-
-
-def noop_handler(func):
-    """Handler to bypass assertions and suggestions.
-
-    Now both assertions and suggestions will become noops.
-    """
-
-    def wrapper(*args, **kwargs):
-        with dspy.settings.context(bypass_assert=True, bypass_suggest=True):
-            return func(*args, **kwargs)
-
-    return wrapper
-
-
-def bypass_suggest_handler(func):
-    """Handler to bypass suggest only.
-
-    If a suggestion fails, it will be logged but not raised.
-    And If an assertion fails, it will be raised.
-    """
-
-    def wrapper(*args, **kwargs):
-        with dspy.settings.context(bypass_suggest=True, bypass_assert=False):
-            return func(*args, **kwargs)
-
-    return wrapper
-
-
-def bypass_assert_handler(func):
-    """Handler to bypass assertion only.
-
-    If a assertion fails, it will be logged but not raised.
-    And If an assertion fails, it will be raised.
-    """
-
-    def wrapper(*args, **kwargs):
-        with dspy.settings.context(bypass_assert=True):
-            return func(*args, **kwargs)
-
-    return wrapper
-
-
-def assert_no_except_handler(func):
-    """Handler to ignore assertion failure and return None."""
-
-    def wrapper(*args, **kwargs):
-        try:
-            return func(*args, **kwargs)
-        except DSPyAssertionError:
-            return None
-
-    return wrapper
-
-
-def backtrack_handler(func, bypass_suggest=True, max_backtracks=2):
-    """Handler for backtracking suggestion and assertion.
-
-    Re-run the latest predictor up to `max_backtracks` times,
-    with updated signature if an assertion fails. updated signature adds a new
-    input field to the signature, which is the feedback.
-    """
-
-    def wrapper(*args, **kwargs):
-        error_msg, result = None, None
-        with dspy.settings.lock:
-            dspy.settings.backtrack_to = None
-            dspy.settings.suggest_failures = 0
-            dspy.settings.assert_failures = 0
-
-            # Predictor -> List[feedback_msg]
-            dspy.settings.predictor_feedbacks = {}
-
-            current_error = None
-            for i in range(max_backtracks + 1):
-                if i > 0 and dspy.settings.backtrack_to is not None:
-                    # generate values for new fields
-                    feedback_msg = _build_error_msg(
-                        dspy.settings.predictor_feedbacks[dspy.settings.backtrack_to],
-                    )
-
-                    dspy.settings.backtrack_to_args = {
-                        "feedback": feedback_msg,
-                        "past_outputs": past_outputs,
-                    }
-
-                # if last backtrack: ignore suggestion errors
-                if i == max_backtracks:
-                    if isinstance(current_error, DSPyAssertionError):
-                        raise current_error
-                    dspy.settings.trace.clear()
-                    result = bypass_suggest_handler(func)(*args, **kwargs) if bypass_suggest else None
-                    break
-                else:
-                    try:
-                        dspy.settings.trace.clear()
-                        result = func(*args, **kwargs)
-                        break
-                    except (DSPySuggestionError, DSPyAssertionError) as e:
-                        if not current_error:
-                            current_error = e
-                        _error_id, error_msg, error_target_module, error_state = (
-                            e.id,
-                            e.msg,
-                            e.target_module,
-                            e.state[-1],
-                        )
-
-                        # increment failure count depending on type of error
-                        if isinstance(e, DSPySuggestionError) and e.is_metric:
-                            dspy.settings.suggest_failures += 1
-                        elif isinstance(e, DSPyAssertionError) and e.is_metric:
-                            dspy.settings.assert_failures += 1
-
-                        if dspy.settings.trace:
-                            if error_target_module:
-                                for i in range(len(dspy.settings.trace) - 1, -1, -1):
-                                    trace_element = dspy.settings.trace[i]
-                                    mod = trace_element[0]
-                                    if mod == error_target_module:
-                                        error_state = e.state[i]
-                                        dspy.settings.backtrack_to = mod
-                                        break
-                            else:
-                                dspy.settings.backtrack_to = dspy.settings.trace[-1][0]
-
-                            if dspy.settings.backtrack_to is None:
-                                logger.error("Module not found in trace. If passing a DSPy Signature, please specify the intended module for the assertion (e.g., use `target_module = self.my_module(my_signature)` instead of `target_module =  my_signature`).")
-
-                            # save unique feedback message for predictor
-                            if error_msg not in dspy.settings.predictor_feedbacks.setdefault(
-                                dspy.settings.backtrack_to,
-                                [],
-                            ):
-                                dspy.settings.predictor_feedbacks[dspy.settings.backtrack_to].append(error_msg)
-
-                            # use `new_signature` if available (CoT)
-                            if hasattr(error_state[0], 'new_signature'):
-                                output_fields = error_state[0].new_signature.output_fields
-                            else:
-                                output_fields = error_state[0].signature.output_fields
-                            past_outputs = {}
-                            for field_name in output_fields.keys():
-                                past_outputs[field_name] = getattr(
-                                    error_state[2],
-                                    field_name,
-                                    None,
-                                )
-
-                            # save latest failure trace for predictor per suggestion
-                            error_state[1]
-                            error_op = error_state[2].__dict__["_store"]
-                            error_op.pop("_assert_feedback", None)
-                            error_op.pop("_assert_traces", None)
-
-                        else:
-                            logger.error(
-                                "UNREACHABLE: No trace available, this should not happen. Is this run time?",
-                            )
-
-            return result
-
-    return wrapper
-
-
-def handle_assert_forward(assertion_handler, **handler_args):
-    def forward(self, *args, **kwargs):
-        args_to_vals = inspect.getcallargs(self._forward, *args, **kwargs)
-
-        # if user has specified a bypass_assert flag, set it
-        if "bypass_assert" in args_to_vals:
-            dspy.settings.configure(bypass_assert=args_to_vals["bypass_assert"])
-
-        wrapped_forward = assertion_handler(self._forward, **handler_args)
-        return wrapped_forward(*args, **kwargs)
-
-    return forward
-
-
-default_assertion_handler = backtrack_handler
-
-
-def assert_transform_module(
-    module,
-    assertion_handler=default_assertion_handler,
-    **handler_args,
-):
-    """
-    Transform a module to handle assertions.
-    """
-    if not getattr(module, "forward", False):
-        raise ValueError(
-            "Module must have a forward method to have assertions handled.",
-        )
-    if getattr(module, "_forward", False):
-        logger.info(
-            f"Module {module.__class__.__name__} already has a _forward method. Skipping...",
-        )
-        pass  # TODO warning: might be overwriting a previous _forward method
-
-    module._forward = module.forward
-    module.forward = handle_assert_forward(assertion_handler, **handler_args).__get__(
-        module,
-    )
-
-    if all(
-        map(lambda p: isinstance(p[1], dspy.retry.Retry), module.named_predictors()),
-    ):
-        pass  # we already applied the Retry mapping outside
-    elif all(
-        map(lambda p: not isinstance(p[1], dspy.retry.Retry), module.named_predictors()),
-    ):
-        module.map_named_predictors(dspy.retry.Retry)
-    else:
-        raise RuntimeError("Module has mixed predictors, can't apply Retry mapping.")
-
-    module._assert_transformed = True
-
-    return module
+#     def __call__(self) -> Any:
+#         if isinstance(self.result, bool):
+#             if self.result:
+#                 return True
+#             elif dspy.settings.bypass_suggest:
+#                 logger.info(f"SuggestionFailed: {self.msg}")
+#                 return True
+#             else:
+#                 logger.info(f"SuggestionFailed: {self.msg}")
+#                 raise DSPySuggestionError(
+#                     id=self.id,
+#                     msg=self.msg,
+#                     target_module=self.target_module,
+#                     state=dspy.settings.trace,
+#                     is_metric=self.is_metric,
+#                 )
+#         else:
+#             raise ValueError("Suggestion function should always return [bool]")
+
+
+# #################### Assertion Handlers ####################
+
+
+# def noop_handler(func):
+#     """Handler to bypass assertions and suggestions.
+
+#     Now both assertions and suggestions will become noops.
+#     """
+
+#     def wrapper(*args, **kwargs):
+#         with dspy.settings.context(bypass_assert=True, bypass_suggest=True):
+#             return func(*args, **kwargs)
+
+#     return wrapper
+
+
+# def bypass_suggest_handler(func):
+#     """Handler to bypass suggest only.
+
+#     If a suggestion fails, it will be logged but not raised.
+#     And If an assertion fails, it will be raised.
+#     """
+
+#     def wrapper(*args, **kwargs):
+#         with dspy.settings.context(bypass_suggest=True, bypass_assert=False):
+#             return func(*args, **kwargs)
+
+#     return wrapper
+
+
+# def bypass_assert_handler(func):
+#     """Handler to bypass assertion only.
+
+#     If a assertion fails, it will be logged but not raised.
+#     And If an assertion fails, it will be raised.
+#     """
+
+#     def wrapper(*args, **kwargs):
+#         with dspy.settings.context(bypass_assert=True):
+#             return func(*args, **kwargs)
+
+#     return wrapper
+
+
+# def assert_no_except_handler(func):
+#     """Handler to ignore assertion failure and return None."""
+
+#     def wrapper(*args, **kwargs):
+#         try:
+#             return func(*args, **kwargs)
+#         except DSPyAssertionError:
+#             return None
+
+#     return wrapper
+
+
+# def backtrack_handler(func, bypass_suggest=True, max_backtracks=2):
+#     """Handler for backtracking suggestion and assertion.
+
+#     Re-run the latest predictor up to `max_backtracks` times,
+#     with updated signature if an assertion fails. updated signature adds a new
+#     input field to the signature, which is the feedback.
+#     """
+
+#     def wrapper(*args, **kwargs):
+#         error_msg, result = None, None
+#         with dspy.settings.lock:
+#             dspy.settings.backtrack_to = None
+#             dspy.settings.suggest_failures = 0
+#             dspy.settings.assert_failures = 0
+
+#             # Predictor -> List[feedback_msg]
+#             dspy.settings.predictor_feedbacks = {}
+
+#             current_error = None
+#             for i in range(max_backtracks + 1):
+#                 if i > 0 and dspy.settings.backtrack_to is not None:
+#                     # generate values for new fields
+#                     feedback_msg = _build_error_msg(
+#                         dspy.settings.predictor_feedbacks[dspy.settings.backtrack_to],
+#                     )
+
+#                     dspy.settings.backtrack_to_args = {
+#                         "feedback": feedback_msg,
+#                         "past_outputs": past_outputs,
+#                     }
+
+#                 # if last backtrack: ignore suggestion errors
+#                 if i == max_backtracks:
+#                     if isinstance(current_error, DSPyAssertionError):
+#                         raise current_error
+#                     dspy.settings.trace.clear()
+#                     result = bypass_suggest_handler(func)(*args, **kwargs) if bypass_suggest else None
+#                     break
+#                 else:
+#                     try:
+#                         dspy.settings.trace.clear()
+#                         result = func(*args, **kwargs)
+#                         break
+#                     except (DSPySuggestionError, DSPyAssertionError) as e:
+#                         if not current_error:
+#                             current_error = e
+#                         _error_id, error_msg, error_target_module, error_state = (
+#                             e.id,
+#                             e.msg,
+#                             e.target_module,
+#                             e.state[-1],
+#                         )
+
+#                         # increment failure count depending on type of error
+#                         if isinstance(e, DSPySuggestionError) and e.is_metric:
+#                             dspy.settings.suggest_failures += 1
+#                         elif isinstance(e, DSPyAssertionError) and e.is_metric:
+#                             dspy.settings.assert_failures += 1
+
+#                         if dspy.settings.trace:
+#                             if error_target_module:
+#                                 for i in range(len(dspy.settings.trace) - 1, -1, -1):
+#                                     trace_element = dspy.settings.trace[i]
+#                                     mod = trace_element[0]
+#                                     if mod == error_target_module:
+#                                         error_state = e.state[i]
+#                                         dspy.settings.backtrack_to = mod
+#                                         break
+#                             else:
+#                                 dspy.settings.backtrack_to = dspy.settings.trace[-1][0]
+
+#                             if dspy.settings.backtrack_to is None:
+#                                 logger.error("Module not found in trace. If passing a DSPy Signature, please specify the intended module for the assertion (e.g., use `target_module = self.my_module(my_signature)` instead of `target_module =  my_signature`).")
+
+#                             # save unique feedback message for predictor
+#                             if error_msg not in dspy.settings.predictor_feedbacks.setdefault(
+#                                 dspy.settings.backtrack_to,
+#                                 [],
+#                             ):
+#                                 dspy.settings.predictor_feedbacks[dspy.settings.backtrack_to].append(error_msg)
+
+#                             # use `new_signature` if available (CoT)
+#                             if hasattr(error_state[0], 'new_signature'):
+#                                 output_fields = error_state[0].new_signature.output_fields
+#                             else:
+#                                 output_fields = error_state[0].signature.output_fields
+#                             past_outputs = {}
+#                             for field_name in output_fields.keys():
+#                                 past_outputs[field_name] = getattr(
+#                                     error_state[2],
+#                                     field_name,
+#                                     None,
+#                                 )
+
+#                             # save latest failure trace for predictor per suggestion
+#                             error_state[1]
+#                             error_op = error_state[2].__dict__["_store"]
+#                             error_op.pop("_assert_feedback", None)
+#                             error_op.pop("_assert_traces", None)
+
+#                         else:
+#                             logger.error(
+#                                 "UNREACHABLE: No trace available, this should not happen. Is this run time?",
+#                             )
+
+#             return result
+
+#     return wrapper
+
+
+# def handle_assert_forward(assertion_handler, **handler_args):
+#     def forward(self, *args, **kwargs):
+#         args_to_vals = inspect.getcallargs(self._forward, *args, **kwargs)
+
+#         # if user has specified a bypass_assert flag, set it
+#         if "bypass_assert" in args_to_vals:
+#             dspy.settings.configure(bypass_assert=args_to_vals["bypass_assert"])
+
+#         wrapped_forward = assertion_handler(self._forward, **handler_args)
+#         return wrapped_forward(*args, **kwargs)
+
+#     return forward
+
+
+# default_assertion_handler = backtrack_handler
+
+
+# def assert_transform_module(
+#     module,
+#     assertion_handler=default_assertion_handler,
+#     **handler_args,
+# ):
+#     """
+#     Transform a module to handle assertions.
+#     """
+#     if not getattr(module, "forward", False):
+#         raise ValueError(
+#             "Module must have a forward method to have assertions handled.",
+#         )
+#     if getattr(module, "_forward", False):
+#         logger.info(
+#             f"Module {module.__class__.__name__} already has a _forward method. Skipping...",
+#         )
+#         pass  # TODO warning: might be overwriting a previous _forward method
+
+#     module._forward = module.forward
+#     module.forward = handle_assert_forward(assertion_handler, **handler_args).__get__(
+#         module,
+#     )
+
+#     if all(
+#         map(lambda p: isinstance(p[1], dspy.retry.Retry), module.named_predictors()),
+#     ):
+#         pass  # we already applied the Retry mapping outside
+#     elif all(
+#         map(lambda p: not isinstance(p[1], dspy.retry.Retry), module.named_predictors()),
+#     ):
+#         module.map_named_predictors(dspy.retry.Retry)
+#     else:
+#         raise RuntimeError("Module has mixed predictors, can't apply Retry mapping.")
+
+#     module._assert_transformed = True
+
+#     return module
diff --git a/dspy/primitives/module.py b/dspy/primitives/module.py
index c2146301..8bec9162 100644
--- a/dspy/primitives/module.py
+++ b/dspy/primitives/module.py
@@ -256,14 +256,15 @@ class BaseModule:
 
 
 def postprocess_parameter_name(name, value):
-    # For ChainOfThought backward compatibility, remove ending ._predict if it's there
-    if name.endswith("._predict"):
-        name = name[:-9]
+    return name
+    # # For ChainOfThought backward compatibility, remove ending ._predict if it's there
+    # if name.endswith("._predict"):
+    #     name = name[:-9]
 
-    if name.endswith(".self"):
-        name = name[:-5]
+    # if name.endswith(".self"):
+    #     name = name[:-5]
 
-    if name == "_predict":
-        return "self"
+    # if name == "_predict":
+    #     return "self"
 
-    return name
+    # return name
diff --git a/dspy/primitives/program.py b/dspy/primitives/program.py
index a8b2d134..a8b9676d 100644
--- a/dspy/primitives/program.py
+++ b/dspy/primitives/program.py
@@ -1,10 +1,10 @@
 from dspy.utils.callback import with_callbacks
 import magicattr
 
-import dspy
 from dspy.predict.parallel import Parallel
-from dspy.primitives.assertions import *
 from dspy.primitives.module import BaseModule
+# import dspy
+# from dspy.primitives.assertions import *
 
 
 class ProgramMeta(type):
@@ -32,28 +32,16 @@ class Module(BaseModule, metaclass=ProgramMeta):
         return [param for _, param in self.named_predictors()]
 
     def set_lm(self, lm):
-        if not dspy.settings.experimental:
-            raise ValueError(
-                "Setting or getting the LM of a program is an experimental feature. Please enable the "
-                "'dspy.settings.experimental' flag to use these features."
-            )
-
         for _, param in self.named_predictors():
             param.lm = lm
 
     def get_lm(self):
-        if not dspy.settings.experimental:
-            raise ValueError(
-                "Setting or getting the LM of a program is an experimental feature. Please enable the "
-                "'dspy.settings.experimental' flag to use these features."
-            )
-
         all_used_lms = [param.lm for _, param in self.named_predictors()]
 
         if len(set(all_used_lms)) == 1:
             return all_used_lms[0]
 
-        raise ValueError("Multiple LMs are being used in the module.")
+        raise ValueError("Multiple LMs are being used in the module. There's no unique LM to return.")
 
     def __repr__(self):
         s = []
@@ -69,13 +57,13 @@ class Module(BaseModule, metaclass=ProgramMeta):
             set_attribute_by_name(self, name, func(predictor))
         return self
 
-    def activate_assertions(self, handler=backtrack_handler, **handler_args):
-        """
-        Activates assertions for the module.
-        The default handler is the backtrack_handler.
-        """
-        assert_transform_module(self, handler, **handler_args)
-        return self
+    # def activate_assertions(self, handler=backtrack_handler, **handler_args):
+    #     """
+    #     Activates assertions for the module.
+    #     The default handler is the backtrack_handler.
+    #     """
+    #     assert_transform_module(self, handler, **handler_args)
+    #     return self
 
     # def __deepcopy__(self, memo):
     #     # memo is a dict of id's to copies already made during the current call
@@ -103,7 +91,7 @@ class Module(BaseModule, metaclass=ProgramMeta):
         return_failed_examples: bool = False,
         provide_traceback: bool = False,
         disable_progress_bar: bool = False,
-    ) -> Any:
+    ):
         """
         Processes a list of dspy.Example instances in parallel using the Parallel module.
 
diff --git a/dspy/teleprompt/copro_optimizer.py b/dspy/teleprompt/copro_optimizer.py
index de8cc9f0..feeecf9e 100644
--- a/dspy/teleprompt/copro_optimizer.py
+++ b/dspy/teleprompt/copro_optimizer.py
@@ -113,16 +113,12 @@ class COPRO(Teleprompter):
         logger.debug(f"p: {list(signature.fields.values())[-1].json_schema_extra['prefix']}")
 
     def _get_signature(self, predictor):
-        if hasattr(predictor, "extended_signature"):
-            return predictor.extended_signature
-        elif hasattr(predictor, "signature"):
-            return predictor.signature
+        assert hasattr(predictor, "signature")
+        return predictor.signature
 
     def _set_signature(self, predictor, updated_signature):
-        if hasattr(predictor, "extended_signature"):
-            predictor.extended_signature = updated_signature
-        elif hasattr(predictor, "signature"):
-            predictor.signature = updated_signature
+        assert hasattr(predictor, "signature")
+        predictor.signature = updated_signature
 
     def compile(self, student, *, trainset, eval_kwargs):
         """
diff --git a/dspy/teleprompt/utils.py b/dspy/teleprompt/utils.py
index d763d74d..bdc6a529 100644
--- a/dspy/teleprompt/utils.py
+++ b/dspy/teleprompt/utils.py
@@ -257,18 +257,13 @@ def get_prompt_model(prompt_model):
         return dspy.settings.lm
 
 def get_signature(predictor):
-    if hasattr(predictor, "extended_signature"):
-        return predictor.extended_signature
-    elif hasattr(predictor, "signature"):
-        return predictor.signature
-    return None
+    assert hasattr(predictor, "signature")
+    return predictor.signature
 
 
 def set_signature(predictor, updated_signature):
-    if hasattr(predictor, "extended_signature"):
-        predictor.extended_signature = updated_signature
-    elif hasattr(predictor, "signature"):
-        predictor.signature = updated_signature
+    assert hasattr(predictor, "signature")
+    predictor.signature = updated_signature
 
 
 def create_n_fewshot_demo_sets(


Test patch:
diff --git a/tests/predict/test_chain_of_thought.py b/tests/predict/test_chain_of_thought.py
index 04f73bb5..c050819e 100644
--- a/tests/predict/test_chain_of_thought.py
+++ b/tests/predict/test_chain_of_thought.py
@@ -7,7 +7,7 @@ def test_initialization_with_string_signature():
     lm = DummyLM([{"reasoning": "find the number after 1", "answer": "2"}])
     dspy.settings.configure(lm=lm)
     predict = ChainOfThought("question -> answer")
-    assert list(predict.extended_signature.output_fields.keys()) == [
+    assert list(predict.predict.signature.output_fields.keys()) == [
         "reasoning",
         "answer",
     ]
diff --git a/tests/predict/test_retry.py b/tests/predict/test_retry.py
index ba28e904..eb1c48b0 100644
--- a/tests/predict/test_retry.py
+++ b/tests/predict/test_retry.py
@@ -1,86 +1,48 @@
-import functools
+# import functools
 
-import pydantic
+# import pydantic
 
-import dspy
-from dspy.primitives.assertions import assert_transform_module, backtrack_handler
-from dspy.utils import DummyLM
+# import dspy
+# from dspy.primitives.assertions import assert_transform_module, backtrack_handler
+# from dspy.utils import DummyLM
 
 
-def test_retry_simple():
-    predict = dspy.Predict("question -> answer")
-    retry_module = dspy.Retry(predict)
+# def test_retry_simple():
+#     predict = dspy.Predict("question -> answer")
+#     retry_module = dspy.Retry(predict)
 
-    # Test Retry has created the correct new signature
-    for field in predict.signature.output_fields:
-        assert f"past_{field}" in retry_module.new_signature.input_fields
-    assert "feedback" in retry_module.new_signature.input_fields
+#     # Test Retry has created the correct new signature
+#     for field in predict.signature.output_fields:
+#         assert f"past_{field}" in retry_module.new_signature.input_fields
+#     assert "feedback" in retry_module.new_signature.input_fields
 
-    lm = DummyLM([{"answer": "blue"}])
-    dspy.settings.configure(lm=lm)
-    result = retry_module.forward(
-        question="What color is the sky?",
-        past_outputs={"answer": "red"},
-        feedback="Try harder",
-    )
-    assert result.answer == "blue"
-
-
-def test_retry_forward_with_feedback():
-    # First we make a mistake, then we fix it
-    lm = DummyLM([{"answer": "red"}, {"answer": "blue"}])
-    dspy.settings.configure(lm=lm, trace=[])
-
-    class SimpleModule(dspy.Module):
-        def __init__(self):
-            super().__init__()
-            self.predictor = dspy.Predict("question -> answer")
-
-        def forward(self, **kwargs):
-            result = self.predictor(**kwargs)
-            print(f"SimpleModule got {result.answer=}")
-            dspy.Suggest(result.answer == "blue", "Please think harder")
-            return result
-
-    program = SimpleModule()
-    program = assert_transform_module(
-        program.map_named_predictors(dspy.Retry),
-        functools.partial(backtrack_handler, max_backtracks=1),
-    )
-
-    result = program(question="What color is the sky?")
-
-    assert result.answer == "blue"
+#     lm = DummyLM([{"answer": "blue"}])
+#     dspy.settings.configure(lm=lm)
+#     result = retry_module.forward(
+#         question="What color is the sky?",
+#         past_outputs={"answer": "red"},
+#         feedback="Try harder",
+#     )
+#     assert result.answer == "blue"
 
 
-# def test_retry_forward_with_typed_predictor():
+# def test_retry_forward_with_feedback():
 #     # First we make a mistake, then we fix it
-#     lm = DummyLM([{"output": '{"answer":"red"}'}, {"output": '{"answer":"blue"}'}])
+#     lm = DummyLM([{"answer": "red"}, {"answer": "blue"}])
 #     dspy.settings.configure(lm=lm, trace=[])
 
-#     class AnswerQuestion(dspy.Signature):
-#         """Answer questions with succinct responses."""
-
-#         class Input(pydantic.BaseModel):
-#             question: str
-
-#         class Output(pydantic.BaseModel):
-#             answer: str
-
-#         input: Input = dspy.InputField()
-#         output: Output = dspy.OutputField()
-
-#     class QuestionAnswerer(dspy.Module):
+#     class SimpleModule(dspy.Module):
 #         def __init__(self):
 #             super().__init__()
-#             self.answer_question = dspy.TypedPredictor(AnswerQuestion)
+#             self.predictor = dspy.Predict("question -> answer")
 
 #         def forward(self, **kwargs):
-#             result = self.answer_question(input=AnswerQuestion.Input(**kwargs)).output
+#             result = self.predictor(**kwargs)
+#             print(f"SimpleModule got {result.answer=}")
 #             dspy.Suggest(result.answer == "blue", "Please think harder")
 #             return result
 
-#     program = QuestionAnswerer()
+#     program = SimpleModule()
 #     program = assert_transform_module(
 #         program.map_named_predictors(dspy.Retry),
 #         functools.partial(backtrack_handler, max_backtracks=1),
@@ -89,3 +51,41 @@ def test_retry_forward_with_feedback():
 #     result = program(question="What color is the sky?")
 
 #     assert result.answer == "blue"
+
+
+# # def test_retry_forward_with_typed_predictor():
+# #     # First we make a mistake, then we fix it
+# #     lm = DummyLM([{"output": '{"answer":"red"}'}, {"output": '{"answer":"blue"}'}])
+# #     dspy.settings.configure(lm=lm, trace=[])
+
+# #     class AnswerQuestion(dspy.Signature):
+# #         """Answer questions with succinct responses."""
+
+# #         class Input(pydantic.BaseModel):
+# #             question: str
+
+# #         class Output(pydantic.BaseModel):
+# #             answer: str
+
+# #         input: Input = dspy.InputField()
+# #         output: Output = dspy.OutputField()
+
+# #     class QuestionAnswerer(dspy.Module):
+# #         def __init__(self):
+# #             super().__init__()
+# #             self.answer_question = dspy.TypedPredictor(AnswerQuestion)
+
+# #         def forward(self, **kwargs):
+# #             result = self.answer_question(input=AnswerQuestion.Input(**kwargs)).output
+# #             dspy.Suggest(result.answer == "blue", "Please think harder")
+# #             return result
+
+# #     program = QuestionAnswerer()
+# #     program = assert_transform_module(
+# #         program.map_named_predictors(dspy.Retry),
+# #         functools.partial(backtrack_handler, max_backtracks=1),
+# #     )
+
+# #     result = program(question="What color is the sky?")
+
+# #     assert result.answer == "blue"
diff --git a/tests/primitives/test_module.py b/tests/primitives/test_module.py
index 5e6790d8..d0532fca 100644
--- a/tests/primitives/test_module.py
+++ b/tests/primitives/test_module.py
@@ -51,8 +51,8 @@ def test_deepcopy_with_nested_modules():
 
 def test_save_and_load_with_json(tmp_path):
     model = dspy.ChainOfThought(dspy.Signature("q -> a"))
-    model._predict.signature = model._predict.signature.with_instructions("You are a helpful assistant.")
-    model._predict.demos = [
+    model.predict.signature = model.predict.signature.with_instructions("You are a helpful assistant.")
+    model.predict.demos = [
         dspy.Example(q="What is the capital of France?", a="Paris", reasoning="n/a").with_inputs("q", "a")
     ]
     save_path = tmp_path / "model.json"
@@ -60,8 +60,8 @@ def test_save_and_load_with_json(tmp_path):
     new_model = dspy.ChainOfThought(dspy.Signature("q -> a"))
     new_model.load(save_path)
 
-    assert str(new_model.signature) == str(model.signature)
-    assert new_model.demos[0] == model.demos[0].toDict()
+    assert str(new_model.predict.signature) == str(model.predict.signature)
+    assert new_model.predict.demos[0] == model.predict.demos[0].toDict()
 
 
 def test_save_and_load_with_pkl(tmp_path):
@@ -96,7 +96,7 @@ def test_save_and_load_with_pkl(tmp_path):
 
     optimizer = dspy.BootstrapFewShot(max_bootstrapped_demos=4, max_labeled_demos=4, max_rounds=5, metric=dummy_metric)
     compiled_cot = optimizer.compile(cot, trainset=trainset)
-    compiled_cot._predict.signature = compiled_cot._predict.signature.with_instructions("You are a helpful assistant.")
+    compiled_cot.predict.signature = compiled_cot.predict.signature.with_instructions("You are a helpful assistant.")
 
     save_path = tmp_path / "program.pkl"
     compiled_cot.save(save_path)
@@ -104,5 +104,5 @@ def test_save_and_load_with_pkl(tmp_path):
     new_cot = dspy.ChainOfThought(MySignature)
     new_cot.load(save_path)
 
-    assert str(new_cot.signature) == str(compiled_cot.signature)
-    assert new_cot.demos == compiled_cot.demos
+    assert str(new_cot.predict.signature) == str(compiled_cot.predict.signature)
+    assert new_cot.predict.demos == compiled_cot.predict.demos