diff --git a/dspy/predict/best_of_n.py b/dspy/predict/best_of_n.py index 70b59e0e3a..4b359d7358 100644 --- a/dspy/predict/best_of_n.py +++ b/dspy/predict/best_of_n.py @@ -52,6 +52,7 @@ def forward(self, **kwargs): start = lm.kwargs.get("rollout_id", 0) rollout_ids = [start + i for i in range(self.N)] best_pred, best_trace, best_reward = None, None, -float("inf") + fail_count = self.fail_count for idx, rid in enumerate(rollout_ids): lm_ = lm.copy(rollout_id=rid, temperature=1.0) @@ -74,9 +75,9 @@ def forward(self, **kwargs): except Exception as e: print(f"BestOfN: Attempt {idx + 1} failed with rollout id {rid}: {e}") - if idx > self.fail_count: + fail_count -= 1 + if fail_count < 0: raise e - self.fail_count -= 1 if best_trace: dspy.settings.trace.extend(best_trace) diff --git a/dspy/predict/refine.py b/dspy/predict/refine.py index da49c7ff6b..86dd7c89a5 100644 --- a/dspy/predict/refine.py +++ b/dspy/predict/refine.py @@ -102,6 +102,7 @@ def forward(self, **kwargs): best_pred, best_trace, best_reward = None, None, -float("inf") advice = None adapter = dspy.settings.adapter or dspy.ChatAdapter() + fail_count = self.fail_count for idx, rid in enumerate(rollout_ids): lm_ = lm.copy(rollout_id=rid, temperature=1.0) @@ -169,9 +170,9 @@ def __call__(self, lm, lm_kwargs, signature, demos, inputs): except Exception as e: print(f"Refine: Attempt failed with rollout id {rid}: {e}") - if idx > self.fail_count: + fail_count -= 1 + if fail_count < 0: raise e - self.fail_count -= 1 if best_trace: dspy.settings.trace.extend(best_trace) return best_pred diff --git a/tests/predict/test_best_of_n.py b/tests/predict/test_best_of_n.py index 684bc0629d..2e116ae458 100644 --- a/tests/predict/test_best_of_n.py +++ b/tests/predict/test_best_of_n.py @@ -55,8 +55,12 @@ def always_raise(self, **kwargs): predict = DummyModule("question -> answer", always_raise) best_of_n = BestOfN(module=predict, N=3, reward_fn=lambda _, __: 1.0, threshold=0.0) - with pytest.raises(ValueError): - best_of_n(question="What is the capital of Belgium?") + # fail_count defaults to N: up to N failures are tolerated, so all N failing + # rollouts are absorbed and the best prediction seen so far (None) is returned. + result = best_of_n(question="What is the capital of Belgium?") + assert result is None + # the instance budget must not leak across forward() calls + assert best_of_n.fail_count == 3 def test_refine_module_custom_fail_count(): @@ -78,3 +82,73 @@ def raise_on_second_call(self, **kwargs): assert module_call_count[0] == 2, ( "Module should have been called exactly 2 times, but was called %d times" % module_call_count[0] ) + + +def test_best_of_n_fail_count_tolerates_single_late_failure(): + # threshold=2.0 is unreachable (reward capped at 1.0), so all N rollouts run and + # the single failure on the final rollout is actually reached. With + # fail_count=2 a single late failure must be tolerated (not re-raised). + dspy.configure(lm=DummyLM([{"answer": "Brussels"}] * 4)) + + rollout_count = [0] + + def succeed_four_then_fail(self, **kwargs): + rollout_count[0] += 1 + if rollout_count[0] <= 4: + return self.predictor(**kwargs) + raise ValueError("transient failure on rollout %d" % rollout_count[0]) + + predict = DummyModule("question -> answer", succeed_four_then_fail) + best_of_n = BestOfN(module=predict, N=5, reward_fn=lambda _, __: 1.0, threshold=2.0, fail_count=2) + + result = best_of_n(question="What is the capital of Belgium?") + + assert result.answer == "Brussels" + assert best_of_n.fail_count == 2 + + +def test_best_of_n_fail_count_still_raises_when_exceeded(): + # A single late failure is tolerated with fail_count=1, but the *second* + # late failure exceeds the budget and must be re-raised. + dspy.configure(lm=DummyLM([{"answer": "Brussels"}] * 4)) + + rollout_count = [0] + + def succeed_four_then_fail_twice(self, **kwargs): + rollout_count[0] += 1 + if rollout_count[0] <= 4: + return self.predictor(**kwargs) + raise ValueError("failure on rollout %d" % rollout_count[0]) + + predict = DummyModule("question -> answer", succeed_four_then_fail_twice) + best_of_n = BestOfN(module=predict, N=6, reward_fn=lambda _, __: 1.0, threshold=2.0, fail_count=1) + + with pytest.raises(ValueError, match="failure on rollout 6"): + best_of_n(question="What is the capital of Belgium?") + assert best_of_n.fail_count == 1 + + +def test_best_of_n_fail_count_does_not_leak_across_calls(): + # Each forward() call fails twice then succeeds; with fail_count=2 both + # leading failures must be tolerated on *every* call (no cross-call leak). + dspy.configure(lm=DummyLM([{"answer": "Brussels"}] * 20)) + + state = {"rollout_in_call": 0} + + def fail_twice_per_call(self, **kwargs): + if state["rollout_in_call"] < 2: + state["rollout_in_call"] += 1 + raise ValueError("leading failure #%d" % state["rollout_in_call"]) + return self.predictor(**kwargs) + + predict = DummyModule("question -> answer", fail_twice_per_call) + best_of_n = BestOfN(module=predict, N=5, reward_fn=lambda _, __: 1.0, threshold=2.0, fail_count=2) + + r1 = best_of_n(question="Q?") + assert r1.answer == "Brussels" + assert best_of_n.fail_count == 2 + + state["rollout_in_call"] = 0 + r2 = best_of_n(question="Q?") + assert r2.answer == "Brussels" + assert best_of_n.fail_count == 2 diff --git a/tests/predict/test_refine.py b/tests/predict/test_refine.py index 7a5c2ece13..db08b4ec63 100644 --- a/tests/predict/test_refine.py +++ b/tests/predict/test_refine.py @@ -55,8 +55,12 @@ def always_raise(self, **kwargs): 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?") + # fail_count defaults to N: up to N failures are tolerated, so all N failing + # rollouts are absorbed and the best prediction seen so far (None) is returned. + result = refine(question="What is the capital of Belgium?") + assert result is None + # the instance budget must not leak across forward() calls + assert refine.fail_count == 3 def test_refine_module_custom_fail_count(): @@ -78,3 +82,82 @@ def raise_on_second_call(self, **kwargs): assert module_call_count[0] == 2, ( "Module should have been called exactly 2 times, but was called %d times" % module_call_count[0] ) + + +# A single LM response that satisfies both the `question -> answer` predictor +# (which extracts `answer`) and the `OfferFeedback` predictor (which extracts +# `discussion` and `advice`). Using one combined response makes the LM answer +# order irrelevant, so consumption stays correct across multiple forward() calls +# (where each call consumes an odd number of responses: predictor + advice pairs +# plus a final predictor without advice). +ANSWER_AND_ADVICE = {"answer": "Brussels", "discussion": "no blame", "advice": {"predictor": "N/A"}} + + +def test_refine_fail_count_tolerates_single_late_failure(): + # threshold=2.0 is unreachable (reward capped at 1.0), so all N rollouts run and + # the single failure on the final rollout is actually reached. With + # fail_count=2 a single late failure must be tolerated (not re-raised). + dspy.configure(lm=DummyLM([ANSWER_AND_ADVICE] * 8)) + + rollout_count = [0] + + def succeed_four_then_fail(self, **kwargs): + rollout_count[0] += 1 + if rollout_count[0] <= 4: + return self.predictor(**kwargs) + raise ValueError("transient failure on rollout %d" % rollout_count[0]) + + predict = DummyModule("question -> answer", succeed_four_then_fail) + refine = Refine(module=predict, N=5, reward_fn=lambda _, __: 1.0, threshold=2.0, fail_count=2) + + result = refine(question="What is the capital of Belgium?") + + assert result.answer == "Brussels" + assert refine.fail_count == 2 + + +def test_refine_fail_count_still_raises_when_exceeded(): + # A single late failure is tolerated with fail_count=1, but the *second* + # late failure exceeds the budget and must be re-raised. + dspy.configure(lm=DummyLM([ANSWER_AND_ADVICE] * 8)) + + rollout_count = [0] + + def succeed_four_then_fail_twice(self, **kwargs): + rollout_count[0] += 1 + if rollout_count[0] <= 4: + return self.predictor(**kwargs) + raise ValueError("failure on rollout %d" % rollout_count[0]) + + predict = DummyModule("question -> answer", succeed_four_then_fail_twice) + refine = Refine(module=predict, N=6, reward_fn=lambda _, __: 1.0, threshold=2.0, fail_count=1) + + with pytest.raises(ValueError, match="failure on rollout 6"): + refine(question="What is the capital of Belgium?") + assert refine.fail_count == 1 + + +def test_refine_fail_count_does_not_leak_across_calls(): + # Each forward() call fails twice then succeeds; with fail_count=2 both + # leading failures must be tolerated on *every* call (no cross-call leak). + dspy.configure(lm=DummyLM([ANSWER_AND_ADVICE] * 20)) + + state = {"rollout_in_call": 0} + + def fail_twice_per_call(self, **kwargs): + if state["rollout_in_call"] < 2: + state["rollout_in_call"] += 1 + raise ValueError("leading failure #%d" % state["rollout_in_call"]) + return self.predictor(**kwargs) + + predict = DummyModule("question -> answer", fail_twice_per_call) + refine = Refine(module=predict, N=5, reward_fn=lambda _, __: 1.0, threshold=2.0, fail_count=2) + + r1 = refine(question="Q?") + assert r1.answer == "Brussels" + assert refine.fail_count == 2 + + state["rollout_in_call"] = 0 + r2 = refine(question="Q?") + assert r2.answer == "Brussels" + assert refine.fail_count == 2