feat: allow streaming for sampling results

[jakelorocco]

Give the streaming primitive the instruct-validate-repair retry behavior that SamplingStrategy provides for the non-streaming path, so a streamed generation can retry/repair on validation failure rather than just early-exit.

Two possible interfaces are on the table. One is a dedicated stream_instruct-style entry point, as implemented in Mellea-partials. The other is an optional sampling parameter on stream(), mirroring how validation is already an opt-in parameter; this second option was suggested but not yet evaluated. Both are open; this issue works out both the interface and the mechanism.

Background

This issue began as "allow streaming for functions that return sampling results" (filed 2026-02, referencing the since-closed PR #240). That predates the validated-streaming work and Mellea-partials, arriving independently from the sampling side — it is the origin of the output-shape question below. The body here supersedes that original one-line description with the full scope.

Framing: the streaming counterpart of the existing sampling modes

The non-streaming path already has two sampling-related return modes on act()/instruct(). This issue is their streaming equivalents:

Non-streaming (today)	Streaming (this issue)
act(..., strategy=...) retries/repairs internally, returns the final output	stream(...) gains retry/repair — the control-flow half
act(..., return_sampling_results=True) returns a SamplingResult with every attempt (passed and failed)	streaming surfaces those attempts as they happen — the output-shape half

So this issue has two dimensions: how the retry loop drives a caller's async for (control flow), and what a caller sees when they want the full multi-attempt view rather than just the final stream (output shape). The two map to strategy= and return_sampling_results=True respectively.

Intent (from Mellea-partials)

Mellea-partials provides a stream_instruct: a streaming counterpart to instruct() that runs the whole loop —

1. Generate a streaming MOT.
2. Chunk + quick-check each chunk during streaming (per-chunk requirements).
3. Quick-check failure → cancel, retry immediately (regenerate, without strategy.repair()).
4. Full-output validation once chunks pass (full-output requirements).
5. Full-validation failure → strategy.repair() produces an updated instruction/context, retry.
6. Loop up to strategy.loop_budget; emit events throughout.

Mellea-partials predates the current stream() and was built against the original stream_with_chunking, so its structure reflects that older design. Use it as the reference for what streaming sampling needs to do, not as a design to port — the goal is for this to fall out of the new stream() cleanly.

How the current codebase does sampling (non-streaming)

• BaseSamplingStrategy.sample(action, context, backend, requirements, ...) owns the entire generate/validate/repair loop and returns a SamplingResult. It calls generation, awaits the full result, validates, repairs, loops up to loop_budget.
• act()/aact() take a strategy= param (default RejectionSamplingStrategy(loop_budget=2)) and delegate the loop to sample().

The core tension to resolve

sample() is built to own the generate-and-validate cycle: it drives generation, awaits the full output, validates, repairs, loops. stream()'s whole point is that the caller drives consumption via async for. These are opposed control-flow models:

• sample() wants to be in control of the loop.
• stream() hands the loop to the caller.

This tension holds regardless of the interface chosen (param on stream() vs. a wrapper): either way, the open work is how a caller-driven async for interoperates with a loop-owning sampling strategy — specifically, what happens to iteration when a streamed attempt fails validation and the strategy wants to repair and retry. The interface question (param vs. wrapper) is the smaller, later decision; this control-flow mechanism is the substance.

Open questions

1. Who drives the retry loop? Does stream() internally call strategy.sample() (letting it own the loop, with streaming as an implementation detail inside), or is the sampling loop re-expressed for the streaming/caller-driven model?
2. What does the caller see across a retry? When attempt 1 fails and attempt 2 begins, does the single async for transparently continue into attempt 2's chunks? Restart? Surface a retry event and re-yield from the top? Mellea-partials emitted RetryEvent and continued; the new model must define this explicitly.
3. Two retry triggers or one? Mellea-partials distinguished quick-check failure (regenerate, no repair) from full-validation failure (repair then retry). Does the new model keep both, or unify?
4. Does this reuse BaseSamplingStrategy or bypass it? Phase 1's primitive deliberately bypassed BaseSamplingStrategy (callers own the retry loop, per epic feat: streaming validation — per-chunk requirement checking with early exit #891). Reintroducing sampling here means deciding whether to re-integrate with sample() or keep a parallel streaming-specific loop.
5. Relationship to full-output validation. stream() already runs full-output validate() on natural completion (refactor: replace stream_with_chunking with single-task stream() #1440). Sampling adds the repair-and-retry layer on top of that — confirm the layering (validation is the check; sampling is what acts on a failed check).
6. Output shape — what does streaming a multi-attempt run yield? The non-streaming return_sampling_results=True returns a SamplingResult holding every attempt (passed and failed) after completion. Its streaming equivalent has to define what a caller sees live across attempts: a flat chunk stream of only the final/winning attempt, or a richer stream that surfaces each attempt's chunks, validation, and retry as they happen. This is distinct from Q2 (Q2 is "does the async for continue"; this is "what typed view of the whole run is available"). May need a streaming counterpart to SamplingResult, or Streamer fields exposing attempts.

Notes

• Depends on refactor: replace stream_with_chunking with single-task stream() #1440 (stream() primitive) only; independent of per-validator chunking (feat: per-requirement chunking strategies #1441) and stream_parsed_repr (feat(core): stream_parsed_repr — streaming parsed representation on the MOT #1442).
• Interface is undecided: a param on stream() vs. a stream_instruct-style wrapper (the Mellea-partials form). Both face the same control-flow mechanism question, which is the real design work.

[ajbozarth]

I spent some time working out an actual design from the above issue outline. I'll paste it below for reference. I intend to start executing it either later today or Monday.

#403 — Streaming sampling: design decisions

This supplements the #403 issue body with the design we settled after investigation. It maps a decision onto each open question in the body, states the core constraint that drives the whole shape, and lists what's deliberately deferred to implementation planning. It is written to be pasted as an issue comment and to seed an implementation-planning session.

Summary

• Extend the existing stream() primitive with an opt-in strategy: SamplingStrategy | None = None parameter. No separate stream_instruct entry point (kept only as a fallback if implementation hits something unreconcilable).
• The retry/repair loop lives in the stream() layer, not in sampling. stream() reuses the strategy's pure decision pieces (repair(), select_from_failure(), loop_budget) and adds one new streaming-native piece (stream_repair). It never calls sample()/_sample().
• Two failure triggers: a mid-stream (partial) failure repairs via the new stream_repair; a completed-but-failed final validate() repairs via the existing repair().
• Every strategy is accepted. Repair-family strategies stream live; strategies that override the sampling loop (MajorityVoting*, BudgetForcing, SOFAI) degrade to a blocking pass-through. No capability flag.
• Multi-attempt results are exposed through a new streaming-native sibling to SamplingResult, held as an optional sub-object on Streamer.

The core constraint: why sample() cannot be reused

sample()/_sample() own the generate_from_context call (and then avalue() the full result before validating). stream() must own that same call to stream it (STREAM=True, iterating the model output thunk). Two owners of one generation cannot coexist, so stream() can never delegate into sample() to borrow a strategy's algorithm.

PR #1584's sample/_sample split lifts the lifecycle out (sampling id, loop start/end hooks) but leaves generation inside _sample, so it does not reconcile this — the conflict is identical before and after #1584. Reusing sample() would require a further refactor that lifts generation out of _sample (so a streaming generator could be injected), which is out of scope for #403.

Consequence: stream() can only reuse the strategy's pure decision functions — the ones that don't touch generation — and must re-express the loop itself. This is what makes the whole design fall out.

Decisions (mapped to the issue's open questions)

Q1 — Who drives the retry loop? → the stream() layer

The loop is re-expressed in stream(). It owns generation and validation and calls the strategy only for its pure decision pieces. This also matches the phase-1 (epic #891) stance that streaming owns its retry loop and uses the streaming event vocabulary for observability rather than the sampling hooks.

Q2 — What does the caller see across a retry? → one async for, continued flat

The single async for continues transparently across attempts. Because an attempt streams its chunks live before it can fail, the caller sees the failed attempt's partial chunks; a RetryEvent is emitted before each re-attempt (lighting up the reserved-but-unemitted event), and the attempt field on subsequent events increments.

Inherent caveat to document: a consumer rendering live text will see discarded attempts back-to-back unless it watches for RetryEvent (available via stream(..., as_events=True)) to reset its display. This is intrinsic to live streaming + retry, not a flaw to hide.

Q3 — Two retry triggers or one? → two, because repair() is only well-defined at completion

• Completed attempt, final validate() fails → the existing repair(), which needs a ComputedModelOutputThunk + full ValidationResults (both available at natural completion).
• Broke early mid-stream → a new stream_repair, which accepts the partial result (cancelled thunk, PartialValidationResult failures) that repair() cannot consume.

stream_repair is a real per-strategy method with a safe default (regenerate, i.e. rejection-style). Feedback strategies (RepairTemplateStrategy, MultiTurnStrategy) override it to fold the partial's failure reason into the next instruction/turn. (Note: the codebase already unified Mellea-partials' quick_check_requirements / requirements into a single requirements list — each requirement does per-chunk stream_validate and a final validate() — so this is two triggers, not two requirement lists. We are not porting Mellea-partials' per-chunk quick_repair callback; repair is strategy-driven.)

Q4 — Reuse BaseSamplingStrategy or bypass it? → hybrid: reuse the pure pieces, never call sample()

stream() reuses repair(), select_from_failure(), and loop_budget (none of which own generation), and adds stream_repair. It does not call sample()/_sample() (see the core constraint). This is the only sampling surface composable with a stream.

Q5 — Relationship to full-output validation → validation is the check, sampling acts on it

Validation is unchanged as the check: per-chunk stream_validate early-exit plus a final validate() on natural completion. Sampling is what acts on a failed check (repair + retry). With strategy=None, behavior is exactly today's — the final validate() result is recorded but not acted on — so streaming sampling is strictly opt-in and backward-compatible.

Q6 — Output shape (multi-attempt view) → a SamplingResult sibling on Streamer

Streamer gains an optional sub-object (only populated when a strategy was used) that is the streaming-native sibling of SamplingResult, holding the whole run: per-attempt records, the selected index, and overall success. Streamer's existing top-level fields (full_text, mot, final_validations, streaming_failures) continue to describe the winning attempt, so single-attempt consumers are unaffected.

This mirrors the requirement precedent — a full/streaming result pair (ValidationResult ↔ PartialValidationResult) — giving sampling the same pairing (SamplingResult ↔ the new sibling). The sibling exists precisely because streaming attempts are not uniform: an attempt can be won, completed-but-failed (full ValidationResults, like SamplingResult), or broken-early (partial text + PartialValidationResult failures, which SamplingResult can't represent). Keeping this in the sibling leaves SamplingResult untouched.

Interface → strategy= parameter on stream()

stream() gains strategy: SamplingStrategy | None = None, mirroring act(strategy=…) and sitting beside the already-opt-in requirements=. It still returns a Streamer / EventStreamer; the new sampling sub-object is simply populated when a strategy was passed. This also keeps the module dependency one-way (streaming → sampling, via the pure pieces), avoiding a cycle.

Strategy compatibility & graceful degradation

Every strategy is accepted by stream(); support varies in degree (the same principle as requirements that return "unknown" throughout streaming and do their real work at the final validate()).

• Repair-family (strategies that use the standard repair loop): stream live via the stream() loop + stream_repair.
• Loop-overriders (MajorityVoting*, BudgetForcing, SOFAI): their algorithm is welded to generation inside their overridden sample()/_sample(), so it cannot run inside a caller-owned stream. They degrade to a blocking pass-through — stream() runs the strategy's real sample() to completion and emits the result. Correct algorithm, no live chunks. (Consider surfacing a signal that the run was non-live so this isn't silently misleading.)

There is no supports_streaming flag. The distinction between the two is read from the strategy's existing structure — a loop-overrider already declares its difference by overriding the sampling loop — not from a new boolean. This avoids the inheritance foot-gun where MajorityVoting/BudgetForcing (which subclass RejectionSamplingStrategy) would otherwise inherit repair-family streaming behavior.

Why the requirement analogy holds for one and not the other: validate() is called by the generation driver (composable); sample() is a generation driver (not composable with another one).

Dependencies & coordination

• refactor: replace stream_with_chunking with single-task stream() #1440 (the stream() primitive) — merged; this builds directly on it.
• fix(sampling): sampling span missing for strategies that override sam… #1584 / fix(sampling): sampling span missing for strategies that override sample() #1487 — coordination, not a hard dependency. Both feat: allow streaming for sampling results #403 and fix(sampling): sampling span missing for strategies that override sam… #1584 modify SamplingStrategy (feat: allow streaming for sampling results #403 adds stream_repair; fix(sampling): sampling span missing for strategies that override sam… #1584 restructures sample/_sample), so they need merge coordination, and the repair-family-vs-overrider detection interacts with fix(sampling): sampling span missing for strategies that override sam… #1584's _sample split. But the core reconciliation (generation ownership) is identical before and after fix(sampling): sampling span missing for strategies that override sam… #1584, so feat: allow streaming for sampling results #403's design does not block on it.
• Independent of feat: per-requirement chunking strategies #1441 (per-validator chunking, merged) and feat(core): stream_parsed_repr — streaming parsed representation on the MOT #1442 (stream_parsed_repr).

Deferred to implementation planning

These do not change the design but need answers when building it:

• Un-wiring the inner early-exit. Today _drive's mid-stream failure is a bare return → _finalize → mot.aclose() (terminal, cancels generation). Retry requires routing that early-exit to a retry decision instead of finalization, with a fresh model output thunk per attempt (the single-consumer guard forbids re-iteration). The existing single-attempt drive becomes the "one attempt" unit wrapped by an outer attempt loop.
• Exact detection of repair-family vs loop-overrider (e.g. whether the strategy's _sample resolves to the standard base implementation), consistent with fix(sampling): sampling span missing for strategies that override sam… #1584's shape — no new flag.
• The new result type: its name, whether it offers a projection to a real SamplingResult (completed attempts only) for parity with return_sampling_results, and whether it is populated progressively during the run or finalized at the end.
• select_from_failure inputs now include partial (early-broken) attempts, so a strategy that ranks failures non-trivially needs its inputs normalized; the repair-family's is trivial (return first).
• stream_repair default + overrides for each repair-family strategy.