feat(stdlib): add stream_with_chunking() with per-chunk validation (#901)

[planetf1]

Misc PR

Type of PR

• Bug Fix
• New Feature
• Documentation
• Other

Description

• Link to Issue: Closes feat(stdlib): implement stream_with_chunking() with per-chunk validation #901

Implements stream_with_chunking() — the core streaming orchestration primitive for the streaming validation epic (#891), closing issue #901.

Part of epic #891 · Wave 3 of 4

This PR closes #901. The following are explicitly out of scope and tracked in #902 (the final wave — nothing further is planned after #902):

• Streaming event types (ChunkEvent, QuickCheckEvent, StreamingDoneEvent, FullValidationEvent, CompletedEvent, ErrorEvent)
• Application-level OTEL span and span events for the orchestrator
• record_requirement_check / record_requirement_failure / record_sampling_outcome / record_error calls
• ErrorEvent replacing the MelleaLogger.warning stopgap (see TODO(#902) in streaming.py)
• Narrative documentation: how-to section, requirements-system concept update, end-to-end tutorial

Builds on the now-merged #925 (squash-merged as upstream commit 7912a1df), which added Requirement.stream_validate(). This branch has been rebased directly onto upstream/main; the 13 Wave 3 commits are 8128dfad..3fb501ef.

What changed

• mellea/core/base.py — adds ModelOutputThunk.cancel_generation(error=None): cancels in-progress _generate / _generate_extra tasks, drains the internal async queue (to release any blocked put() calls), closes the open telemetry span (recording the provided error if given, else a generic RuntimeError("Generation cancelled")), and marks the MOT as computed. Uses .pop() on _meta["_telemetry_span"] to prevent KeyError if two coroutines race before _computed is set.

• mellea/stdlib/streaming.py (new) — StreamChunkingResult and stream_with_chunking():

  • Starts a background asyncio.Task that consumes the MOT's async stream, splits accumulated text via ChunkingStrategy, and calls stream_validate() once per complete chunk, passing that single chunk (not the accumulation).
  • When an astream() iteration produces multiple new chunks, they are validated sequentially in order so early exit prevents later chunks in the same batch from being validated or emitted to the consumer.
  • Early exit: on first "fail" result, cancel_generation() is called and StreamChunkingResult.completed is set to False. The failing chunk is not emitted to the consumer; use streaming_failures to inspect what failed.
  • End-of-stream flush: when the stream ends naturally, any trailing fragment withheld by the chunker (e.g. a final sentence with no trailing whitespace) is released via the new ChunkingStrategy.flush() method and run through stream_validate on the same terms as regular chunks. Skipped on early exit (the fragment is mid-token and incomplete).
  • Final validation: after natural completion, validate() is called on all non-failed requirements. Skipped on early exit.
  • Clone-per-call: copy(req) clones each requirement before use; originals are never mutated.
  • String aliases "sentence", "word", "paragraph" resolve to the corresponding ChunkingStrategy subclasses.
  • Exception in stream_validate → orchestrator calls cancel_generation(error=exc) so backend telemetry records the real cause, then surfaces the exception to the consumer via astream() / acomplete().
  • finally block calls cancel_generation() when not already computed, so the backend producer is stopped even on external CancelledError (which bypasses except Exception).

• mellea/stdlib/chunking.py — adds ChunkingStrategy.flush(accumulated_text) -> list[str] (default returns [] — backward-compatible for external chunkers). SentenceChunker, WordChunker, and ParagraphChunker each override to return the withheld trailing fragment.

• mellea/stdlib/__init__.py — re-exports StreamChunkingResult and stream_with_chunking.

• docs/examples/streaming/streaming_chunking.py (new) — end-to-end example with a stateful MaxSentencesReq showing the canonical accumulate-on-self pattern. Marked # pytest: ollama, e2e.

Spec adherence and deliberate variations

For reviewer attention. Items (a)–(d) are spec-compliant (explicit or, in (d), one of the options the spec names); items (e)–(h) are design decisions taken where the spec is silent.

(a) Chunk semantics — spec-compliant. Addresses @jakelorocco's review on #925 (now approved). stream_validate receives a single complete chunk from the chunking strategy, not the accumulated output. Matches the contract in the #891 epic and the #900 spec.

(b) Clone-per-attempt — spec-compliant. copy(req) clones each requirement before use; originals are never mutated (per #891 and #901).

(c) Early-exit cancellation — spec-compliant. cancel_generation() is called on first "fail" to stop the backend producer before it blocks on mot._async_queue (per #891).

(d) End-of-stream flush — spec-compliant design choice. The #899 ABC docstring offers two options for the trailing fragment: "pass to a final validator or discard". This PR takes the first by adding ChunkingStrategy.flush(accumulated_text) and running the returned fragment through stream_validate on the same terms as regular chunks, so the "no unvalidated content reaches the consumer" invariant extends to the trailing fragment.

(e) Exception handling in stream_validate — variation; spec-silent. The spec covers cancel_generation() on explicit "fail" but not on unhandled exceptions. We extend the same resource-leak reasoning: any orchestrator exit must stop the producer. The original exception is passed to cancel_generation(error=...) so the backend telemetry span records the real cause. If cancel_generation() itself raises, we log via MelleaLogger.warning with a TODO(#902) marker and still surface the original exception to the consumer. The finally block additionally covers external CancelledError (which bypasses except Exception) for the same reason.

(f) astream() single-consumer — variation; spec-silent. #901 does not say whether re-iteration is supported. Current implementation is single-consumer (queue drained to the None sentinel). Docstring updated to state the contract explicitly.

(g) Validator latency — variation; spec-silent. Chunks are emitted only after all active validators return for that chunk. A slow stream_validate therefore adds latency; the preserved invariant is that the consumer never sees unvalidated content. A concurrent-emission fast path may be added if a concrete use case drives it.

(h) Orchestrator-level OTEL deferred to #902 — per the epic's explicit instruction: "these event types are the equivalent observability mechanism for the streaming path". Backend-layer instrumentation (mot._meta["_telemetry_span"], token/latency/error metrics via plugins) remains intact. Not in this PR: application-level trace span for stream_with_chunking, span events for chunk lifecycle, record_requirement_check / record_requirement_failure / record_sampling_outcome / record_error calls, and the ErrorEvent that will replace the MelleaLogger.warning stopgap. All enumerated in the #902 acceptance criteria.

Testing

• Tests added to the respective file if code was changed
  • test/stdlib/test_streaming.py (new) — 12 unit tests via StreamingMockBackend (no Ollama needed) covering normal completion, early exit on fail, clone isolation, quick_check_backend routing (asserts clone saw val_backend for every call), deadlock prevention, as_thunk correctness, astream() chunk granularity, no-requirements passthrough, per-chunk contract (exact-match capture of individual chunks), trailing-fragment flush, early exit on trailing fragment, multi-chunk batch with mid-batch fail, cancel-on-fail spy verification, and exception-in-validator cancellation.
  • test/stdlib/test_chunking.py — 13 new tests covering ChunkingStrategy.flush (ABC default + each built-in chunker's fragment logic).
  • test/core/test_stream_validate.py (on stacked feat: add stream_validate() hook to Requirement (#900) #925) — 9 tests.
  • Full related suite: 69 passed locally (uv run pytest test/stdlib/test_streaming.py test/stdlib/test_chunking.py test/core/test_astream_mock.py -q).
• New code coverage: ~92% on streaming.py (16 tests). Remaining uncovered lines are all error-within-error cleanup paths (cancel_generation raising during exception handling, external CancelledError cleanup failing, acomplete re-raising an orchestration exception) and the TOCTOU RuntimeError race — these require fault injection beyond unit scope and are acceptable gaps.
• Ensure existing tests and github automation passes (a maintainer will kick off the github automation when the rest of the PR is populated)
• Integration test against local Ollama (granite4:micro) — example runs cleanly; both sentences of a non-terminated response reach the consumer via the flush path.

Attribution

• AI coding assistants used

[planetf1]

FYI @jakelorocco — next part of the streaming validation epic is ready for review.

[planetf1]

Resolved — all issues addressed across Wave 3 and Wave 4 fix rounds. Ready for review.

[planetf1]

All of @psschwei's review comments are now addressed — details in the individual threads, but the short version:

• Raise-once guard — replaced the stash-as-sentinel pattern with a dedicated _exception_surfaced flag so the two "None" cases (never set vs. already cleared) can't be confused
• _copy_from / __copy__ / __deepcopy__ — _cancelled is now propagated in all three paths
• full_text on early exit — tracks emitted_text separately so it reflects only what the consumer actually received, not the raw accumulated delta
• Indentation question — inline comment added; the outer break is targeting the while loop intentionally

All covered by regression tests. @psschwei @nrfulton — would really appreciate a re-review when you get a chance. There's a stacked PR (#958) waiting behind this one, so getting this merged would unblock that too.

[psschwei]

edit: sigh... comments posted on wrong lines 😢
will repost review

[psschwei]

Thanks, this is close and the focused streaming/chunking suite passes locally.

I found a few lifecycle issues that seem worth fixing before merge: one setup-path leak before the orchestrator exists, one validator-concurrency cleanup issue, and one HF-specific cancellation gap where the asyncio task cancellation does not necessarily stop the underlying generation thread.

[psschwei]

This setup order can leak an already-started backend stream if requirement cloning fails.

generate_from_context() starts the backend _generate task, then copy(req) runs before the orchestration task/finally block exists. If a requirement’s __copy__ raises, nothing cancels or drains the MOT queue, so real backend producers can remain pending against the maxsize=20 queue.

Could we either clone requirements before starting generation, or wrap the post-generate_from_context()setup intry/exceptand callawait mot.cancel_generation(error=exc)` before re-raising?

[planetf1]

Confirmed. The leak window is as described: backend generation is in flight from generate_from_context, copy(req) then runs with no cleanup wrapping it, and a raising __copy__ leaves the MOT feeder stuck against the maxsize=20 queue with no consumer.

Addressed in 7fc40a4 by reordering copy(req) ahead of generate_from_context, so the backend simply never starts if cloning raises. Preferred reordering over wrap-and-cancel because it eliminates the failure case rather than cleaning up after it.

The window after generate_from_context returns and before the orchestration task is armed (StreamChunkingResult(...) plus asyncio.create_task(...)) is covered by try/except BaseException; BaseException rather than Exception so loop-shutdown and cancellation paths are caught. Cleanup calls mot.cancel_generation() inside a nested try/except Exception, so a cleanup failure cannot mask the original exception — same shape as streaming.py:288–294. On a failing create_task, coro.close() suppresses the "coroutine was never awaited" warning.

Scope check: the in-tree Requirement subclasses do not override __copy__, so no current caller trips the bug. However, the base-class docstring at core/requirement.py:302 advertises __copy__ as the recommended isolation mechanism for stateful stream_validate overrides, so user-defined subclasses are the intended fix target.

Docstring changes in the same commit:

• stream_with_chunking: the existing "before use" sentence now says "before backend generation begins" and records the no-leak guarantee; a new Note: block documents that __copy__ exceptions propagate with no backend started.
• Requirement.stream_validate gets a sentence at the __copy__ paragraph explaining the failure contract for override authors.

Test: parametrised over _PlainReq (no override) and _RaisingCopyReq (__copy__ raises ValueError). The failure branch asserts generate_from_context_call_count == 0 — the hard invariant: backend never started.

[psschwei]

This cancels the asyncio task wrapper, but for the HF backend _generate_extra is created from asyncio.to_thread(model.generate, ...). Cancelling the asyncio task does not stop an already-running generation thread, so early-exit validation can return while HF/GPU generation continues in the background.

Can we add a cooperative cancellation hook for backend producers, at least for HF? For example, HF could install a StoppingCriteria backed by a cancellation event, and cancel_generation() could trigger that before awaiting/cancelling the task wrappers.

[planetf1]

Confirmed. cancel_generation today cancels the asyncio task wrapping asyncio.to_thread(model.generate, ...), but the thread itself runs to completion — the GPU stays busy and the MOT reports cancelled=True while work is still ongoing.

Addressed in d8018dd by adding a generic cooperative-cancel hook on ModelOutputThunk and wiring HuggingFace to it via StoppingCriteria:

• ModelOutputThunk gains _cancel_hook: Callable[[], None] | None (default None). cancel_generation calls the hook before cancelling _generate/_generate_extra, so the thread receives its stop signal first and the subsequent task cancellation then unblocks the awaiter. The hook is called inside a try/except Exception so a faulty hook cannot mask the existing cancel path.
• HuggingFaceBackend builds a threading.Event per streaming request, attaches a private _EventStoppingCriteria (merged with any user-supplied stopping_criteria via _install_cancel_stopping_criteria()), and assigns output._cancel_hook = cancel_event.set. Both streaming paths (KV-cache and standard) are updated; the hook is armed before the task is created so a cancel racing with task creation still fires cleanly.
• All three MOT copy sites (_copy_from, __copy__, __deepcopy__) reset _cancel_hook to None on the copy — a copied MOT is a distinct computation and must not share the original's thread signal.

Chose the generic hook rather than an HF-specific cancel event because a future litellm or vLLM streaming path has the same shape (blocking work wrapped in to_thread), and keeping the contract on MOT avoids a second round of changes to core/base.py when that lands.

Tests in test/core/test_base.py:

• test_cancel_generation_invokes_cancel_hook_before_task_cancel: hook fires and cancel_generation() returns within 1 s when the thread only exits on event.set().
• test_cancel_hook_not_forwarded_by_copy_methods: all three copy sites verified to reset the hook.
• test_cancel_generation_hook_exception_is_suppressed: faulty hook does not mask cancellation.

[psschwei]

asyncio.gather() propagates the first validator exception, but sibling validators are not cancelled. That means if one stream_validate() raises while another validator is still doing work, the orchestrator starts cleanup and surfaces the error while the sibling validator can continue running in the background.

That is especially risky if validators call an LLM or mutate external state. Could we switch this to TaskGroup, or explicitly create tasks and cancel/await the remaining tasks on the first exception? The same concern applies to the final validate() gather below.

[planetf1]

Confirmed. Both asyncio.gather call sites (inside _validate_and_emit around line 200, and the final validate() gather around line 270) use the default return_exceptions=False behaviour: the first failing task raises and peer tasks are left detached. If a slow validator is still awaiting a backend response when a sibling fails, it runs on unobserved until it completes or the event loop shuts down.

Addressed in 7fc40a4 by replacing both gather sites with asyncio.TaskGroup. On first failure, the TaskGroup cancels every peer task, awaits their cancellations, and re-raises. Python floor is already 3.11 per pyproject.toml, so this is a direct substitution with no conditional imports.

One adjustment to the outer handler: TaskGroup wraps failures in an ExceptionGroup. ExceptionGroup inherits from Exception, so the existing except Exception as exc: still catches it, but cancel_generation(error=exc) and the downstream _chunk_queue.put(exc) are more readable with the original exception. The handler now unwraps a single-element group before forwarding. If multiple validators fail simultaneously, the suppressed siblings are logged before taking exceptions[0] so they're not invisible in production.

Observable behaviour on success is unchanged (same pvr list, same final_validations). On failure the caller still sees the first exception with the same type — the only change is that peer tasks are now awaited to completion before the handler runs, rather than leaking into the background.

Test: test_stream_with_chunking_cancels_peer_validators — one requirement raises immediately in stream_validate, the second sleeps for 5 s and sets a flag on completion. await result.acomplete() raises the first exception, and the slow sibling's completion flag is asserted to be unset afterwards, proving TaskGroup cancelled it rather than left it detached.

[planetf1]

Two fix commits landed since the last review round — summary for @psschwei and @nrfulton:

bf9a62bc — three pre-merge correctness fixes

1. HF backend scope leak (huggingface.py): _install_cancel_stopping_criteria was unconditionally called on both streaming and non-streaming paths. On the non-streaming path there is no orchestrator calling cancel_generation(), so the StoppingCriteria was dead code and silently wrapped any user-supplied stopping_criteria on every decode step. Now guarded with if stream:.

2. Outer CancelledError absorption (core/base.py): cancel_generation() was catching all CancelledErrors with a bare pass, which silently absorbed external cancellation (e.g. asyncio.wait_for timeout). Fixed to re-raise when asyncio.current_task().cancelling() > 0, i.e. when the outer task itself is being cancelled rather than the inner _generate task we just called .cancel() on.

3. Non-streaming MOT guard (stdlib/streaming.py): stream_with_chunking() now raises RuntimeError immediately if the backend returns an already-computed ModelOutputThunk. Previously this would cause the orchestrator loop to skip entirely, producing empty output and silently passing all final validators against an empty string.

9a715d62 — follow-up polish

• huggingface.py: replaced _cancel_event.set if stream else None ternary with an explicit if stream: block to silence the Pyright possibly-unbound warning
• base.py: corrected misleading comment "Python 3.12+" → "Python 3.11+" (Task.cancelling() was added in 3.11)
• streaming.py: added RuntimeError to the Raises: section of stream_with_chunking()'s docstring
• test/core/test_base.py: added test_cancel_generation_propagates_outer_cancellation — a direct regression test for the re-raise path in fix 2, which had no test coverage

Full fast test suite (1 462 tests) passes.

[psschwei]

one minor question, but otherwise LGTM

[psschwei]

do we also need to flip _exception_surfaced here?

[planetf1]

Good catch — yes, and digging in turned up a related latent issue.

Two problems on this branch:

1. The flag was never set. As you flagged, self._exception_surfaced is not flipped here, so a second acomplete() would raise again.

2. The branch was unreachable for CancelledError anyway. asyncio.Task.exception() raises CancelledError on a cancelled task rather than returning it. So the existing test at test_external_task_cancellation_releases_consumers was passing only by side effect — every call to .exception() re-raised the cancel, and raise task_exc was unreachable on that path.

Fix in 66260fe restructures the branch to short-circuit on _exception_surfaced, check task.cancelled() before calling .exception(), and flip the flag before raising in both sub-branches. Added test_external_cancellation_acomplete_raise_once (calls acomplete() twice on a cancelled task and asserts the second call returns cleanly) to lock the raise-once contract for this path.

[jakelorocco]

looks good; a few minor nits

[jakelorocco]

I think validate and stream_validate should return equivalent results for most requirements.

[planetf1]

Good call — the always-True validate() is technically unreachable when the limit is exceeded (stream_validate cancels first), but it teaches the wrong pattern. Fixed to return ValidationResult(result=self._count <= self._limit) so the two methods agree, which is the pattern users should follow when writing their own requirements.

[jakelorocco]

This doesn't hold if you have it fail on the 3rd chunk. Individual chunks strip whitespace, but the full text has whitespace between sentences.

[planetf1]

Good catch — the test only covers the single-chunk case, which avoids the whitespace discrepancy. When multiple chunks are emitted before early exit, emitted_text (chunk concatenation) loses the inter-sentence spacing from the original token stream. Fixed: replaced emitted_text with an emitted_end cursor into accumulated, so full_text on early exit is accumulated[:emitted_end] — original spacing preserved. Added a 3-chunk test case to cover this.

[jakelorocco]

The fact that post_processing doesn't run makes me a bit nervous. It means that several key chunk concatenation steps won't be done for the mot's values either.

Maybe this just requires adding some documentation that cancelled mots won't have a parsed_repr value, etc...

[planetf1]

This can be confusing — added a doc note to clarify what's reliable after early exit. Also set parsed_repr = value on cancel for plain-text outputs so it's consistent with normal completion.

[jakelorocco]

mots can be typed. We lose that typing information here. I wonder if we should propagate it or just force the caller to investigate the parsed type?

[planetf1]

The streaming path only ever produces text, so the right type here is ModelOutputThunk[str] — updated the return annotation. If the caller wants a typed result they'd need to re-parse full_text themselves, which feels right given the truncated output.

[jakelorocco]

@planetf1, I reviewed with Nathan. I have added a few more comments and I believe some of my initial comments are still unaddressed.

This work is really good and will be nice to integrate into Mellea. As a heads up, we do have future work to revisit sampling strategies and to allow for optimizing "mellea programs" in general. As a result, we may actually combine some of this streaming stuff into some new sort of "sampling result" type in the future. This wasn't part of the initial proposal and doesn't require addressing now, but wanted to keep you in the loop / let you know in case someone mentions something about it.

[jakelorocco]

Since we are no longer allowing requirements to succeed early (and not be validated against again), lets change these parameters to requirements and validation_backend (or something like that).

[planetf1]

Good point — quick_check_* was meaningful when there were two tiers of requirement; that distinction is gone. Renamed quick_check_requirements → requirements and quick_check_backend → validation_backend throughout (signature, docstring, internal uses, tests, example).

[jakelorocco]

Is it also possible to move streaming to core? Or will that introduce too many imports / dependencies there?

[planetf1]

The blocker is ChunkingStrategy and its concrete implementations — streaming.py imports them directly, so moving it to core/ without also moving chunking.py would create a core → stdlib import and a circular dependency. Moving both feels like over-promoting what is essentially text-splitting utility code into the core abstraction layer.

I'd keep streaming.py in stdlib/ for now. Core already provides the primitives (stream_validate, cancel_generation, ModelOutputThunk); the orchestrator composes them — that composition feels like a stdlib concern. The longer-term direction from #1013 (MOT-owned chunking via stream_parsed_repr) would dissolve the orchestrator anyway, so there's limited value in an interim move.

[planetf1]

@jakelorocco ready for review.

fixes / changes have been made