Repository: stanfordnlp/dspy
Pull request #7926 (2025-03-09 06:27:05)

Adds Improved Error Handling, Documentation, and Tests For The Refine Module
Of note: Adds configurable `fail_count` parameter to `Refine` to control how many errors are tolerated before throwing. By default we will throw once if all of the underlying module's calls fails.

Adds in-line documentation to both `Refine` and `BestOfN` and in the cheatsheet.




okhat: This is amazing, thanks so much @cezarc1 ! Timing is interesting because we're considering how to act on this and #7922 together. We will merge both PRs, but the question is what's the final design that emerges from that.
cezarc1: @okhat . I see! I missed #7922. As I think it was mentioned in #7922 I still see the benefits of`Refine` and `BestofN` still existing as my understanding is that the retry mechanism in that PR requires a new adapter, while `Refine` and `BestOfN` do not require a new adapter in order to work.

Patch:
diff --git a/docs/docs/cheatsheet.md b/docs/docs/cheatsheet.md
index 491e4902..e1188acd 100644
--- a/docs/docs/cheatsheet.md
+++ b/docs/docs/cheatsheet.md
@@ -491,34 +491,55 @@ your_dspy_program_compiled = fewshot_optuna_optimizer.compile(student=your_dspy_
 
 Other custom configurations are similar to customizing the `dspy.BootstrapFewShot` optimizer.
 
-## DSPy Assertions
+## DSPy `Refine` and `BestofN`
 
-### Including `dspy.Assert` and `dspy.Suggest` statements
+>`dspy.Suggest` and `dspy.Assert` are replaced by `dspy.Refine` and `dspy.BestofN` in DSPy 2.6.
+
+### BestofN
+
+Runs a module up to `N` times with different temperatures and returns the best prediction, as defined by the `reward_fn`, or the first prediction that passes the `threshold`.
 
 ```python
-dspy.Assert(your_validation_fn(model_outputs), "your feedback message", target_module="YourDSPyModule")
+import dspy
 
-dspy.Suggest(your_validation_fn(model_outputs), "your feedback message", target_module="YourDSPyModule")
+qa = dspy.ChainOfThought("question -> answer")
+def one_word_answer(args, pred):
+    return 1.0 if len(pred.answer) == 1 else 0.0
+best_of_3 = dspy.BestOfN(module=qa, N=3, reward_fn=one_word_answer, threshold=1.0)
+best_of_3(question="What is the capital of Belgium?").answer
+# Brussels
 ```
 
-### Activating DSPy Program with Assertions
+### Refine
 
-**Note**: To use Assertions properly, you must **activate** a DSPy program that includes `dspy.Assert` or `dspy.Suggest` statements from either of the methods above.
+Refines a module by running it up to `N` times with different temperatures and returns the best prediction, as defined by the `reward_fn`, or the first prediction that passes the `threshold`. After each attempt (except the final one), `Refine` automatically generates detailed feedback about the module's performance and uses this feedback as hints for subsequent runs, creating an iterative refinement process.
 
 ```python
-#1. Using `assert_transform_module:
-from dspy.primitives.assertions import assert_transform_module, backtrack_handler
+import dspy
 
-program_with_assertions = assert_transform_module(ProgramWithAssertions(), backtrack_handler)
+qa = dspy.ChainOfThought("question -> answer")
+def one_word_answer(args, pred):
+    return 1.0 if len(pred.answer) == 1 else 0.0
+best_of_3 = dspy.Refine(module=qa, N=3, reward_fn=one_word_answer, threshold=1.0)
+best_of_3(question="What is the capital of Belgium?").answer
+# Brussels
+```
+
+#### Error Handling
 
-#2. Using `activate_assertions()`
-program_with_assertions = ProgramWithAssertions().activate_assertions()
+By default, `Refine` will try to run the module up to N times until the threshold is met. If the module encounters an error, it will keep going up to N failed attempts. You can change this behavior by setting `fail_count` to a smaller number than `N`.
+
+```python
+refine = dspy.Refine(module=qa, N=3, reward_fn=one_word_answer, threshold=1.0, fail_count=1)
+...
+refine(question="What is the capital of Belgium?")
+# If we encounter just one failed attempt, the module will raise an error.
 ```
 
-## Compiling with DSPy Programs with Assertions
+If you want to run the module up to N times without any error handling, you can set `fail_count` to `N`. This is the default behavior.
 
 ```python
-program_with_assertions = assert_transform_module(ProgramWithAssertions(), backtrack_handler)
-fewshot_optimizer = BootstrapFewShotWithRandomSearch(metric = your_defined_metric, max_bootstrapped_demos=2, num_candidate_programs=6)
-compiled_dspy_program_with_assertions = fewshot_optimizer.compile(student=program_with_assertions, teacher = program_with_assertions, trainset=trainset, valset=devset) #student can also be program_without_assertions
+refine = dspy.Refine(module=qa, N=3, reward_fn=one_word_answer, threshold=1.0, fail_count=3)
+...
+refine(question="What is the capital of Belgium?")
 ```
diff --git a/dspy/predict/best_of_n.py b/dspy/predict/best_of_n.py
index 6ca0dd90..452c8617 100644
--- a/dspy/predict/best_of_n.py
+++ b/dspy/predict/best_of_n.py
@@ -6,14 +6,34 @@ from typing import Callable
 
 class BestOfN(Module):
     def __init__(self, module, N: int, reward_fn: Callable, threshold: float):
+        """
+        Runs a module up to `N` times with different temperatures and returns the best prediction or the first prediction that passes the `threshold`.
+
+        Args:
+            module: The module to run.
+            N: The number of times to run the module.
+            reward_fn: The reward function.
+            threshold: The threshold for the reward function.
+
+        Example:
+            >>> import dspy
+            >>> qa = dspy.ChainOfThought("question -> answer")
+            >>> def one_word_answer(args, pred):
+            >>>     return 1.0 if len(pred.answer) == 1 else 0.0
+            >>> best_of_3 = dspy.BestOfN(module=qa, N=3, reward_fn=one_word_answer, threshold=1.0)
+            >>> best_of_3(question="What is the capital of Belgium?").answer
+            >>> # Brussels
+        """
         self.module = module
-        self.reward_fn = lambda *args: reward_fn(*args)  # to prevent this from becoming a parameter
+        self.reward_fn = lambda *args: reward_fn(
+            *args)  # to prevent this from becoming a parameter
         self.threshold = threshold
         self.N = N
 
     def forward(self, **kwargs):
         lm = self.module.get_lm() or dspy.settings.lm
-        temps = [lm.kwargs['temperature']] + [0.5 + i * (0.5 / self.N) for i in range(self.N)]
+        temps = [lm.kwargs['temperature']
+                 ] + [0.5 + i * (0.5 / self.N) for i in range(self.N)]
         temps = list(dict.fromkeys(temps))[:self.N]
         best_pred, best_trace, best_reward = None, None, -float("inf")
 
@@ -29,15 +49,15 @@ class BestOfN(Module):
 
                     # NOTE: Not including the trace of reward_fn.
                     reward = self.reward_fn(kwargs, pred)
-                
+
                 if reward > best_reward:
                     best_reward, best_pred, best_trace = reward, pred, trace
-                
+
                 if self.threshold is not None and reward >= self.threshold:
                     break
-                
+
             except Exception as e:
                 print(f"Attempt failed with temperature {t}: {e}")
-        
+
         dspy.settings.trace.extend(best_trace)
         return best_pred
diff --git a/dspy/predict/refine.py b/dspy/predict/refine.py
index 8d54f8d1..7e421ed2 100644
--- a/dspy/predict/refine.py
+++ b/dspy/predict/refine.py
@@ -1,13 +1,15 @@
-import dspy
-import ujson
 import inspect
 import textwrap
+from typing import Callable, Optional
 
-from .predict import Module
-from dspy.signatures import Signature, InputField, OutputField
+import ujson
 
-from typing import Callable
+import dspy
 from dspy.adapters.chat_adapter import enumerate_fields
+from dspy.predict.predict import Prediction
+from dspy.signatures import InputField, OutputField, Signature
+
+from .predict import Module
 
 
 class OfferFeedback(Signature):
@@ -22,9 +24,7 @@ class OfferFeedback(Signature):
     program_code: str = InputField(desc="The code of the program that we are analyzing")
     modules_defn: str = InputField(desc="The definition of each module in the program, including its I/O")
     program_inputs: str = InputField(desc="The inputs to the program that we are analyzing")
-    program_trajectory: str = InputField(
-        desc="The trajectory of the program's execution, showing each module's I/O"
-    )
+    program_trajectory: str = InputField(desc="The trajectory of the program's execution, showing each module's I/O")
     program_outputs: str = InputField(desc="The outputs of the program that we are analyzing")
     reward_code: str = InputField(desc="The code of the reward function that we are analyzing")
     target_threshold: float = InputField(desc="The target threshold for the reward function")
@@ -39,12 +39,44 @@ class OfferFeedback(Signature):
 
 
 class Refine(Module):
-    def __init__(self, module, N: int, reward_fn: Callable, threshold: float):
+    def __init__(
+        self,
+        module: Module,
+        N: int,
+        reward_fn: Callable[[dict, Prediction], float],
+        threshold: float,
+        fail_count: Optional[int] = None,
+    ):
+        """
+        Refines a module by running it up to N times with different temperatures and returns the best prediction.
+
+        This module runs the provided module multiple times with varying temperature settings and selects
+        either the first prediction that exceeds the specified threshold or the one with the highest reward.
+        If no prediction meets the threshold, it automatically generates feedback to improve future predictions.
+
+
+        Args:
+            module (Module): The module to refine.
+            N (int): The number of times to run the module. must
+            reward_fn (Callable): The reward function.
+            threshold (float): The threshold for the reward function.
+            fail_count (Optional[int], optional): The number of times the module can fail before raising an error
+
+        Example:
+            >>> import dspy
+            >>>
+            >>> qa = dspy.ChainOfThought("question -> answer")
+            >>> def one_word_answer(args, pred):
+            >>>     return 1.0 if len(pred.answer) == 1 else 0.0
+            >>> best_of_3 = dspy.Refine(module=qa, N=3, reward_fn=one_word_answer, threshold=1.0)
+            >>> best_of_3(question="What is the capital of Belgium?").answer
+            >>> # Brussels
+        """
         self.module = module
         self.reward_fn = lambda *args: reward_fn(*args)  # to prevent this from becoming a parameter
         self.threshold = threshold
         self.N = N
-
+        self.fail_count = fail_count or N  # default to N if fail_count is not provided
         self.module_code = inspect.getsource(module.__class__)
         try:
             self.reward_fn_code = inspect.getsource(reward_fn)
@@ -109,14 +141,20 @@ class Refine(Module):
                 advise_kwargs = dict(**modules, **trajectory, **reward, module_names=module_names)
                 # advise_kwargs = {k: ujson.dumps(recursive_mask(v), indent=2) for k, v in advise_kwargs.items()}
                 # only dumps if it's a list or dict
-                advise_kwargs = {k: v if isinstance(v, str) else ujson.dumps(recursive_mask(v), indent=2) for k, v in advise_kwargs.items()}
+                advise_kwargs = {
+                    k: v if isinstance(v, str) else ujson.dumps(recursive_mask(v), indent=2)
+                    for k, v in advise_kwargs.items()
+                }
                 advice = dspy.Predict(OfferFeedback)(**advise_kwargs).advice
                 # print(f"Advice for each module: {advice}")
 
             except Exception as e:
-                print(f"Attempt failed with temperature {t}: {e}")
-
-        dspy.settings.trace.extend(best_trace)
+                print(f"Refine: Attempt failed with temperature {t}: {e}")
+                if idx > self.fail_count:
+                    raise e
+                self.fail_count -= 1
+        if best_trace:
+            dspy.settings.trace.extend(best_trace)
         return best_pred
 
 


Test patch:
diff --git a/tests/predict/test_refine.py b/tests/predict/test_refine.py
new file mode 100644
index 00000000..e63a484a
--- /dev/null
+++ b/tests/predict/test_refine.py
@@ -0,0 +1,82 @@
+from unittest.mock import patch
+
+import pytest
+
+import dspy
+from dspy.predict.predict import Predict
+from dspy.predict.refine import Refine
+from dspy.primitives.prediction import Prediction
+from dspy.utils.dummies import DummyLM
+
+
+class DummyModule(dspy.Module):
+    def __init__(self, signature, forward_fn):
+        super().__init__()
+        self.predictor = Predict(signature)
+        self.forward_fn = forward_fn
+
+    def forward(self, **kwargs) -> Prediction:
+        return self.forward_fn(self, **kwargs)
+
+
+def test_refine_forward_success_first_attempt():
+    lm = DummyLM([{"answer": "Brussels"}, {"answer": "City of Brussels"}, {"answer": "Brussels"}])
+    dspy.settings.configure(lm=lm)
+    module_call_count = [0]
+
+    def count_calls(self, **kwargs):
+        module_call_count[0] += 1
+        return self.predictor(**kwargs)
+
+    reward_call_count = [0]
+
+    def reward_fn(kwargs, pred: Prediction) -> float:
+        reward_call_count[0] += 1
+        # The answer should always be one word.
+        return 1.0 if len(pred.answer) == 1 else 0.0
+
+    predict = DummyModule("question -> answer", count_calls)
+
+    refine = Refine(module=predict, N=3, reward_fn=reward_fn, threshold=1.0)
+    result = refine(question="What is the capital of Belgium?")
+
+    assert result.answer == "Brussels", "Result should be `Brussels`"
+    assert reward_call_count[0] > 0, "Reward function should have been called"
+    assert module_call_count[0] == 3, (
+        "Module should have been called exactly 3 times, but was called %d times" % module_call_count[0]
+    )
+
+
+def test_refine_module_default_fail_count():
+    lm = DummyLM([{"answer": "Brussels"}, {"answer": "City of Brussels"}, {"answer": "Brussels"}])
+    dspy.settings.configure(lm=lm)
+
+    def always_raise(self, **kwargs):
+        raise ValueError("Deliberately failing")
+
+    predict = DummyModule("question -> answer", always_raise)
+
+    refine = Refine(module=predict, N=3, reward_fn=lambda _, __: 1.0, threshold=0.0)
+    with pytest.raises(ValueError):
+        refine(question="What is the capital of Belgium?")
+
+
+def test_refine_module_custom_fail_count():
+    lm = DummyLM([{"answer": "Brussels"}, {"answer": "City of Brussels"}, {"answer": "Brussels"}])
+    dspy.settings.configure(lm=lm)
+    module_call_count = [0]
+
+    def raise_on_second_call(self, **kwargs):
+        if module_call_count[0] < 2:
+            module_call_count[0] += 1
+            raise ValueError("Deliberately failing")
+        return self.predictor(**kwargs)
+
+    predict = DummyModule("question -> answer", raise_on_second_call)
+
+    refine = Refine(module=predict, N=3, reward_fn=lambda _, __: 1.0, threshold=0.0, fail_count=1)
+    with pytest.raises(ValueError):
+        refine(question="What is the capital of Belgium?")
+    assert module_call_count[0] == 2, (
+        "Module should have been called exactly 2 times, but was called %d times" % module_call_count[0]
+    )