diff --git a/docs/dev/streaming-phase2-plan.md b/docs/dev/streaming-phase2-plan.md new file mode 100644 index 000000000..4be3a1b9d --- /dev/null +++ b/docs/dev/streaming-phase2-plan.md @@ -0,0 +1,314 @@ +# Phase 2 Validated Streaming — Epic Breakdown + +> **Purpose:** decompose what remains of the validated-streaming epic (after +> Phase 1 closed) into distinct, separately-shippable pieces, with dependencies +> and a +> recommended ordering. The Phase 1 *refactor* is the first item and it blocks +> the rest. Each item section is written to be liftable into a GitHub issue once +> the shape is agreed. +> +> **Status:** discussion draft to frame the Phase 2 planning conversation. +> Shared alongside a working POC on this branch; nothing is filed yet. +> +> **POC:** `mellea/stdlib/streaming_poc.py` (plus `ModelOutputThunk.__aiter__` in +> `mellea/core/base.py`) is a working prototype of item 1, including hook-based +> event emission. Sections below note where the POC already validates a decision. +> +> **Origin:** the companion proposal [streaming-simplification.md](streaming-simplification.md) +> came out of an offline discussion between @psschwei and @ajbozarth; it is +> included verbatim as the seed for the item-1 direction. + +--- + +## Where things stand + +- **Phase 1 (epic #891) is fully closed.** All children merged: `PartialValidationResult` + (#898), `ChunkingStrategy` (#899), `stream_validate()` (#900), `stream_with_chunking()` + (#901), and streaming event types (#902). Phase 1 shipped the working + `stream_with_chunking` + `StreamChunkingResult` primitive. +- **#1013 is the open Phase 2 epic**, but today it is a *placeholder* scoped + narrowly to one component (`stream_parsed_repr`). In practice "Phase 2" is + larger than that one method — this doc argues for treating the epic as an + umbrella over the items below. + +### Related issues to coordinate with + +- **#403** — streaming for sampling results. May relate to item 4; see that item. +- **#1051** — streaming span events (the OTel-side bridge). It has been waiting on + a streaming-events design; once item 1's event hook is settled, #1051 becomes a + possible *consumer* of that hook rather than separate work. +- **#444** — enhanced tracing epic. The "one span per stream" target (below) lives + here; the item-1 telemetry cleanup should be coordinated with it. + +--- + +## The core reframing + +Phase 1 delivered streaming validation as a **standalone parallel interface** +(`stream_with_chunking` + `StreamChunkingResult`), which epic #891 openly called +a scoped pragmatic choice. Everything left to do divides along two independent +axes that Phase 1 collapsed onto the call site: + +- **How output is consumed / orchestrated** — the *interface* (a separate result + type + background task vs. folding into the normal generation flow). +- **What a "chunk" is and who owns its boundary** — the *type semantics* (an + external string-splitter vs. the MOT knowing its own units). + +The refactor addresses the first axis; `stream_parsed_repr` addresses the +second; the remaining pieces (sampling integration, multi-modal, agent-authoring) +build on one or both. Keeping them as separate items lets each ship and be +reviewed independently. + +--- + +## Item map + +Ordered by dependency. **Item 1 blocks all others.** + +### 1. Refactor: single-task `stream()` primitive ⟶ blocks everything + +Replace `stream_with_chunking()` + `StreamChunkingResult` with a single +`stream()` consumed by a plain `async for`, driven on the caller's task (no +background orchestration task). The only durable state is a thin `Streamer` +handle (`failed_early`, `failure_reason`, `streaming_failures`, `full_text`, +`final_validations`, terminal `mot`). + +```python +streamer = stream(action, backend, ctx, chunking="sentence", requirements=[req]) +async for chunk in streamer: + display(chunk) +if streamer.failed_early: + handle_failure(streamer.failure_reason) +``` + +**Why this matters — streaming shouldn't be a separate interface.** Phase 1 +delivered streaming validation as a standalone parallel path: its own entry +point (`stream_with_chunking`), its own result type (`StreamChunkingResult`, +distinct from `ModelOutputThunk`), and its own consumption protocol +(`astream()`/`events()`/`acomplete()`). Streaming is a *property of a generation +call*, not a different kind of call — the MOT already streams (`astream()`, +`is_computed()`, `cancel_generation()`). The parallel interface wraps a second +orchestration layer around a stream the framework already produces, and the seam +between the two leaks. + +**The leak, concretely — the background-task split.** The parallel path runs on a +background `asyncio` task (`_orchestrate_streaming`); the caller pulls chunks and +events off two queues that task feeds. That *caller task* vs *orchestration task* +split is the root of a recurring class of problems: + +| Symptom | Root | +|---|---| +| Cross-task OTel (tracing) span detach failure, worked around in PR #1361 | the tracing span opens on the caller task but the stream drains on the orchestration task, so the OTel context detach crosses tasks and fails | +| `acomplete()` "works by accident" if you fully drain `astream()` | `_done` is set by the orchestrator task as a drain side effect, not by the documented terminal call | +| `as_thunk` docstring says it raises "before `acomplete()`" but checks `_done.is_set()` | same — `_done` is orchestrator-set, not `acomplete`-set | +| Two text views (`accumulated` vs `emitted_end`) | chunker needs full accumulation; `full_text` should reflect only validated output; they diverge on early exit | +| Raise-once exception plumbing (`_orchestration_exception`, `_exception_surfaced`) | an exception on the orchestrator task must be surfaced to whichever of `astream()`/`acomplete()` the caller drains first | +| Single-consumer guards, `_orchestration_started` event | queue hand-off between tasks needs liveness/ordering coordination | + +None are bugs in isolation — Phase 1 handled each carefully. They are all +consequences of the one architectural choice the refactor removes. + +**Telemetry surfaced this, and single-task fixes it cleanly.** PR #1361 moved +streaming telemetry onto plugin hooks (the #444 direction); this refactor keeps +it there. The problem was never hooks-vs-inline — it is that the two-task split +forces the tracing-span lifecycle to be reconstructed across a task boundary. Per +@psschwei on #1361: *"we should strongly consider 'one span per stream' as the +target design."* A single-task design makes that fall out for free: +`STREAMING_START`/`STREAMING_END` fire at loop entry/exit on one task, so a +subscriber opens the OTel span on START and closes it on END with no cross-task +detach and no re-attach hook. The `STREAMING_ORCHESTRATION_START`/`END` pair — +which existed *only* to bridge the two tasks — can be dropped. + +**Scope of the item:** + +- **Core:** `ModelOutputThunk.__aiter__`/`__anext__` wrapping `astream()` in the + async-iterator protocol. Generic and reusable by plain streaming — could land + independently. +- **Stdlib:** `stream()` + `Streamer`; chunking and per-chunk validation layered + over the shared iteration protocol; full-output validation on natural + completion (this is what runs judge/aLoRA requirements that stream `"unknown"`). +- **Telemetry:** re-point the `STREAMING_*` hooks to fire on the single task; drop + the now-unneeded orchestration-bridge hooks; and move `CompletedEvent` handling + into the `STREAMING_EVENT` (event) plugin rather than finalizing it in the + `STREAMING_END` (OTel span) hook. Coordinate with #444's one-span-per-stream + target. +- **Events (emission):** part of this item — the loop fires `STREAMING_EVENT` for + the full Phase 1 vocabulary (`QuickCheckEvent`, `ChunkEvent`, + `StreamingDoneEvent`, `FullValidationEvent`, `ErrorEvent`, `CompletedEvent`) + uniformly, no event queue. *Event consumption* (an in-band `events()` iterator) + is intentionally **not** carried forward: no integrated consumer exists — + `m serve` streams via raw `mot.astream()`, the telemetry plugins consume the + hook, and the old `events()` iterator appears only in examples. Defer any + convenience consumption API until a real consumer needs it. +- **Migration:** replace `stream_with_chunking()` with `stream()`, via a + deprecation shim or a hard break (see agenda). + +**POC status:** implemented in `mellea/stdlib/streaming_poc.py`. Validates the +single-task shape, per-chunk validation + early exit, full-output validation, +exception propagation, and hook-based event emission — all without the queues, +raise-once plumbing, or `acomplete()` ambiguity above. + +**Deliberate interim behavior:** early-exit `full_text` is delta-granular, not +chunk-exact, until item 2 lands (documented in the POC; see that item). + +### 2. `stream_parsed_repr`: MOT-owned chunking (the epic's literal scope) + +Move chunk-boundary knowledge onto the MOT: each MOT subclass yields complete, +typed units of its own `parsed_repr` type instead of an external +`ChunkingStrategy` splitting a raw string. This is the type-semantics axis from +the reframing — "what is a complete chunk of this output" is a property of the +output type, so it belongs on the MOT, which already owns `parsed_repr`. + +- **Depends on:** item 1 (slots into the loop by changing the chunk *source* — + `chunking.split(accumulated)` → `async for unit in mot.stream_parsed_repr()`). +- **Unblocks / fixes:** the early-exit `full_text` interim behavior from item 1 + (one unit per iteration makes partial output exact — no cursor, no whitespace + edge); external `ChunkingStrategy` can then be deprecated. + +**Example approaches** — these are starting points to react to in the call; the +chosen shape gets locked when the item's PR is opened, not here. + +- **A — async generator yielding the parsed type.** + ```python + async def stream_parsed_repr(self) -> AsyncIterator[S]: ... + ``` + Mirrors `astream()`; typed via the MOT's existing `S`. The `stream()` loop just + iterates it. *For:* smallest surface, reuses the type parameter already on the + MOT. *Against:* `S` is the *final* parsed type — a mid-stream partial (half a + JSON object, a truncated sentence) may not be a valid `S`, so this forces either + "only yield settled units" or a looser `S`. +- **B — pluggable boundary predicate, MOT supplies a default.** + ```python + mot.stream_parsed_repr(boundary=my_predicate) # MOT has a type-appropriate default + ``` + The MOT owns a sensible default boundary per type; callers can override. + *For:* keeps an author extension point, and text vs. multi-modal can ship + different defaults. *Against:* re-introduces a call-site knob the refactor was + trying to remove, and blurs who owns the boundary (MOT vs. requirement). +- **C — typed chunk envelope instead of bare `S`.** + ```python + @dataclass + class ParsedChunk(Generic[S]): + partial: S | None # best-effort parse so far (None if not yet parseable) + raw: bytes | str # the underlying bytes/text for this unit + complete: bool # settled unit vs. a running partial + ``` + *For:* one shape handles multi-modal (bytes/frames), partial-parse state, and + "not enough data yet" (`complete=False`) — it answers the multi-modal and + error-handling questions below in one move. *Against:* heavier; text-only + consumers pay for machinery they don't need, and every consumer unwraps the + envelope. + +The real fork is A vs. C: whether a "chunk" is always a complete `S`, or a +richer object that can carry partial/typed/binary state. Multi-modal (item 3) +pushes toward C. + +**Open questions (from #1013):** + +1. Generator shape & multi-modal — does the signature admit bytes/frames, not just + `str`? (The A-vs-C fork.) +2. Boundary authority — MOT parser, pluggable predicate, or both? +3. Backpressure — if parsed chunks lag raw tokens, where does buffering live? + (`_GenerationState`, added in #909's structural cleanup, is the natural home.) +4. Error handling — partial-parse failure: surface now, wait, or fall back to raw? +5. Backwards-compat — both external `ChunkingStrategy` and MOT-native during + transition? +6. Testability — each MOT type's `stream_parsed_repr` verified against its + non-streaming `parsed_repr`; shared harness shape. + +### 3. Multi-modal streaming chunks + +Extend `stream_parsed_repr` to non-string units — audio segments, image regions, +video. Epic #891 and #1013 name this the *first-class motivation* for MOT-owned +chunking (the `split(accumulated_text: str)` signature forecloses it by design). + +- **Depends on:** item 2 (it is the generalization of the same mechanism to + non-`str` units, and the case that most forces the typed-envelope decision). +- **Note:** likely its own item — the fixture/test story and the per-modality + boundary logic are substantial and separable from the text case. + +### 4. Sampling integration / convenience wrapper + +A session-level entry point (epic #891 sketched `stream_instruct()`) giving the +familiar `instruct()`-style interface with streaming validation + retry handled +internally, so common-case callers never touch the low-level `async for` + +read-fields idiom. + +- **Depends on:** item 1 only (wraps the new primitive, not `stream_with_chunking`). + **Parallelizable with item 2** — a good candidate for a second developer to pick + up alongside `stream_parsed_repr`, since both proceed directly off item 1. +- **Possible overlap with #403 (investigate at pickup):** #403 is an old, + underspecified p2 asking to stream `SamplingResult`s from `BaseSamplingStrategy`. Item 4 also + combines streaming with retry, but via a different path — Phase 1's primitive + bypasses `BaseSamplingStrategy` (callers own the retry loop), so item 4's retry + is not #403's sampling-loop retry. Whether they converge or stay separate is + undecided; #403 has no concrete design to reconcile against yet. Resolve when + item 4 is picked up, not before. +- **Existing unmigrated consumer:** the prototype + [`Mellea-partials`](https://github.com/HendrikStrobelt/Mellea-partials) already + ships a working `stream_instruct` (as a `MelleaSession` powerup, see its + `doc/stream_instruct.md`) plus its own `StreamEvent` types — it predates the + Phase 1 upstream and never migrated onto it. It is the natural reference + implementation and forcing-function consumer: converging it onto the refactored + primitive validates this API against real usage rather than a guess. +- **Open:** exact shape, name, and whether it is the home for a non-iterating + "just give me the validated result" path. Genuinely undesigned — flag as such. + +### 5. Agent-authoring patterns (park / separate track) + +Named directly in the #1013 body as in-scope Phase 2 work (one of the threads the +placeholder epic was raised to track). Both `stream_validate` and +`stream_parsed_repr` have a deterministic check against a non-streaming +counterpart (`validate()` / `parsed_repr`), which makes them candidates for +agent-friendly authoring (potentially skills). Also raised in the PR #942 thread. + +- **Depends on:** item 2 (needs the `stream_parsed_repr` contract to exist first). +- **Recommendation:** its own issue, not folded into the type-semantics work — it + will balloon scope otherwise. + +--- + +## Dependency picture + +```mermaid +flowchart TD + I1["1. Refactor: stream()
(event emission via hook,
__aiter__, telemetry, shim)
— blocks all"] + I2["2. stream_parsed_repr
(the epic's literal scope)"] + I4["4. Sampling / stream_instruct
(rel #403)"] + I3["3. Multi-modal chunks"] + I5["5. Agent-authoring
(own track)"] + + I1 --> I2 + I1 --> I4 + I2 --> I3 + I2 --> I5 +``` + +Items 2 and 4 depend only on item 1 and can run in parallel (separate owners); +items 3 and 5 depend on item 2. + +--- + +## Recommended sequencing + +1. **Land the refactor (item 1) first** — the blocking foundation. Event emission + via hook is part of it, so Phase 1's event surface is not regressed. +2. **Then parallelize across developers:** `stream_parsed_repr` (item 2) and + sampling integration (item 4) both proceed directly off item 1 and are + independent enough for separate owners. +3. **Multi-modal (item 3) and agent-authoring (item 5)** follow item 2. + +Whether the existing epic becomes the umbrella, or a fresh epic is opened and +issue #1013 is demoted to the `stream_parsed_repr` item-2 issue, is an +implementation detail of issue bookkeeping, not a design decision. + +--- + +## Open decisions to settle at the call + +- Does the refactor (item 1) ship standalone, or is it gated on `stream_parsed_repr` + (item 2)? Shipping item 1 alone means an interim mid-epic behavior break: early-exit + `full_text` is delta-granular until item 2 lands. +- Migration: deprecation shim for `stream_with_chunking()`, or a hard break? +- Which `stream_parsed_repr` shape (A / B / C) is the direction — specifically the + A-vs-C "bare `S` vs. typed envelope" fork, which multi-modal leans on. diff --git a/docs/dev/streaming-simplification.md b/docs/dev/streaming-simplification.md new file mode 100644 index 000000000..6f391f3b0 --- /dev/null +++ b/docs/dev/streaming-simplification.md @@ -0,0 +1,87 @@ +# Streaming API Simplification + +--- + +## Proposal + +Replace `stream_with_chunking()` + `StreamChunkingResult` (two async queues, a +background task, `_done`/`_orchestration_started` events, raise-once exception +plumbing, single-consumer guards) with a single `stream()` that you consume with +a plain `async for`. Chunking becomes a parameter; incremental validation +becomes a parameter; full-output validation runs after the loop exactly like the +non-streaming path. The only state that outlives the loop is a two-field +`Streamer` (`failed_early` / `failure_reason`). Migrate via a deprecation shim. + +```python +streamer = stream(action, backend, ctx, chunking="sentence", requirements=[req]) +async for chunk in streamer: + display(chunk) + +if streamer.failed_early: + handle_failure(streamer.failure_reason) +``` + +--- + +## Flow + +The natural-completion path and the early-exit path diverge on a single +`"fail"` result; both converge on the same `try/finally` cleanup that stops the +backend. Full-output validation runs *after* the loop, and only on natural +completion. + +```mermaid +flowchart TD + Start(["stream(action, backend, ctx,
chunking, requirements)"]) --> Gen["Streamer wraps _drive() generator"] + Gen --> Loop{"mot.is_computed()?"} + + Loop -- no --> Delta["await mot.astream() → delta"] + Delta --> Chunk["chunking.split(): new chunks
(chunking=None → raw delta)"] + Chunk --> HasReq{"requirements
set?"} + + HasReq -- no --> Yield["yield chunk to caller"] + HasReq -- yes --> Validate["stream_validate() per requirement
(asyncio.gather)"] + Validate --> Fail{"any result
== 'fail'?"} + + Fail -- no --> Yield + Fail -- yes --> SetFlag["Streamer.failed_early = True
failure_reason = pvr.reason"] + SetFlag --> Cleanup + + Yield --> Loop + + Loop -- yes --> Flush["chunking.flush(): trailing fragment
(skipped in raw-delta mode)"] + Flush --> Cleanup["finally: if not computed →
mot.cancel_generation()"] + + Cleanup --> Done{"failed_early?"} + Done -- "yes (early exit)" --> HandleFail["caller: handle failure_reason
⚠️ context degraded — no avalidate"] + Done -- "no (natural)" --> AValidate["caller: mfuncs.avalidate(
full_output_reqs, ctx)"] + + HandleFail --> End([done]) + AValidate --> End + + Break["caller break / aclose()"] -. "GeneratorExit at yield" .-> Cleanup + + classDef decision fill:#fff3cd,stroke:#856404,color:#000; + classDef cleanup fill:#d1ecf1,stroke:#0c5460,color:#000; + classDef caller fill:#e2e3e5,stroke:#383d41,color:#000; + class Loop,HasReq,Fail,Done decision; + class Cleanup cleanup; + class HandleFail,AValidate,Break caller; +``` + +The dotted edge is the `break`-safety path: abandoning the `async for` delivers +`GeneratorExit` to the suspended `yield`, so the same `finally` cancels the +backend even when the caller never iterates to completion. + +--- + +## Core insight + +Streaming is a sequence of values over time. Chunking is a transformation on +that sequence. Validation is a filter on that sequence. These are all +composable operations on an `AsyncIterator[str]` — they do not require a +dedicated result object or a coupled orchestration layer. + +The data emitted by a raw token stream and a chunked stream is the same type +(`str`). The only difference is granularity. There is no inherent reason they +need different return types or different call sites. diff --git a/docs/examples/streaming/validated_streaming_poc.py b/docs/examples/streaming/validated_streaming_poc.py new file mode 100644 index 000000000..d7ed1978b --- /dev/null +++ b/docs/examples/streaming/validated_streaming_poc.py @@ -0,0 +1,240 @@ +# pytest: skip + +"""Runnable verification harness for the validated-streaming refactor POC. + + uv run python docs/examples/streaming/validated_streaming_poc.py + +Transient scaffolding, not a usage example: it drives the prototype +`mellea.stdlib.streaming_poc` module (which is not yet wired into the framework) +and uses a mock backend that pokes `ModelOutputThunk` internals, so it is marked +`skip` and excluded from CI. It exists only to exercise the POC's interface end +to end while the refactor is unplugged; delete it once the refactor lands and a +real test suite replaces it. + +Runs without Ollama (deterministic mock stream, from test/core/test_astream_mock.py) +and prints one block per case: + + 1. Plain streaming — `async for delta in mot` straight off the MOT (core only). + 2. Per-chunk requirement passes — natural completion + full-output validate. + 3. Per-chunk requirement fails a chunk — early exit, no final validate. + 4. Judge-style requirement (streams "unknown") — passes the stream-end validate. + 5. Judge-style requirement — fails the stream-end validate. + 6. Requirement raises mid-stream — exception propagates to the caller. + +Each case also emits the real `STREAMING_START`/`STREAMING_EVENT`/`STREAMING_END` +hooks; the recorders installed at startup print them to show the single-task +event lifecycle. +""" + +import asyncio +from typing import Any + +from mellea.core.base import CBlock, GenerateType, ModelOutputThunk +from mellea.core.requirement import ( + PartialValidationResult, + Requirement, + ValidationResult, +) +from mellea.plugins import hook, register +from mellea.stdlib.streaming_poc import stream + + +# --- register REAL telemetry hooks (the same ones the prod plugin uses) ------ +def _install_hook_recorders() -> None: + """Subscribe to the actual STREAMING_START/END hooks and print each firing. + + This is what proves the pitch's telemetry claim: the real hooks fire at two + clean points on ONE task. A production span subscriber would open on START, + close on END — no cross-task detach. + """ + + @hook("streaming_start") + async def _on_start(payload: Any, ctx: Any) -> Any: + print( + f" [hook streaming_start] id={payload.streaming_id[:8]} " + f"reqs={payload.requirement_count} chunker={payload.chunking_strategy}" + ) + return None + + @hook("streaming_end") + async def _on_end(payload: Any, ctx: Any) -> Any: + print( + f" [hook streaming_end ] id={payload.streaming_id[:8]} " + f"success={payload.success} reason={payload.failure_reason!r} " + f"exception={payload.exception!r}" + ) + return None + + @hook("streaming_event") + async def _on_event(payload: Any, ctx: Any) -> Any: + print(f" [hook streaming_event] {type(payload.event).__name__}") + return None + + register([_on_start, _on_end, _on_event]) + + +# --- fake MOT plumbing (from test_astream_mock.py) -------------------------- +async def _mock_process(mot: ModelOutputThunk, chunk: Any) -> None: + if mot._underlying_value is None: + mot._underlying_value = "" + if chunk is not None: + mot._underlying_value += chunk + + +async def _mock_post_process(mot: ModelOutputThunk) -> None: + mot.parsed_repr = mot.value # type: ignore[assignment] + + +def _streaming_mot(text: str, *, tokens: int = 6) -> ModelOutputThunk: + """A MOT preloaded to stream `text` in `tokens` slices, then a sentinel.""" + mot: ModelOutputThunk = ModelOutputThunk(value=None) + mot._call.action = CBlock("demo") + mot._gen.generate_type = GenerateType.ASYNC + mot._gen.process = _mock_process + mot._gen.post_process = _mock_post_process + mot._gen.chunk_size = 0 + step = max(1, len(text) // tokens) + for i in range(0, len(text), step): + mot._gen.queue.put_nowait(text[i : i + step]) + mot._gen.queue.put_nowait(None) # completion sentinel + return mot + + +# --- 1. PLAIN streaming: async for straight off the MOT --------------------- +async def demo_plain() -> None: + print("\n=== 1. Plain streaming — `async for delta in mot` (core, no stdlib) ===") + mot = _streaming_mot("The robot learned to cook. It burned the toast.") + got = "" + async for delta in mot: # <-- the NEW protocol; no while/is_computed loop + got += delta + print(f" delta: {delta!r}") + print(f" computed={mot.is_computed()} value={mot.value!r}") + assert got == mot.value + + +# --- 2-6. VALIDATED streaming through stream() ------------------------------ +class _NoBurntToast(Requirement): + """Per-chunk failure: fails the instant a chunk mentions burning (no LLM). + + Its final `validate()` (run on natural completion) re-checks the full text + deterministically via `validation_fn`, so even a per-chunk requirement gets + the stream-end pass — matching Phase 1. + """ + + def __init__(self, desc: str) -> None: + super().__init__( + desc, + validation_fn=lambda ctx: ValidationResult( + "burn" not in str(ctx).lower(), reason="mentions burning" + ), + ) + + async def stream_validate( + self, chunk: str, *, backend: Any, ctx: Any + ) -> PartialValidationResult: + if "burn" in chunk.lower(): + return PartialValidationResult(success="fail", reason="mentions burning") + return PartialValidationResult(success="unknown") + + +class _MaxLength(Requirement): + """Judge-style requirement: can't judge per-chunk, only the full output. + + Streams `"unknown"` throughout, then checks the complete text at stream end + via `validate()` — the exact pattern that breaks if final validation is + dropped. Deterministic (`validation_fn`), no LLM. + """ + + def __init__(self, limit: int) -> None: + super().__init__( + f"at most {limit} chars", + validation_fn=lambda ctx: ValidationResult( + len(str(ctx)) <= limit, reason=f">{limit} chars" + ), + ) + + +class _Explodes(Requirement): + """Raises mid-stream to exercise the exception path (span closes in error).""" + + async def stream_validate( + self, chunk: str, *, backend: Any, ctx: Any + ) -> PartialValidationResult: + raise RuntimeError("validator blew up") + + +class _FakeBackend: + """Minimal backend: generate_from_context returns a preloaded streaming MOT.""" + + def __init__(self, text: str) -> None: + self._text = text + + async def generate_from_context(self, action, ctx, *, model_options=None): + # The "context" passed to validate() is just the text here, so the + # judge-style validation_fn can inspect the full output deterministically. + return _streaming_mot(self._text), self._text + + +async def demo_validated(text: str, label: str, req: Requirement) -> None: + print(f"\n=== {label} ===") + backend = _FakeBackend(text) + streamer = await stream( + CBlock("demo"), + backend, # type: ignore[arg-type] + ctx=None, # type: ignore[arg-type] + chunking="sentence", + requirements=[req], + ) + async for chunk in streamer: + print(f" chunk: {chunk!r}") + print(f" failed_early={streamer.failed_early} reason={streamer.failure_reason!r}") + print( + f" streaming_failures={[(type(r).__name__, p.success) for r, p in streamer.streaming_failures]}" + ) + print(f" full_text={streamer.full_text!r}") + finals = [v.as_bool() for v in streamer.final_validations] + print(f" final_validations (stream-end validate): {finals}") + + +async def main() -> None: + _install_hook_recorders() + await demo_plain() + await demo_validated( + "The robot sliced a tomato. It plated the dish beautifully.", + "2. Per-chunk req passes -> natural completion + final validate", + _NoBurntToast("no burning"), + ) + await demo_validated( + "The robot sliced a tomato. Then it burned the whole kitchen down.", + "3. Per-chunk req fails chunk 2 -> early exit, NO final validate", + _NoBurntToast("no burning"), + ) + await demo_validated( + "Short and sweet.", + "4. Judge-style req: streams 'unknown', PASSES final validate", + _MaxLength(100), + ) + await demo_validated( + "This one is deliberately written to be quite a bit longer than the limit " + "so that the stream-end validate() is the thing that catches it.", + "5. Judge-style req: streams 'unknown', FAILS final validate", + _MaxLength(50), + ) + + print("\n=== 6. Requirement raises mid-stream -> exception propagates ===") + streamer = await stream( + CBlock("demo"), + _FakeBackend("The robot sliced a tomato. It plated the dish."), # type: ignore[arg-type] + ctx=None, # type: ignore[arg-type] + chunking="sentence", + requirements=[_Explodes("boom")], + ) + try: + async for chunk in streamer: + print(f" chunk: {chunk!r}") + except RuntimeError as exc: + print(f" caller caught: {exc!r}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/mellea/core/base.py b/mellea/core/base.py index 89cfc378f..d3a41146e 100644 --- a/mellea/core/base.py +++ b/mellea/core/base.py @@ -1214,6 +1214,34 @@ async def astream(self) -> str: else self._underlying_value[beginning_length:] # type: ignore ) + def __aiter__(self) -> ModelOutputThunk[S]: + """Iterate the streamed deltas with `async for`. + + Wraps `astream()` in the async-iterator protocol so callers can write + `async for delta in mot:` rather than the manual + `while not mot.is_computed(): await mot.astream()` loop. Each iteration + yields a delta (the new text since the previous one); iteration ends when + generation completes. Subject to the same single-consumer constraint as + `astream()` — do not iterate the same thunk from multiple tasks. + + Returns: + ModelOutputThunk[S]: This thunk, acting as its own iterator. + """ + return self + + async def __anext__(self) -> str: + """Return the next streamed delta, or stop when generation is complete. + + Returns: + str: The new text received since the previous delta. + + Raises: + StopAsyncIteration: When the thunk is already computed. + """ + if self._computed: + raise StopAsyncIteration + return await self.astream() + def __str__(self) -> str: """Stringifies the thunk value.""" return self.value if self.value else "" diff --git a/mellea/stdlib/streaming_poc.py b/mellea/stdlib/streaming_poc.py new file mode 100644 index 000000000..2c719e51f --- /dev/null +++ b/mellea/stdlib/streaming_poc.py @@ -0,0 +1,368 @@ +"""Validated streaming: a thin layer over `ModelOutputThunk.__aiter__`. + +Functional replacement for `mellea/stdlib/streaming.py` (`stream_with_chunking` + +`StreamChunkingResult`). The design rests on two properties: + +1. Plain and validated streaming share one iteration protocol (`async for`). The + shared part (`ModelOutputThunk.__aiter__`) lives in core; chunking and + validation layer on top here in stdlib, so core keeps no dependency on stdlib. + +2. The whole stream runs on the caller's task — no background orchestration task, + so no queues, no raise-once exception plumbing, and no + `acomplete()`-vs-`astream()` finalization ambiguity. +""" + +from __future__ import annotations + +import asyncio +import uuid +from collections.abc import AsyncIterator, Sequence +from copy import copy +from typing import Any + +from ..backends.model_options import ModelOption +from ..core.backend import Backend +from ..core.base import CBlock, Component, Context, ModelOutputThunk +from ..core.requirement import PartialValidationResult, Requirement, ValidationResult +from ..plugins.hooks.streaming import ( + StreamingEndPayload, + StreamingEventPayload, + StreamingStartPayload, +) +from ..plugins.manager import has_plugins, invoke_hook +from ..plugins.types import HookType +from .chunking import ChunkingStrategy, ParagraphChunker, SentenceChunker, WordChunker + +# POC: reuse the Phase 1 event dataclasses rather than redefining them here. +from .streaming import ( + ChunkEvent, + CompletedEvent, + ErrorEvent, + FullValidationEvent, + QuickCheckEvent, + StreamEvent, + StreamingDoneEvent, +) + +_CHUNKING_ALIASES: dict[str, type[ChunkingStrategy]] = { + "sentence": SentenceChunker, + "word": WordChunker, + "paragraph": ParagraphChunker, +} + + +async def _emit_event( + streaming_id: str, ev: StreamEvent, *, requirements: list[Requirement] | None = None +) -> None: + """Fire the STREAMING_EVENT hook for `ev`. + + For a `QuickCheckEvent`, `requirements` carries the active requirement + instances in result order so a subscriber can attribute each result. + """ + if has_plugins(HookType.STREAMING_EVENT): + await invoke_hook( + HookType.STREAMING_EVENT, + StreamingEventPayload( + streaming_id=streaming_id, event=ev, requirements=requirements or [] + ), + ) + + +class Streamer: + """Async-iterable handle for a `stream` call. + + Consume the chunks with `async for`, then read `failed_early` / + `failure_reason`. On early exit, `streaming_failures` holds every + `(requirement, result)` that failed the offending chunk and `full_text` + holds the validated output through the last fully-passed delta. On natural + completion, `mot` holds the computed output and `final_validations` holds the + stream-end `validate()` results. + + Args: + mot: The in-flight streaming thunk from the backend generation call. + ctx: The generation context, used for validation calls. + chunking: Resolved chunking strategy, or `None` for raw deltas. + requirements: Requirements to validate against; pre-copied by `stream`. + validation_backend: Backend used for validation calls. + """ + + def __init__( + self, + mot: ModelOutputThunk, + ctx: Context, + chunking: ChunkingStrategy | None, + requirements: list[Requirement], + validation_backend: Backend, + ) -> None: + """Wrap an in-flight generation; iterating the `Streamer` drives it.""" + self.failed_early: bool = False + self.failure_reason: str | None = None + self.streaming_failures: list[tuple[Requirement, PartialValidationResult]] = [] + self.full_text: str = "" + self.mot: ModelOutputThunk | None = None + self.final_validations: list[ValidationResult] = [] + # Correlates this stream's START/END hooks; unique per concurrent stream. + self.streaming_id: str = str(uuid.uuid4()) + self._chunks: AsyncIterator[str] = _drive( + self, mot, ctx, chunking, requirements, validation_backend + ) + + def __aiter__(self) -> AsyncIterator[str]: + """Return the generator that drives generation and yields chunks.""" + return self._chunks + + +async def _validate_chunk( + streamer: Streamer, + chunk: str, + chunk_index: int, + requirements: list[Requirement], + validation_backend: Backend, + ctx: Context, + *, + on_flush: bool = False, +) -> bool: + """Run every requirement's `stream_validate` on `chunk`. + + Returns `True` when `chunk` passed and may be emitted (no requirements, or + all returned `"pass"`/`"unknown"`). Returns `False` when any requirement + fails — every failing `(requirement, result)` is recorded on `streamer` and + the caller should stop before yielding `chunk`. `on_flush` distinguishes a + failure on the trailing flushed fragment (stream already ended) from a + mid-stream one in the recorded reason. + """ + if not requirements: + return True + results = list( + await asyncio.gather( + *[ + req.stream_validate(chunk, backend=validation_backend, ctx=ctx) + for req in requirements + ] + ) + ) + failures = [ + (req, r) for req, r in zip(requirements, results) if r.success == "fail" + ] + await _emit_event( + streamer.streaming_id, + QuickCheckEvent( + chunk_index=chunk_index, attempt=1, passed=not failures, results=results + ), + requirements=requirements, + ) + if not failures: + return True + streamer.failed_early = True + streamer.streaming_failures.extend(failures) + where = " on flush" if on_flush else "" + streamer.failure_reason = ( + f"Streaming validation failed{where}: {failures[-1][1].reason or ''}" + ) + return False + + +async def _drive( + streamer: Streamer, + mot: ModelOutputThunk, + ctx: Context, + chunking: ChunkingStrategy | None, + requirements: list[Requirement], + validation_backend: Backend, +) -> AsyncIterator[str]: + """Drive the whole stream from one generator on the caller's task. + + A caller `break`/`aclose()` delivers `GeneratorExit` to the suspended `yield`, + so the single `finally` always runs — cleanup and STREAMING_END fire on every + exit path (natural end, early exit, caller break, exception). + + On natural completion every requirement's `validate()` runs on the full output + (early exit already returned, so all requirements reached the end unfailed); + this is what checks judge/aLoRA requirements that streamed only `"unknown"`. + """ + accumulated = "" + prev_chunk_count = 0 + chunk_index = 0 + success = False + error: Exception | None = None + + if has_plugins(HookType.STREAMING_START): + await invoke_hook( + HookType.STREAMING_START, + StreamingStartPayload( + streaming_id=streamer.streaming_id, + has_requirements=bool(requirements), + requirement_count=len(requirements), + chunking_strategy=type(chunking).__name__ if chunking else "none", + ), + ) + + try: + async for delta in mot: + accumulated += delta + + # chunking=None -> yield each raw delta as its own chunk. + if chunking is None: + new_chunks = [delta] if delta else [] + else: + chunks = chunking.split(accumulated) + new_chunks = chunks[prev_chunk_count:] + prev_chunk_count = len(chunks) + + for c in new_chunks: + if not await _validate_chunk( + streamer, c, chunk_index, requirements, validation_backend, ctx + ): + return + yield c + await _emit_event( + streamer.streaming_id, + ChunkEvent(text=c, chunk_index=chunk_index, attempt=1), + ) + chunk_index += 1 + + # Snapshot after a delta fully passes so `full_text` excludes any + # unvalidated chunk on early exit. + # TODO(#1013): delta-granular, not chunk-exact; MOT-owned chunking + # (one unit per iteration) makes it exact. + streamer.full_text = accumulated + + # Flush the trailing fragment the chunker withheld (skipped in raw mode). + if chunking is not None: + for c in chunking.flush(accumulated): + if not await _validate_chunk( + streamer, + c, + chunk_index, + requirements, + validation_backend, + ctx, + on_flush=True, + ): + return + yield c + await _emit_event( + streamer.streaming_id, + ChunkEvent(text=c, chunk_index=chunk_index, attempt=1), + ) + chunk_index += 1 + + # Natural completion: capture the flushed fragment the snapshot missed. + streamer.full_text = accumulated + streamer.mot = mot + await _emit_event( + streamer.streaming_id, StreamingDoneEvent(attempt=1, full_text=accumulated) + ) + + # Reached only on natural completion, so every requirement is still + # unfailed and gets a full-output validate(). + if requirements: + streamer.final_validations = list( + await asyncio.gather( + *[req.validate(validation_backend, ctx) for req in requirements] + ) + ) + await _emit_event( + streamer.streaming_id, + FullValidationEvent( + attempt=1, + passed=all(v.as_bool() for v in streamer.final_validations), + results=streamer.final_validations, + ), + ) + success = True + except Exception as exc: + # Record for the STREAMING_END span, then re-raise so the exception + # still propagates to the caller through the `async for`. + error = exc + await _emit_event( + streamer.streaming_id, + ErrorEvent(exception_type=type(exc).__name__, detail=str(exc)), + ) + raise + finally: + # Cancel on any early/broken exit so the backend producer never wedges + # on a full queue; no-op once the stream is fully drained. + if not mot.is_computed(): + await mot.cancel_generation() + + # Always the last StreamEvent, on every exit path. + await _emit_event( + streamer.streaming_id, + CompletedEvent( + success=success, full_text=streamer.full_text, attempts_used=1 + ), + ) + + if has_plugins(HookType.STREAMING_END): + await invoke_hook( + HookType.STREAMING_END, + StreamingEndPayload( + streaming_id=streamer.streaming_id, + success=success, + failure_reason=streamer.failure_reason, + exception=error, + model=mot.generation.model, + provider=mot.generation.provider, + full_text_length=len(accumulated), + ), + ) + + +async def stream( + action: Component[Any] | CBlock, + backend: Backend, + ctx: Context, + *, + chunking: str | ChunkingStrategy | None = "sentence", + requirements: Sequence[Requirement] | None = None, + validation_backend: Backend | None = None, +) -> Streamer: + """Start a streaming generation, optionally chunked and validated per chunk. + + Consume the returned `Streamer` with `async for`. Each iteration yields a + chunk once it has passed every requirement's `stream_validate`; a `"fail"` + stops the stream early and cancels the backend. On natural completion, + `validate()` runs on the full output. With no `requirements`, chunks are + yielded without validation. + + Args: + action: The component or content block to generate from. + backend: Backend used for generation and, unless `validation_backend` + is set, validation. + ctx: The generation context. + chunking: A `ChunkingStrategy`, one of the aliases `"sentence"`, + `"word"`, `"paragraph"`, or `None` to yield raw deltas unchunked. + requirements: Requirements validated against each chunk during + streaming and against the full output at stream end. `None` yields + chunks without validation. + validation_backend: Backend for validation calls; defaults to `backend`. + + Returns: + Streamer: An async-iterable handle over the validated chunks. + + Raises: + ValueError: If `chunking` is a string that is not a known alias. + RuntimeError: If the backend returns an already-computed thunk instead + of a streaming one (it is not honouring `ModelOption.STREAM`). + """ + if isinstance(chunking, str): + cls = _CHUNKING_ALIASES.get(chunking) + if cls is None: + raise ValueError(f"Unknown chunking alias {chunking!r}") + chunking = cls() + + # Copy so a raising __copy__ surfaces before generation starts, and the + # caller's requirement instances are never mutated by streaming state. + cloned_reqs = [copy(req) for req in (requirements or [])] + resolved_backend = validation_backend if validation_backend is not None else backend + + mot, gen_ctx = await backend.generate_from_context( + action, ctx, model_options={ModelOption.STREAM: True} + ) + if mot.is_computed(): + raise RuntimeError( + "stream() requires a streaming backend; got an already-computed MOT." + ) + + return Streamer(mot, gen_ctx, chunking, cloned_reqs, resolved_backend)