Introduce dspy.Refine, dspy.BestOfN and fix thread safety for dspy.ChainOfThoughtWithHint (#1959)

* Add DSPy Refine

* Introduce BestOfN and make Refine & CoTWithHint thread safe

* Improve Refine implementation

---------

Co-authored-by: Omar Khattab <[email protected]>
Co-authored-by: Omar Khattab <[email protected]>

Author: chenmoneygithub

dspy/predict/__init__.py

@@ -1,21 +1,26 @@
 from dspy.predict.aggregation import majority
+from dspy.predict.best_of_n import BestOfN
 from dspy.predict.chain_of_thought import ChainOfThought
 from dspy.predict.chain_of_thought_with_hint import ChainOfThoughtWithHint
 from dspy.predict.knn import KNN
 from dspy.predict.multi_chain_comparison import MultiChainComparison
 from dspy.predict.predict import Predict
 from dspy.predict.program_of_thought import ProgramOfThought
-from dspy.predict.react import ReAct
+from dspy.predict.react import ReAct, Tool
+from dspy.predict.refine import Refine
 from dspy.predict.parallel import Parallel
 
 __all__ = [
     "majority",
+    "BestOfN",
     "ChainOfThought",
     "ChainOfThoughtWithHint",
     "KNN",
     "MultiChainComparison",
     "Predict",
     "ProgramOfThought",
     "ReAct",
+    "Refine",
+    "Tool",
     "Parallel",
 ]

dspy/predict/best_of_n.py

@@ -0,0 +1,43 @@
+import dspy
+
+from .predict import Module
+from typing import Callable
+
+
+class BestOfN(Module):
+    def __init__(self, module, N: int, reward_fn: Callable, threshold: float):
+        self.module = module
+        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 = list(dict.fromkeys(temps))[:self.N]
+        best_pred, best_trace, best_reward = None, None, -float("inf")
+
+        for t in temps:
+            lm_ = lm.copy(temperature=t)
+            mod = self.module.deepcopy()
+            mod.set_lm(lm_)
+
+            try:
+                with dspy.context(trace=[]):
+                    pred = mod(**kwargs)
+                    trace = dspy.settings.trace.copy()
+
+                    # 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

dspy/predict/chain_of_thought_with_hint.py

@@ -1,9 +1,6 @@
 import dspy
-
 from .predict import Module
 
-# TODO: FIXME: Insert this right before the *first* output field. Also rewrite this to use the new signature system.
-
 class ChainOfThoughtWithHint(Module):
     def __init__(self, signature, rationale_type=None, **config):
         self.signature = dspy.ensure_signature(signature)
@@ -19,16 +16,11 @@ def forward(self, **kwargs):
             kwargs[last_key] = str(kwargs[last_key]) + hint
 
             # Run CoT then update the trace with original kwargs, i.e. without the hint.
-            pred = self.module(**kwargs)
-            this_trace = dspy.settings.trace[-1]
-            dspy.settings.trace[-1] = (this_trace[0], original_kwargs, this_trace[2])
+            with dspy.context(trace=[]):
+                pred = self.module(**kwargs)
+                this_trace = dspy.settings.trace[-1]
+
+            dspy.settings.trace.append((this_trace[0], original_kwargs, this_trace[2]))
             return pred
 
         return self.module(**kwargs)
-
-
-"""
-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.
-"""