Output Refinement with BestOfN and Refine

This example wraps a ChainOfThought program in BestOfN and Refine to score its outputs against a reward function and keep the best result. It covers reward functions over predictions, early stopping with failCount, and a reward that calls a language model to judge the output.

For a short introduction to these two modules, see Modules.

Scoring with a reward function

val bestOf3 = BestOfN(
  module = qa,
  n = AttemptCount(3),
  rewardFn = (_, pred) => oneWord(pred.output.answer),
  threshold = 1.0
)

def call(question: String)(using RuntimeContext): Either[DspyError, String] =
  bestOf3((question = question)).map(_.output.answer)

BestOfN samples n completions from the inner program in parallel and returns the one with the highest reward. The reward is (input, prediction) => Double. Because the inner program is a ChainOfThought("question -> answer"), its prediction carries the augmented output (reasoning, answer), so the reward reads pred.output.answer. The threshold lets sampling stop early: once a completion scores at or above it, that completion is returned without drawing the rest. Here oneWord returns 1.0 only when the answer is a single word.

Stopping early on failures

val bestOf3 = BestOfN(
  module = qa,
  n = AttemptCount(3),
  rewardFn = (_, pred) => oneWord(pred.output.answer),
  threshold = 1.0,
  failCount = Some(FailureCount(1))
)

failCount caps how many completions may score below the threshold before sampling gives up and returns the best result seen so far. With failCount = Some(1), a single sub-threshold completion ends the search. This bounds the number of model calls when the program is unlikely to reach the threshold. Refine accepts the same parameter and refines sequentially, feeding each attempt's score into the next try.

Judging with a language model

private val judge = ChainOfThought(Signature.of[FactualityJudge])

def call(question: String)(using ctx: RuntimeContext): Either[DspyError, String] =
  def reward(answer: String): Double =
    judge((statement = answer)).map(r => if r.output.is_factual then 1.0 else 0.0).getOrElse(0.0)
  val refinedQa = Refine(
    module = qa,
    n = AttemptCount(3),
    rewardFn = (_, pred) => reward(pred.output.answer),
    threshold = 1.0
  )
  refinedQa((question = question)).map(_.output.answer)

A reward function can itself call a language model. Here judge is a ChainOfThought over a FactualityJudge signature that maps a statement to a boolean is_factual. The reward runs the judge on the candidate answer and returns 1.0 when it is judged factual. The reward signature is (input, prediction) only, with no implicit context threaded in, so the judge call captures the ambient RuntimeContext from the enclosing call. Refine then drives the inner program until the judged reward reaches the threshold.

Running it

OPENAI_API_KEY=sk-... sbt "examples/runMain dspy4s.examples.tutorials.output_refinement.bestOfNAndRefineMain"

Notes

Full source: BestOfNAndRefine.scala