Repository navigation
feat: allow streaming for sampling results #403
Description
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 —
- Generate a streaming MOT.
- Chunk + quick-check each chunk during streaming (per-chunk requirements).
- Quick-check failure → cancel, retry immediately (regenerate, without
strategy.repair()). - Full-output validation once chunks pass (full-output requirements).
- Full-validation failure →
strategy.repair()produces an updated instruction/context, retry. - 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 aSamplingResult. It calls generation, awaits the full result, validates, repairs, loops up toloop_budget.act()/aact()take astrategy=param (defaultRejectionSamplingStrategy(loop_budget=2)) and delegate the loop tosample().
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
- Who drives the retry loop? Does
stream()internally callstrategy.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? - What does the caller see across a retry? When attempt 1 fails and attempt 2 begins, does the single
async fortransparently continue into attempt 2's chunks? Restart? Surface a retry event and re-yield from the top?Mellea-partialsemittedRetryEventand continued; the new model must define this explicitly. - Two retry triggers or one?
Mellea-partialsdistinguished quick-check failure (regenerate, no repair) from full-validation failure (repair then retry). Does the new model keep both, or unify? - Does this reuse
BaseSamplingStrategyor bypass it? Phase 1's primitive deliberately bypassedBaseSamplingStrategy(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 withsample()or keep a parallel streaming-specific loop. - Relationship to full-output validation.
stream()already runs full-outputvalidate()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). - Output shape — what does streaming a multi-attempt run yield? The non-streaming
return_sampling_results=Truereturns aSamplingResultholding 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 theasync forcontinue"; this is "what typed view of the whole run is available"). May need a streaming counterpart toSamplingResult, orStreamerfields 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) andstream_parsed_repr(feat(core): stream_parsed_repr — streaming parsed representation on the MOT #1442). - Interface is undecided: a param on
stream()vs. astream_instruct-style wrapper (theMellea-partialsform). Both face the same control-flow mechanism question, which is the real design work.
Activity
ajbozarth commentedon Sep 11, 2026
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
stream()primitive with an opt-instrategy: SamplingStrategy | None = Noneparameter. No separatestream_instructentry point (kept only as a fallback if implementation hits something unreconcilable).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 callssample()/_sample().stream_repair; a completed-but-failed finalvalidate()repairs via the existingrepair().MajorityVoting*,BudgetForcing,SOFAI) degrade to a blocking pass-through. No capability flag.SamplingResult, held as an optional sub-object onStreamer.The core constraint: why
sample()cannot be reusedsample()/_sample()own thegenerate_from_contextcall (and thenavalue()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, sostream()can never delegate intosample()to borrow a strategy's algorithm.PR #1584's
sample/_samplesplit 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. Reusingsample()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()layerThe 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 flatThe single
async forcontinues transparently across attempts. Because an attempt streams its chunks live before it can fail, the caller sees the failed attempt's partial chunks; aRetryEventis emitted before each re-attempt (lighting up the reserved-but-unemitted event), and theattemptfield 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 viastream(..., 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 completionvalidate()fails → the existingrepair(), which needs aComputedModelOutputThunk+ fullValidationResults (both available at natural completion).stream_repair, which accepts the partial result (cancelled thunk,PartialValidationResultfailures) thatrepair()cannot consume.stream_repairis 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/requirementsinto a singlerequirementslist — each requirement does per-chunkstream_validateand a finalvalidate()— so this is two triggers, not two requirement lists. We are not porting Mellea-partials' per-chunkquick_repaircallback; repair is strategy-driven.)Q4 — Reuse
BaseSamplingStrategyor bypass it? → hybrid: reuse the pure pieces, never callsample()stream()reusesrepair(),select_from_failure(), andloop_budget(none of which own generation), and addsstream_repair. It does not callsample()/_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_validateearly-exit plus a finalvalidate()on natural completion. Sampling is what acts on a failed check (repair + retry). Withstrategy=None, behavior is exactly today's — the finalvalidate()result is recorded but not acted on — so streaming sampling is strictly opt-in and backward-compatible.Q6 — Output shape (multi-attempt view) → a
SamplingResultsibling onStreamerStreamergains an optional sub-object (only populated when a strategy was used) that is the streaming-native sibling ofSamplingResult, 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 (fullValidationResults, likeSamplingResult), or broken-early (partial text +PartialValidationResultfailures, whichSamplingResultcan't represent). Keeping this in the sibling leavesSamplingResultuntouched.Interface →
strategy=parameter onstream()stream()gainsstrategy: SamplingStrategy | None = None, mirroringact(strategy=…)and sitting beside the already-opt-inrequirements=. It still returns aStreamer/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 finalvalidate()).stream()loop +stream_repair.MajorityVoting*,BudgetForcing,SOFAI): their algorithm is welded to generation inside their overriddensample()/_sample(), so it cannot run inside a caller-owned stream. They degrade to a blocking pass-through —stream()runs the strategy's realsample()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_streamingflag. 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 whereMajorityVoting/BudgetForcing(which subclassRejectionSamplingStrategy) 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
stream()primitive) — merged; this builds directly on it.SamplingStrategy(feat: allow streaming for sampling results #403 addsstream_repair; fix(sampling): sampling span missing for strategies that override sam… #1584 restructuressample/_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_samplesplit. 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.stream_parsed_repr).Deferred to implementation planning
These do not change the design but need answers when building it:
_drive's mid-stream failure is a barereturn→_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._sampleresolves to the standard base implementation), consistent with fix(sampling): sampling span missing for strategies that override sam… #1584's shape — no new flag.SamplingResult(completed attempts only) for parity withreturn_sampling_results, and whether it is populated progressively during the run or finalized at the end.select_from_failureinputs 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_repairdefault + overrides for each repair-family strategy.ajbozarth commentedon Sep 15, 2026
@jakelorocco I probably won't return to implement this until later in the week or start of next week if you wanted to review my above design note before I actually implemented it
jakelorocco commentedon Sep 21, 2026
I think this design looks good; two thoughts: