feat: add sampling-strategy retry/repair to stream()

[ajbozarth]

Issue

Fixes #403

Description

stream() previously aborted on a failed requirement with no recourse. This adds the retry/repair behavior of non-streaming sampling to streaming by accepting a SamplingStrategy: a failed attempt is repaired and retried up to loop_budget, and the winning attempt is projected onto the Streamer's terminal state (full_text, mot, final_validations). Requirements set on the strategy are merged with those passed to stream(). With strategy=None, behavior is unchanged.

Two triggers drive a retry:

• a mid-stream "fail" cancels the attempt at the failing chunk and repairs via the new stream_repair();
• a completed attempt whose final validate() fails repairs via the existing repair().

Per-attempt history is recorded on Streamer.attempts (a new StreamAttempt dataclass), and a RetryEvent is emitted before each re-attempt. CompletedEvent gains attempts_used.

stream_repair() is a non-abstract method on BaseSamplingStrategy that raises NotImplementedError by default (custom strategies stay instantiable; streaming repair is opt-in); RejectionSamplingStrategy, RepairTemplateStrategy, and MultiTurnStrategy implement it. repair()/stream_repair() past_results is widened to Sequence[ComputedModelOutputThunk | None] so streaming can pass its per-attempt results (None for early-broken attempts) without narrowing existing callers.

Testing

• Tests added to the respective file if code was changed
• New code has 100% coverage if code was added
• Ensure existing tests and github automation passes (a maintainer will kick off the github automation when the rest of the PR is populated)

Attribution

• AI coding assistants used

---

Adding a new component, requirement, sampling strategy, or tool?

If your PR adds or modifies one of the types below, check the matching box. A checklist of type-specific review items will be posted as a comment.

• Component
• Requirement
• Sampling Strategy
• Tool

NOTE: Please ensure you have an issue that has been acknowledged by a core contributor and routed you to open a pull request against this repository. Otherwise, please open an issue before continuing with this pull request.

[planetf1]

The main ones are the requirement merge, abandoned streams counting as sampling failures, and the repair() signature change. The rest are smaller.

[planetf1]

Was the union here deliberate? It's different from what the non-streaming path does with the same strategy.

_merge_requirements (mellea/core/sampling.py:385-403) uses the strategy's requirements instead of the per-call ones whenever the strategy has any, whereas stream() validates both. So with RejectionSamplingStrategy(requirements=[a]) and requirements=[b], instruct() checks a only and stream() checks a and b.

The docstrings disagree too. stream() says per-call requirements are "merged with any set on strategy", but BaseSamplingStrategy and SamplingStrategy both say strategy requirements "override per-call requirements". test_strategy_requirements_merge_with_call_requirements pins the union, so I'm guessing it was a choice.

Could you clarify which behaviour is intended when both are set, and have the docstrings on both sides say so? Otherwise anyone moving an instruct() call over to stream() is in for a surprise.

[jakelorocco]

We should keep the existing instruct / sampling strategy behavior.

[ajbozarth]

as noted and discussed in scrum this was intentional and addresses a bug in our code that I will fix elsewhere in a follow up #1698

[planetf1]

What should the sampling metrics record when a caller abandons a strategy stream? At the moment it counts as a failure.

On a break, an early aclose() or a cancellation, _finalize gets exception=None, so this produces sampling_success=False and record_streaming_outcome (mellea/telemetry/metrics_plugins.py:497-504) adds it to mellea.sampling.failures. I tried a loop_budget=3 stream whose output would have passed. Drained, it reports sampling_success=True. With a break after the first chunk it reports False, and its one attempt isn't selected and has an empty full_text.

The non-streaming path doesn't count this. A cancelled act() reaches SAMPLING_LOOP_END with exception set (mellea/core/sampling.py:239), and the metric handler skips it (metrics_plugins.py:443). The streaming handler's own docstring says "a crash is not a sampling verdict" (:490), and abandoning a stream looks like the same case to me.

One way round it would be to report sampling_success=None unless the loop actually reached a verdict. It'd also be worth saying in the StreamAttempt docstring (:232) what attempts looks like after an early exit, since "exactly one attempt in a run is selected" doesn't hold there.

[ajbozarth]

addressed in c9c6806

[planetf1]

Widening past_results here means custom strategies written against the current signature stop type-checking. repair() is the documented extension point, and 0.8.0 ships it as past_results: list[ComputedModelOutputThunk]. Parameter types can't be narrowed in an override, so a subclass with that signature passes mypy on main and fails on this branch:

Argument 4 of "repair" is incompatible with supertype "BaseSamplingStrategy";
supertype defines the argument type as "Sequence[ComputedModelOutputThunk[Any] | None]"  [override]

test/core/test_component_typing.py needed the same change in this PR. Anyone with their own strategy who runs mypy will hit it, whether or not they use stream().

It shows up at runtime too. The completed-fail branch passes every attempt (streaming.py:1000-1019), so for early-broken ones a custom repair() gets None in past_results and an empty list in past_val. The built-in strategies are fine because they only read [-1].

I think this is avoidable. If repair() only gets the completed attempts, the way _select_failure already handles select_from_failure (streaming.py:775-781), then [-1] doesn't change for any of the three built-in strategies (repair() only runs when the latest attempt completed), repair() keeps its 0.8.0 signature, and the None-bearing type stays on the new stream_repair(). The type: ignore[misc] at :779 could go as well. If you'd rather keep the wider type, it's worth a line in the release notes for people with custom strategies.

Separately, the past_results docstrings (base.py:174, :611, :677, :780) still say "List of (unsuccessful) generation results", with no mention of None.

[ajbozarth]

The | None widening is deliberate and I'd like to keep it: past_results, past_actions, and past_val are index-aligned per attempt, and None holds an early-broken attempt's slot so list index stays equal to attempt number. I'd treat it as an intentional breaking change and flag it in the PR title and commit message.

It's a cheap break in practice — no built-in reads past_results (they only use past_actions[-1] / past_val[-1]), so nothing we ship is affected; it only touches custom repair() overrides.

Looping in @jakelorocco — I'm inclined to keep this as an intentional break. Does that sit right with you both, or would you rather I avoid it by filtering repair() to completed attempts only? That keeps the lists mutually aligned but drops the early-broken slots, so list index would no longer equal attempt number.

[planetf1]

stream() accepts any BaseSamplingStrategy, but only uses loop_budget, repair, stream_repair and select_from_failure, so a strategy that does its real work in _sample quietly runs as plain rejection sampling:

• BudgetForcingSamplingStrategy and BaseMBRDSampling (majority voting, MBRD-ROUGE-L) override _sample but inherit RejectionSamplingStrategy.stream_repair, so the thinking budget and the voting are skipped
• ModelFriendlyRepairStrategy overrides repair but inherits RepairTemplateStrategy.stream_repair, so a mid-stream failure gets the generic repair text rather than its ModelFriendlyFeedbackFormatter output

Related: the new docs say "pass a SamplingStrategy to stream()" (use-async-and-streaming.md:354, requirements-system.md:390, and twice in 06-streaming-validation.md), but the parameter is BaseSamplingStrategy. SOFAISamplingStrategy subclasses SamplingStrategy directly, so mypy rejects it here, and at runtime it fails with AttributeError: 'SOFAISamplingStrategy' object has no attribute 'stream_repair' at the first mid-stream failure.

The how-to page lists the supported strategies (use-async-and-streaming.md:395-397), but nothing at the call site tells you. A warning or TypeError for strategies that override _sample would catch it, or at least name the unsupported ones here next to the concurrency_budget note.

[ajbozarth]

I fixed the docs in c9c6806

As for the guard, I investigated adding it and it was non-trivial and hacky. A run-time guard already existing via the NotImplemented error, and the strategies that don't have that are all valid to their particular use case and code path

[planetf1]

One gap for callers: there's no simple way to tell whether sampling actually succeeded. CompletedEvent.success only says the stream finished, so when the budget runs out and every attempt failed validation you still get success=True (that's what the test at test_streaming_sampling.py:257 checks). Right now you'd have to work it out from s.attempts.

_finalize already calculates this as sampling_success for the hook payload, so it'd be nice to expose it on Streamer/EventStreamer too. That would also cover the SamplingResult.success equivalent #403 asks for.

[ajbozarth]

fixed in c9c6806 by adding sampling_success to CompletedEvent and Streamer

[planetf1]

A custom strategy that doesn't implement stream_repair only fails at the first mid-stream failure, after the caller has already received some chunks. It'd be friendlier to check at stream() entry (e.g. type(strategy).stream_repair is BaseSamplingStrategy.stream_repair) and raise straight away.

[ajbozarth]

I actually investigated and tried adding the guard and it was less trivial than expected. The suggested case of checking is not valid as a strategy would not implement stream_repair, but still implement repair, and only trigger of final validation and not mid-stream. That would be a valid use of stream() but this check would block it. The more I investigated the more I found that it was impossible to put a guard in without causing issues with some edge case and the only true clean guard is the one I already have here

[planetf1]

nit: RetryEvent is now emitted but not exported here, while StreamAttempt and the other event types are.

[ajbozarth]

fixed in c9c6806

[planetf1]

nit: "else empty" isn't quite right. A plain stream records one attempt too, so this is only empty before the stream starts.

[ajbozarth]

fixed in c9c6806

[planetf1]

nit: event_queue lost its default and the new parameters are all required. Streamer is exported, so defaults would be the safer choice, even though nobody's likely to construct it directly.

[ajbozarth]

the fact that a user should init a Streamer themselves is clearly documented and it is only exported for typing purposes. This update cleans up a redundant nested default alongside other improvements

[planetf1]

nit: when every attempt breaks early, select_from_failure() isn't called at all. _select_failure only passes it completed attempts and otherwise falls back to the last one (streaming.py:775-785). Worth saying so here.

[jakelorocco]

Should we change this behavior? It seems reasonable that select_from_failure should be able to select from any failed run? Should there be a two-tier system: one tier has early failures and one tier has things that only failed the last requirement?

[ajbozarth]

I fixed the docs in c9c6806

As for jakes comment about changing the behavior, I had Claude write up the following explanation:

---

Why select_from_failure only sees completed attempts today: its signature takes completed outputs (list[ComputedModelOutputThunk]) and final validation results (list[tuple[Requirement, ValidationResult]]). An early-broken attempt has neither — no completed output, only mid-stream PartialValidationResults — so it can't be passed as the contract stands. _select_failure selects among the completed attempts and falls back to the last attempt when none completed.

A two-tier version (letting it also select from early-broken attempts) would take:

• A breaking change to the public select_from_failure signature — widen sampled_results to allow None and carry partial validations — which breaks every custom strategy's override, the same compatibility class as the repair()/past_results widening.
• Little built-in upside: the built-in selectors ignore that data (they return 0/-1), so the plumbing is trivial — but the value is only for a strategy that actually scores failures (e.g. picking the attempt with the most passing requirements), and such a strategy can't rank an early-broken attempt on equal footing, since it has no final verdicts or complete output.
• So the real cost is a design decision, not plumbing: how do you compare a truncated-output-plus-mid-stream-fail attempt against a full-output-plus-final-verdicts one? There's no obvious ranking.

So I'd keep the current completed-only behavior now and treat the two-tier as a deliberate, API-breaking refactor to revisit if there's real demand.

[jakelorocco]

Agree with Nigel's comments as well.

[jakelorocco]

Should we change this behavior? It seems reasonable that select_from_failure should be able to select from any failed run? Should there be a two-tier system: one tier has early failures and one tier has things that only failed the last requirement?

[jakelorocco]

We should keep the existing instruct / sampling strategy behavior.

[jakelorocco]

Should we also override .stream_repair for this one so that it uses the repair strategy for mid-stream repairs?

[ajbozarth]

I dug into this and the short answer is no, primarily because this strategy is designed to work a specific set of requirements that all don't implement stream_validate and thus would never trigger stream_repair. and if a user gave a separate requirement outside those defaults it safely falls back to the default RepairTemplateStrategy.stream_repair

[ajbozarth]

I should have review feedback fixes pushed tomorrow morning along with responses inline, I've just run out of time today to wrap it up

[ajbozarth]

@planetf1 @jakelorocco I believe I have responded to all your feedback, either addressing it or pushing back with a response or question.

If you could re-review please