diff --git a/docs/dev/migrate-streaming-v0.8.md b/docs/dev/migrate-streaming-v0.8.md new file mode 100644 index 0000000000..68e051728d --- /dev/null +++ b/docs/dev/migrate-streaming-v0.8.md @@ -0,0 +1,181 @@ +# Migrating streaming from v0.7 to the single-task `stream()` API + +v0.7 exposed streaming validation through `stream_with_chunking()`, which +returned a `StreamChunkingResult` driven by a background orchestration task. You +observed it through `result.events()` (typed events) or `result.astream()` (raw +chunks), and called `result.acomplete()` to wait for the background task to +finish. + +That two-task model is replaced by a single-task primitive: `stream()` returns a +`Streamer` you consume directly with `async for`, on your own task. There is no +background orchestrator and no `acomplete()`. Typed events now come from the +`streaming_event` plugin hook rather than an iterator on the result. This is a +**breaking change with no deprecation shim** — call sites must be updated. + +## API mapping + +| v0.7 | v0.8 | Notes | +| --- | --- | --- | +| `stream_with_chunking(...)` | `stream(...)` | Same arguments, except `chunking` now defaults to `None` (raw deltas) instead of `"sentence"`; pass `chunking="sentence"` to preserve v0.7 chunk boundaries | +| returns `StreamChunkingResult` | returns `Streamer` | Consume with `async for`, ideally inside `async with` | +| `async for chunk in result.astream()` | `async for chunk in streamer` | Iterate the `Streamer` directly | +| `async for event in result.events()` | `@hook("streaming_event")` plugin | Events move to the hook (see below) | +| `await result.acomplete()` | *(removed)* | Consuming the stream drives it to completion | +| `STREAMING_ORCHESTRATION_START`/`_END` hooks | *(removed)* | No replacement — the whole run is on one task now, so there is nothing to reattach a span across | +| `result.completed` | `not streamer.failed_early` | | +| *(new)* | `streamer.completed_normally` | `True` only on natural completion; unlike `not failed_early`, it is `False` after an early `break` | +| `result.full_text` | `streamer.full_text` | Same | +| `result.streaming_failures` | `streamer.streaming_failures` | Same | +| `result.final_validations` | `streamer.final_validations` | Same | +| `result.as_thunk` | `streamer.mot` | Set on natural completion | +| `SentenceChunker` | `SentenceChunking` | Strategy classes renamed | +| `WordChunker` | `WordChunking` | | +| `ParagraphChunker` | `ParagraphChunking` | | + +Wrap consumption in `async with` so the generation is cancelled on every exit +path — an early `break` or an exception — instead of leaking an abandoned +background stream. + +## Chunking strategy renames + +The three built-in strategy classes were renamed from `...Chunker` to +`...Chunking`, freeing the `Chunker` name for the new stateful driver: + +```python +# v0.7 +from mellea.stdlib.chunking import SentenceChunker, WordChunker, ParagraphChunker + +# v0.8 +from mellea.stdlib.chunking import SentenceChunking, WordChunking, ParagraphChunking +``` + +This only affects code that imports a strategy class by name (for example, to +subclass it or pass an instance). Passing a string alias — `chunking="sentence"`, +`"word"`, or `"paragraph"` — is unchanged. + +## Before and after: observing events + +If you consumed typed events with `result.events()`, the events move to a +`streaming_event` plugin. Both snippets below produce the same output — they are +the `main()` from `docs/examples/streaming/validated_streaming.py`, v0.7 then +v0.8, using the same requirement, prompt, and chunking. + +### v0.7 + +```python +result = await stream_with_chunking( + action, backend, ctx, requirements=[req], chunking="sentence" +) + +print("Streaming events as they arrive:") +async for event in result.events(): + match event: + case ChunkEvent(): + print(f" CHUNK[{event.chunk_index}]: {event.text!r}") + case QuickCheckEvent(passed=False): + print( + f" QUICK_CHECK[{event.chunk_index}]: FAIL — " + f"{event.results[0].reason if event.results else 'unknown reason'}" + ) + case QuickCheckEvent(): + print(f" QUICK_CHECK[{event.chunk_index}]: pass") + case StreamingDoneEvent(): + print(f" STREAMING_DONE: {len(event.full_text)} chars accumulated") + case FullValidationEvent(): + print(f" FULL_VALIDATION: {'PASS' if event.passed else 'FAIL'}") + case CompletedEvent(): + print(f" COMPLETED: success={event.success}") + case _: + pass + +await result.acomplete() + +print(f"\nCompleted normally: {result.completed}") +print(f"Full text: {result.full_text!r}") +``` + +### v0.8 + +```python +@hook("streaming_event") +async def print_events(payload, ctx) -> None: + event = payload.event + match event: + case ChunkEvent(): + print(f" CHUNK[{event.chunk_index}]: {event.text!r}") + case QuickCheckEvent(passed=False): + print( + f" QUICK_CHECK[{event.chunk_index}]: FAIL — " + f"{event.results[0].reason if event.results else 'unknown reason'}" + ) + case QuickCheckEvent(): + print(f" QUICK_CHECK[{event.chunk_index}]: pass") + case StreamingDoneEvent(): + print(f" STREAMING_DONE: {len(event.full_text)} chars accumulated") + case FullValidationEvent(): + print(f" FULL_VALIDATION: {'PASS' if event.passed else 'FAIL'}") + case CompletedEvent(): + print(f" COMPLETED: success={event.success}") + case _: + pass + + +register(print_events) + +print("Stream events as they arrive:") +async with await stream( + action, backend, ctx, requirements=[req], chunking="sentence" +) as streamer: + # Draining the stream fires the events; the hook does the printing. + async for _chunk in streamer: + pass + +print(f"\nCompleted normally: {streamer.completed_normally}") +print(f"Full text: {streamer.full_text!r}") +``` + +The three transformations to make: + +1. **`stream_with_chunking(...)` → `async with await stream(...) as streamer:`**, + and iterate `streamer` directly. The `async with` replaces `acomplete()` for + cleanup. +2. **The `result.events()` match loop → a `@hook("streaming_event")` plugin.** + The event vocabulary and payloads are unchanged; register the plugin (with + `register()` or a `plugin_scope`) before consuming. Draining the stream is + what fires the events. +3. **`result.` → `streamer.`**. `result.completed` maps directly to + `not streamer.failed_early`; prefer the new `streamer.completed_normally` if + you need a signal that also excludes an early `break` (used above). + +## Before and after: consuming raw chunks + +If you never used `events()` and only consumed the validated chunk text with +`result.astream()`, the migration is smaller — iterate the `Streamer` directly +and drop `acomplete()`: + +### v0.7 + +```python +result = await stream_with_chunking( + action, backend, ctx, requirements=[req], chunking="sentence" +) +async for chunk in result.astream(): + print(chunk) +await result.acomplete() +``` + +### v0.8 + +```python +async with await stream( + action, backend, ctx, requirements=[req], chunking="sentence" +) as streamer: + async for chunk in streamer: + print(chunk) +``` + +## See also + +- [Streaming validation tutorial](../docs/tutorials/06-streaming-validation.md) +- [How-to: Async and streaming](../docs/how-to/use-async-and-streaming.md) +- [`docs/examples/streaming/`](../examples/streaming/) diff --git a/docs/docs/concepts/requirements-system.md b/docs/docs/concepts/requirements-system.md index 42def3fe38..a998a2ef9f 100644 --- a/docs/docs/concepts/requirements-system.md +++ b/docs/docs/concepts/requirements-system.md @@ -342,8 +342,8 @@ field: - `"fail"` — the stream is cancelled immediately; no further chunks reach the consumer; `validate()` is skipped for this requirement. -State isolation is per-clone: `stream_with_chunking()` copies each requirement -with `copy()` before starting the orchestrator, so the original objects are never +State isolation is per-clone: `stream()` copies each requirement +with `copy()` before generation starts, so the original objects are never mutated. Requirements that accumulate state across chunks (e.g. a running word count) should reassign mutable containers rather than mutate in place, since clones share the original's `__dict__` values at copy time. diff --git a/docs/docs/examples/index.md b/docs/docs/examples/index.md index cdda1b8bb4..5cecdd8886 100644 --- a/docs/docs/examples/index.md +++ b/docs/docs/examples/index.md @@ -32,7 +32,7 @@ to run. | `context/` | Context inspection, sampling with context trees, parallel context branches | | `sessions/` | Custom session types and backend selection | | `async/` | How to utilize basic async capabilities | -| `streaming/` | `stream_with_chunking()` with per-chunk validation, typed event vocabulary, early-exit on fail | +| `streaming/` | `stream()` with per-chunk validation, typed event vocabulary, early-exit on fail | ### Data and documents diff --git a/docs/docs/how-to/use-async-and-streaming.md b/docs/docs/how-to/use-async-and-streaming.md index 33989e2706..bca862bf17 100644 --- a/docs/docs/how-to/use-async-and-streaming.md +++ b/docs/docs/how-to/use-async-and-streaming.md @@ -142,6 +142,40 @@ How `astream()` behaves: > **Warning:** Do not call `astream()` from multiple coroutines simultaneously on > the same thunk. Each thunk should have a single reader. +### Iterating a thunk with `async for` + +`astream()` is the low-level primitive. To consume a thunk as an async iterator, +use `async for` — ideally inside `async with`, which cancels the generation if +you leave the loop early (an exception or `break`) so an abandoned stream does +not keep running: + +```python +# Requires: mellea +# Returns: None +import asyncio +import mellea +from mellea.backends import ModelOption + +async def main(): + m = mellea.start_session() + mot = await m.ainstruct( + "Write a short story about a robot learning to cook.", + model_options={ModelOption.STREAM: True}, + ) + + async with mot: + async for delta in mot: + print(delta, end="", flush=True) + print() # newline after streaming completes + +asyncio.run(main()) +# Output will vary — LLM responses depend on model and temperature. +``` + +Each iteration yields the same delta `astream()` would return; iteration ends +when the thunk is computed. Like `astream()`, a thunk has a single reader — a +second `async for` over the same thunk raises rather than splitting the stream. + ### Streaming timeout Mellea waits up to 120 seconds for each chunk by default, including the first @@ -207,14 +241,15 @@ For parallel generation, use `SimpleContext`. ## Streaming with per-chunk validation -`stream_with_chunking()` adds per-chunk validation to a streaming generation. -It splits the accumulated text into semantic units (sentences, words, or -paragraphs), calls `stream_validate()` on each chunk in parallel, and can -exit early if any requirement returns `"fail"` — preventing the consumer from -seeing invalid content mid-stream. +`stream()` adds per-chunk validation to a streaming generation. It splits the +accumulated text into semantic units (sentences, words, or paragraphs), calls +`stream_validate()` on each chunk in parallel, and can exit early if any +requirement returns `"fail"` — preventing the consumer from seeing invalid +content mid-stream. -The primary way to observe a `stream_with_chunking()` run is via typed -`StreamEvent` objects from `result.events()`: +`stream()` returns a `Streamer` you consume with `async for`, ideally inside +`async with` so the generation is released on every exit path (including an +early `break` or exception): ```python # Requires: mellea @@ -225,14 +260,7 @@ from mellea.core.backend import Backend from mellea.core.base import Context from mellea.core.requirement import PartialValidationResult, Requirement, ValidationResult from mellea.stdlib.components import Instruction -from mellea.stdlib.streaming import ( - ChunkEvent, - CompletedEvent, - FullValidationEvent, - QuickCheckEvent, - StreamingDoneEvent, - stream_with_chunking, -) +from mellea.stdlib.streaming import stream class MaxSentencesReq(Requirement): @@ -267,48 +295,24 @@ async def main() -> None: action = Instruction("Write a two-sentence summary of the water cycle.") req = MaxSentencesReq(limit=3) - result = await stream_with_chunking( + async with await stream( action, m.backend, m.ctx, requirements=[req], chunking="sentence" - ) + ) as streamer: + async for chunk in streamer: + print(chunk) - async for event in result.events(): - match event: - case ChunkEvent(): - print(f" chunk[{event.chunk_index}]: {event.text!r}") - case QuickCheckEvent(passed=False): - print(f" FAIL at chunk {event.chunk_index}: {event.results}") - case StreamingDoneEvent(): - print(f" stream done — {len(event.full_text)} chars") - case FullValidationEvent(): - print(f" final: {'pass' if event.passed else 'fail'}") - case CompletedEvent(): - print(f" completed — success={event.success}") - case _: - pass # ErrorEvent and other future types - - await result.acomplete() - print(f"completed={result.completed}, failures={len(result.streaming_failures)}") + # Terminal state on the Streamer, after the loop. + print(f"Completed normally: {streamer.completed_normally}") + for _req, result in streamer.streaming_failures: + print(f"Streaming failure: {result.reason}") asyncio.run(main()) ``` -If you only need the raw validated text without event metadata, use -`result.astream()` instead: - -```python -result = await stream_with_chunking( - action, m.backend, m.ctx, requirements=[req], chunking="sentence" -) -async for chunk in result.astream(): - print(chunk) -await result.acomplete() -``` - -Both `astream()` (raw chunks) and `events()` are available on the same result -object. They use independent queues, so you can run them concurrently with -`asyncio.gather`. Both are **single-consumer** — a second iteration on either -will block indefinitely. +To observe the run through typed `StreamEvent` objects instead, register a +plugin on the `streaming_event` hook — see the +[streaming validation tutorial](../tutorials/06-streaming-validation.md). ### The `stream_validate` tri-state diff --git a/docs/docs/observability/tracing.md b/docs/docs/observability/tracing.md index 9bcfaef05c..2161c28ed6 100644 --- a/docs/docs/observability/tracing.md +++ b/docs/docs/observability/tracing.md @@ -193,16 +193,16 @@ One per requirement-validation batch, wherever requirements are checked. | `mellea.validation.failed_count` | Number of requirements that failed | | `mellea.validation.failure_reasons` | List of failing requirements' reasons; recorded only when `MELLEA_TRACES_CONTENT=true` | -#### `stream_with_chunking` span +#### `stream` span -One per `stream_with_chunking()` run, wrapping the backend generation and any -per-chunk validation. +One per `stream()` run, wrapping the backend generation and any per-chunk +validation. | Attribute | Description | | --------- | ----------- | | `mellea.streaming.has_requirements` | Whether requirements were supplied | | `mellea.streaming.requirement_count` | Number of requirements supplied | -| `mellea.streaming.chunking_strategy` | `ChunkingStrategy` class name (e.g., `SentenceChunker`) | +| `mellea.streaming.chunking_strategy` | `ChunkingStrategy` class name (e.g., `SentenceChunking`) | | `mellea.streaming.full_text_length` | Length of the accumulated text at completion | | `gen_ai.request.model` | Model ID, when known | | `gen_ai.provider.name` | Provider name, when known | @@ -283,21 +283,20 @@ Tool execution happens after the generating call completes, so `execute_tool` spans are not nested inside the `action` that requested them. Outside a session they are root spans. -In a `stream_with_chunking` run, the backend generation and each per-chunk -validation call nest under the `stream_with_chunking` span as sibling `chat` -spans: +In a `stream` run, the backend generation and each per-chunk validation call +nest under the `stream` span as sibling `chat` spans: ```text -stream_with_chunking (mellea.application) -│ [mellea.streaming.chunking_strategy=SentenceChunker] +stream (mellea.application) +│ [mellea.streaming.chunking_strategy=SentenceChunking] ├── chat (mellea.backend) ← streaming generation │ [gen_ai.request.model=granite4.1:3b] └── chat (mellea.backend) ← per-chunk validation [gen_ai.request.model=granite4.1:3b] ``` -The `stream_with_chunking` span itself parents under whatever span is active -when the run starts, or is a root span when none is. +The `stream` span itself parents under whatever span is active when the run +starts, or is a root span when none is. > **Note:** Full span nesting requires Python 3.12+. On Python 3.11 some spans > may appear flattened rather than nested; all spans and attributes are still diff --git a/docs/docs/tutorials/02-streaming-and-async.md b/docs/docs/tutorials/02-streaming-and-async.md index b0b33896ad..9844163a1c 100644 --- a/docs/docs/tutorials/02-streaming-and-async.md +++ b/docs/docs/tutorials/02-streaming-and-async.md @@ -111,6 +111,48 @@ How `astream()` works: a pre-computed thunk and `is_computed()` is already `True` before the loop runs. - Do not call `astream()` from multiple coroutines on the same thunk simultaneously. +### Iterating with `async for` + +The thunk is also an async iterator, so you can consume the same chunks with +`async for` instead of the manual `is_computed()` loop. Iterate inside +`async with`, which cancels the generation if you leave the loop early (an +exception or `break`) so an abandoned stream never keeps running: + +```python +# Requires: mellea +# Returns: str +import asyncio +import mellea +from mellea.backends import ModelOption + +async def stream_summary(feedback: str) -> str: + m = mellea.start_session() + mot = await m.ainstruct( + "Summarise this customer feedback in one sentence: {{text}}", + user_variables={"text": feedback}, + model_options={ModelOption.STREAM: True}, + strategy=None, + ) + + chunks = [] + async with mot: + async for chunk in mot: + print(chunk, end="", flush=True) + chunks.append(chunk) + print() # newline after streaming completes + + return "".join(chunks) + +asyncio.run(stream_summary( + "The onboarding was confusing and took far too long. " + "Support was helpful once I got through." +)) +``` + +Each iteration yields the same chunk `astream()` would return; iteration ends +when the thunk is computed. Like `astream()`, a thunk has a single reader — a +second `async for` over the same thunk raises rather than splitting the stream. + --- ## Step 3: Concurrent batch processing diff --git a/docs/docs/tutorials/06-streaming-validation.md b/docs/docs/tutorials/06-streaming-validation.md index 220ce789ec..8490d9ac8c 100644 --- a/docs/docs/tutorials/06-streaming-validation.md +++ b/docs/docs/tutorials/06-streaming-validation.md @@ -13,12 +13,13 @@ moment a requirement fails. By the end you will have covered: -- `stream_with_chunking()` — the streaming validation entry point -- The typed event vocabulary (`ChunkEvent`, `QuickCheckEvent`, …) from `result.events()` +- `stream()` — the streaming validation entry point +- Consuming validated chunks with `async for` inside `async with` - Early-exit cancellation and reading `streaming_failures` - Choosing between `"word"`, `"sentence"`, and `"paragraph"` chunking +- Observing the typed event vocabulary (`ChunkEvent`, `QuickCheckEvent`, …) + through the `streaming_event` hook - Subclassing `ChunkingStrategy` to define a custom split boundary -- `result.astream()` for consumers that only need the validated chunks **Prerequisites:** [Tutorial 02](./streaming-and-async) (async and streaming), [Tutorial 04](./making-agents-reliable) (requirements and validation), @@ -28,11 +29,11 @@ By the end you will have covered: ## Step 1: Your first streaming validation call -`stream_with_chunking()` returns a `StreamChunkingResult` immediately. The -orchestrator runs in the background, splitting accumulated text into chunks and -calling `stream_validate()` on each one. Consume events with `result.events()`, -then call `result.acomplete()` to wait for the orchestrator to finish and raise -any exception it stored. +`stream()` starts a streaming generation and returns a `Streamer`. Consume it +with `async for` inside `async with`: each iteration yields the next chunk, +already validated against every requirement, and the `async with` block cancels +the generation on every exit path — including an early `break` or an exception — +so an abandoned stream never keeps running in the background. ```python # Requires: mellea @@ -44,14 +45,7 @@ from mellea.core.backend import Backend from mellea.core.base import Context from mellea.core.requirement import PartialValidationResult, Requirement, ValidationResult from mellea.stdlib.components import Instruction -from mellea.stdlib.streaming import ( - ChunkEvent, - CompletedEvent, - FullValidationEvent, - QuickCheckEvent, - StreamingDoneEvent, - stream_with_chunking, -) +from mellea.stdlib.streaming import stream _SENTENCE_END = re.compile(r"[.!?]+") @@ -88,64 +82,47 @@ async def main() -> None: m = start_session() - result = await stream_with_chunking( + async with await stream( Instruction("Write a two-sentence summary of how photosynthesis works."), m.backend, m.ctx, requirements=[MaxSentencesReq(limit=3)], chunking="sentence", - ) - - async for event in result.events(): - match event: - case ChunkEvent(): - print(f" chunk[{event.chunk_index}]: {event.text!r}") - case QuickCheckEvent(passed=False): - print(f" FAIL at chunk {event.chunk_index}: {event.results[0].reason}") - case StreamingDoneEvent(): - print(f" stream done — {len(event.full_text)} chars") - case FullValidationEvent(): - print(f" final validation: {'pass' if event.passed else 'fail'}") - case CompletedEvent(): - print(f" completed — success={event.success}") - case _: - pass - - await result.acomplete() - print(f"\nFull text: {result.full_text!r}") + ) as streamer: + async for chunk in streamer: + print(f" chunk: {chunk!r}") + + print(f"\nCompleted normally: {streamer.completed_normally}") + print(f"Full text: {streamer.full_text!r}") asyncio.run(main()) ``` ```text Sample output - chunk[0]: 'Photosynthesis is the process by which plants use sunlight, water, and carbon dioxide to produce glucose and oxygen.' - chunk[1]: 'This reaction takes place in the chloroplasts and is essential to nearly all life on Earth.' - stream done — 222 chars - final validation: pass - completed — success=True + chunk: 'Photosynthesis is the process by which plants use sunlight, water, and carbon dioxide to produce glucose and oxygen.' + chunk: 'This reaction takes place in the chloroplasts and is essential to nearly all life on Earth.' +Completed normally: True Full text: 'Photosynthesis is the process by which plants use sunlight, water, and carbon dioxide to produce glucose and oxygen. This reaction takes place in the chloroplasts and is essential to nearly all life on Earth.' ``` > **Note:** LLM output is non-deterministic. Your result will vary in wording. -Three things to notice: +Two things to notice: -- `stream_with_chunking()` is called with `await` but returns immediately — the - orchestrator runs as a background task. -- `result.events()` is an async iterator that yields one event per semantic - unit. The loop ends when the `CompletedEvent` is delivered. -- `result.acomplete()` must be called after the event loop drains to propagate - any orchestrator exception and to ensure the background task has fully settled. +- `stream()` is awaited to obtain the `Streamer`, but generation begins eagerly — + it is already running by the time you enter the `async for` loop. +- Terminal state is read from the `Streamer` **after** the loop: `failed_early`, + `full_text`, `streaming_failures`, and `final_validations`. --- ## Step 2: Early exit on failure -When `stream_validate()` returns `"fail"`, the orchestrator cancels the backend -immediately and stops the stream. No further chunks are delivered, and the -failure is recorded in `result.streaming_failures`. +When `stream_validate()` returns `"fail"`, the backend generation is cancelled +immediately and the loop ends. No further chunks are delivered, `failed_early` +becomes `True`, and the failure is recorded in `streamer.streaming_failures`. Lower the sentence limit so the model is likely to exceed it: @@ -159,7 +136,7 @@ from mellea.core.backend import Backend from mellea.core.base import Context from mellea.core.requirement import PartialValidationResult, Requirement, ValidationResult from mellea.stdlib.components import Instruction -from mellea.stdlib.streaming import ChunkEvent, CompletedEvent, QuickCheckEvent, stream_with_chunking +from mellea.stdlib.streaming import stream _SENTENCE_END = re.compile(r"[.!?]+") @@ -196,55 +173,41 @@ async def main() -> None: # Ask for five sentences but cap the requirement at two. # The stream should be cancelled after the third sentence arrives. - result = await stream_with_chunking( + async with await stream( Instruction("Write five sentences about the history of the internet."), m.backend, m.ctx, requirements=[MaxSentencesReq(limit=2)], chunking="sentence", - ) - - async for event in result.events(): - match event: - case ChunkEvent(): - print(f" chunk[{event.chunk_index}]: {event.text[:60]!r}...") - case QuickCheckEvent(passed=False): - print(f" CANCELLED at chunk {event.chunk_index}") - case CompletedEvent(): - print(f" completed — success={event.success}") - case _: - pass - - await result.acomplete() - - if result.streaming_failures: - req, pvr = result.streaming_failures[0] + ) as streamer: + async for chunk in streamer: + print(f" chunk: {chunk[:60]!r}...") + + if streamer.streaming_failures: + _req, pvr = streamer.streaming_failures[0] print(f"\nStreaming failure: {pvr.reason}") - print(f"Text at cancellation:\n{result.full_text!r}") + print(f"Text at cancellation:\n{streamer.full_text!r}") else: - print(f"\nFull text: {result.full_text!r}") + print(f"\nFull text: {streamer.full_text!r}") asyncio.run(main()) ``` ```text Sample output - chunk[0]: 'The internet began as ARPANET, a U.S. Defense Department pr'... - chunk[1]: 'In the 1980s, the network expanded beyond government use and'... - chunk[2]: 'Tim Berners-Lee invented the World Wide Web in 1989, transfo'... - CANCELLED at chunk 2 - completed — success=False + chunk: 'The internet began as ARPANET, a U.S. Defense Department pr'... + chunk: 'In the 1980s, the network expanded beyond government use an'... Streaming failure: Exceeded 2-sentence limit Text at cancellation: -'The internet began as ARPANET, a U.S. Defense Department project in the late 1960s. In the 1980s, the network expanded beyond government use and began connecting universities and research centres. Tim Berners-Lee invented the World Wide Web in 1989...' +'The internet began as ARPANET, a U.S. Defense Department project in the late 1960s. In the 1980s, the network expanded beyond government use and began connecting universities and research centres.' ``` > **Note:** Whether the stream is cancelled depends on whether the model > exceeds the limit. If the model happens to comply, `streaming_failures` will -> be empty and `result.completed` will be `True`. +> be empty and `failed_early` will be `False`. -`result.full_text` always contains the text accumulated up to the point where +`streamer.full_text` always contains the text accumulated up to the point where generation stopped — useful for debugging what the model produced before the requirement failed. @@ -260,15 +223,14 @@ The built-in strategies cover a coarse-to-fine spectrum: | `"sentence"` | `.`, `!`, `?` followed by whitespace | Grammar, coherence, per-sentence content rules | | `"paragraph"` | Two or more consecutive newlines | Topic coherence, citation presence, heading structure | -The trade-off is **latency vs context**. Word chunking fires after every token — +The trade-off is **latency vs context**. Word chunking fires after every word — maximum reaction speed, but each chunk carries only a single word. Paragraph chunking waits for blank lines — full paragraph context for the validator, but detection is later and may happen after the model has produced a large amount of invalid content. -To see the granularity difference concretely, switch to word chunking and print -every fifth word — so you can count how many more validation events fire compared -to Step 1's two sentences: +To see the granularity difference concretely, switch to word chunking and count +how many chunks arrive compared to Step 1's two sentences: ```python # Requires: mellea @@ -279,7 +241,7 @@ from mellea.core.backend import Backend from mellea.core.base import Context from mellea.core.requirement import PartialValidationResult, Requirement, ValidationResult from mellea.stdlib.components import Instruction -from mellea.stdlib.streaming import ChunkEvent, CompletedEvent, QuickCheckEvent, stream_with_chunking +from mellea.stdlib.streaming import stream _FORBIDDEN = {"deprecated", "legacy", "obsolete"} @@ -309,7 +271,8 @@ async def main() -> None: m = start_session() - result = await stream_with_chunking( + word_count = 0 + async with await stream( Instruction( "Describe three advantages of cloud-native development in two sentences." ), @@ -317,30 +280,17 @@ async def main() -> None: m.ctx, requirements=[ForbiddenWordReq()], chunking="word", - ) - - word_count = 0 - async for event in result.events(): - match event: - case ChunkEvent(): - word_count += 1 - # Print every fifth word to show how many events fire. - if word_count % 5 == 1: - print(f" word {word_count:>3}: {event.text!r}") - case QuickCheckEvent(passed=False): - print(f" CANCELLED at word {event.chunk_index}: {event.results[0].reason}") - case CompletedEvent(): - status = "CANCELLED" if not event.success else "ok" - print(f" {status} — {word_count} word events total") - case _: - pass - - await result.acomplete() - - if result.streaming_failures: - print(f"Failure: {result.streaming_failures[0][1].reason}") + ) as streamer: + async for word in streamer: + word_count += 1 + # Print every fifth word to show how many chunks arrive. + if word_count % 5 == 1: + print(f" word {word_count:>3}: {word!r}") + + if streamer.streaming_failures: + print(f"Failure: {streamer.streaming_failures[0][1].reason}") else: - print(f"Full text: {result.full_text!r}") + print(f"{word_count} word chunks total") asyncio.run(main()) @@ -355,28 +305,28 @@ asyncio.run(main()) word 26: 'costs,' word 31: 'deployments,' word 36: 'services.' - ok — 38 word events total -Full text: 'Cloud-native development enables scalable, resilient ...' +38 word chunks total ``` > **Note:** LLM output is non-deterministic. Your result will vary in wording. -The same two-sentence response that produced **2** `ChunkEvent` items with sentence -chunking now produces **38**. The validator fires on every word — maximum reaction -speed at the cost of per-chunk context. +The same two-sentence response that produced **2** chunks with sentence chunking +now produces **38**. The validator fires on every word — maximum reaction speed +at the cost of per-chunk context. -If a forbidden word appears, the stream stops at that word and no further -`ChunkEvent` items are emitted. To see early exit in action, change `_FORBIDDEN` -to include a common English word like `"and"` or `"the"`. +If a forbidden word appears, the stream stops at that word and no further chunks +are delivered. To see early exit in action, change `_FORBIDDEN` to include a +common English word like `"and"` or `"the"`. --- -## Step 4: Raw chunk access with `astream()` +## Step 4: Observing the event lifecycle -If you only need the validated chunks and do not want event metadata, use -`result.astream()` instead of `result.events()`. It yields the text of each -validated chunk as a plain string — useful for streaming output directly to a -UI buffer or building the response incrementally without a `match` dispatch: +The `async for` loop gives you validated chunks. To observe the full lifecycle — +per-chunk validation results, stream completion, final validation, errors — +subscribe to the `streaming_event` hook. `stream()` fires one typed `StreamEvent` +per lifecycle moment through this hook, so a plugin can watch a run without +touching the chunk iterator. ```python # Requires: mellea @@ -387,8 +337,16 @@ import re from mellea.core.backend import Backend from mellea.core.base import Context from mellea.core.requirement import PartialValidationResult, Requirement, ValidationResult +from mellea.plugins import hook, register from mellea.stdlib.components import Instruction -from mellea.stdlib.streaming import stream_with_chunking +from mellea.stdlib.streaming import ( + ChunkEvent, + CompletedEvent, + FullValidationEvent, + QuickCheckEvent, + StreamingDoneEvent, + stream, +) _SENTENCE_END = re.compile(r"[.!?]+") @@ -407,7 +365,9 @@ class MaxSentencesReq(Requirement): ) -> PartialValidationResult: self._count += len(_SENTENCE_END.findall(chunk)) if self._count > self._limit: - return PartialValidationResult("fail", reason=f"Exceeded {self._limit}-sentence limit") + return PartialValidationResult( + "fail", reason=f"Exceeded {self._limit}-sentence limit" + ) return PartialValidationResult("unknown") async def validate( @@ -416,45 +376,69 @@ class MaxSentencesReq(Requirement): return ValidationResult(result=self._count <= self._limit) +@hook("streaming_event") +async def print_events(payload, ctx) -> None: + event = payload.event + match event: + case ChunkEvent(): + print(f" chunk[{event.chunk_index}]: {event.text!r}") + case QuickCheckEvent(passed=False): + print(f" FAIL at chunk {event.chunk_index}: {event.results[0].reason}") + case StreamingDoneEvent(): + print(f" stream done — {len(event.full_text)} chars") + case FullValidationEvent(): + print(f" final validation: {'pass' if event.passed else 'fail'}") + case CompletedEvent(): + print(f" completed — success={event.success}") + case _: + pass + + async def main() -> None: from mellea.stdlib.session import start_session m = start_session() + register(print_events) - result = await stream_with_chunking( + # Draining the stream drives generation; print_events fires per event. + async with await stream( Instruction("Write a two-sentence summary of the water cycle."), m.backend, m.ctx, requirements=[MaxSentencesReq(limit=3)], chunking="sentence", - ) - - # astream() yields only validated chunk text — no event wrapper. - async for chunk in result.astream(): - print(chunk, end=" ", flush=True) - print() - - await result.acomplete() - print(f"completed={result.completed}") + ) as streamer: + async for _chunk in streamer: + pass asyncio.run(main()) ``` ```text Sample output -Water evaporates from oceans and lakes, rises into the atmosphere, and -condenses into clouds. Precipitation falls back to Earth as rain or -snow, replenishing rivers, lakes, and groundwater. -completed=True + chunk[0]: 'Water evaporates from oceans and lakes, rises into the atmosphere, and condenses into clouds.' + chunk[1]: 'Precipitation then falls back to Earth as rain or snow, replenishing rivers, lakes, and groundwater.' + stream done — 195 chars + final validation: pass + completed — success=True ``` > **Note:** LLM output is non-deterministic. Your result will vary in wording. -`astream()` and `events()` are independent — both are available on the same -result object and can even be consumed concurrently with `asyncio.gather`. Each -is **single-consumer**: calling either iterator a second time raises -`RuntimeError`. If you need chunks after the fact, capture them to a list -during iteration or read `result.full_text` after `acomplete()`. +The event vocabulary: + +- `ChunkEvent` — a validated chunk was delivered to the consumer. +- `QuickCheckEvent` — the result of validating one chunk; `passed=False` marks + the requirement failure that ends the stream. +- `StreamingDoneEvent` — the token stream finished (natural completion only). +- `FullValidationEvent` — the final `validate()` pass over the whole output. +- `CompletedEvent` — the stream exited; always the last event, on every path. +- `ErrorEvent` — an exception occurred mid-stream. + +Because the hook is global, a single plugin can observe many concurrent streams +at once — use `payload.streaming_id` to tell their events apart. See +[`docs/examples/streaming/multi_stream_events.py`](https://github.com/generative-computing/mellea/blob/main/docs/examples/streaming/multi_stream_events.py) +for a multi-stream consumer. --- @@ -466,14 +450,13 @@ and define your own split boundary. Two methods to implement: -- **`split(accumulated_text)`** — called on every new token delta. Return all - complete chunks found so far; withhold any trailing fragment. Must be - stateless: it receives the full accumulated text each time, not a delta. -- **`flush(accumulated_text)`** — called once at natural end of stream. Release - the withheld trailing fragment, or return `[]` to discard it. +- **`split(text)`** — return all complete chunks in `text`, withholding any + trailing fragment. Must be stateless and idempotent. +- **`flush(text)`** — called once at natural end of stream. Release the withheld + trailing fragment, or return `[]` to discard it. -Here is a `LineChunker` that splits on single newlines — natural for numbered -list output where each line is one item: +Here is a `LineChunking` strategy that splits on single newlines — natural for +numbered list output where each line is one item: ```python # Requires: mellea @@ -486,27 +469,25 @@ from mellea.core.base import Context from mellea.core.requirement import PartialValidationResult, Requirement, ValidationResult from mellea.stdlib.chunking import ChunkingStrategy from mellea.stdlib.components import Instruction -from mellea.stdlib.streaming import ChunkEvent, CompletedEvent, QuickCheckEvent, stream_with_chunking +from mellea.stdlib.streaming import stream _NUMBERED_LINE = re.compile(r"^\s*\d+[\.\)]\s") -class LineChunker(ChunkingStrategy): +class LineChunking(ChunkingStrategy): """Emits one complete line per chunk, splitting on single newlines.""" - def split(self, accumulated_text: str) -> list[str]: - if "\n" not in accumulated_text: + def split(self, text: str) -> list[str]: + if "\n" not in text: return [] - last_nl = accumulated_text.rfind("\n") - return [line for line in accumulated_text[:last_nl].split("\n") if line.strip()] + last_nl = text.rfind("\n") + return [line for line in text[:last_nl].split("\n") if line.strip()] - def flush(self, accumulated_text: str) -> list[str]: - if not accumulated_text: + def flush(self, text: str) -> list[str]: + if not text: return [] - last_nl = accumulated_text.rfind("\n") - trailing = ( - accumulated_text if last_nl == -1 else accumulated_text[last_nl + 1 :] - ).strip() + last_nl = text.rfind("\n") + trailing = (text if last_nl == -1 else text[last_nl + 1 :]).strip() return [trailing] if trailing else [] @@ -538,7 +519,7 @@ async def main() -> None: m = start_session() - result = await stream_with_chunking( + async with await stream( Instruction( "List five world capitals, one per line, numbered 1 through 5. " "Use the format: '1. City'. Output only the numbered list, nothing else." @@ -546,33 +527,27 @@ async def main() -> None: m.backend, m.ctx, requirements=[NumberedLineReq()], - chunking=LineChunker(), - ) - - async for event in result.events(): - match event: - case ChunkEvent(): - print(f" line[{event.chunk_index}]: {event.text.strip()!r}") - case QuickCheckEvent(passed=False): - print(f" FAIL: {event.results[0].reason}") - case CompletedEvent(): - print(f" completed — success={event.success}") - case _: - pass + chunking=LineChunking(), + ) as streamer: + async for line in streamer: + print(f" line: {line.strip()!r}") - await result.acomplete() + if streamer.streaming_failures: + print(f"FAIL: {streamer.streaming_failures[0][1].reason}") + else: + print("Completed normally") asyncio.run(main()) ``` ```text Sample output - line[0]: '1. London' - line[1]: '2. Paris' - line[2]: '3. Tokyo' - line[3]: '4. Ottawa' - line[4]: '5. Canberra' - completed — success=True + line: '1. London' + line: '2. Paris' + line: '3. Tokyo' + line: '4. Ottawa' + line: '5. Canberra' +Completed normally ``` > **Note:** LLM output is non-deterministic. Your result will vary in wording. @@ -585,11 +560,11 @@ before reaching `validate()`. Lines that do reach it have already passed individual chunks rather than the full output. Pass a `ChunkingStrategy` **instance** (not a string alias) to use a custom -chunker. The built-in chunkers (`WordChunker`, `SentenceChunker`, -`ParagraphChunker`) are also available as instances if you need to pass one +chunker. The built-in strategies (`WordChunking`, `SentenceChunking`, +`ParagraphChunking`) are also available as instances if you need to pass one explicitly or subclass to override `flush()`. -> **See also:** [`docs/examples/streaming/custom_chunking.py`](/examples) +> **See also:** [`docs/examples/streaming/custom_chunking.py`](https://github.com/generative-computing/mellea/blob/main/docs/examples/streaming/custom_chunking.py) > for an annotated version of this pattern with a more detailed `split()`/`flush()` > contract walkthrough. @@ -599,17 +574,16 @@ explicitly or subclass to override `flush()`. | Concept | What it gives you | | --- | --- | -| `stream_with_chunking()` + `requirements=` | Per-chunk validation with automatic early exit | -| `result.events()` | Typed event stream — observe every chunk, validation result, and lifecycle signal | -| `QuickCheckEvent(passed=False)` | Detect the moment a requirement fails, mid-stream | -| `result.streaming_failures` | List of `(requirement, PartialValidationResult)` pairs for failed checks | +| `stream()` + `requirements=` | Per-chunk validation with automatic early exit | +| `async for chunk in streamer` | Validated chunks as they arrive, inside `async with` for safe cleanup | +| `streamer.failed_early` / `streamer.streaming_failures` | Detect and inspect a mid-stream requirement failure | +| `streaming_event` hook | Typed event stream — observe every chunk, validation result, and lifecycle signal | | `"word"` / `"sentence"` / `"paragraph"` | Built-in chunking strategies trading reaction speed for context | | `ChunkingStrategy` subclass | Custom split boundaries for structured output (lists, code, CSV) | -| `result.astream()` | Raw validated chunks without event metadata | --- > **See also:** > [How-to: Streaming with per-chunk validation](../how-to/use-async-and-streaming#streaming-with-per-chunk-validation) | > [Concepts: The Requirements System — Streaming validation](../concepts/requirements-system#streaming-validation) | -> [Examples: streaming/](/examples) +> [Examples: streaming/](https://github.com/generative-computing/mellea/tree/main/docs/examples/streaming) diff --git a/docs/examples/async/README.md b/docs/examples/async/README.md index decf766bda..793aad298d 100644 --- a/docs/examples/async/README.md +++ b/docs/examples/async/README.md @@ -35,6 +35,16 @@ Demonstrates: - Creating ModelOutputThunk objects - Deferred computation patterns +### Streaming with the Async Iterator + +```bash +uv run async-iterator.py +``` + +Demonstrates: +- Consuming a streamed response with `async for delta in response` +- Using `async with` so an early exit cancels the in-flight generation + ## Key Concepts **Async Backend Operations**: Mellea backends support async methods (`ainstruct`, `aact`, `achat`) that allow concurrent execution without blocking. diff --git a/docs/examples/async/async-iterator.py b/docs/examples/async/async-iterator.py new file mode 100644 index 0000000000..59656989c0 --- /dev/null +++ b/docs/examples/async/async-iterator.py @@ -0,0 +1,31 @@ +# pytest: ollama, e2e + +"""Example of streaming a response with the `async for` iterator.""" + +import asyncio + +from mellea.backends.model_options import ModelOption +from mellea.core.base import ModelOutputThunk +from mellea.stdlib.session import start_session + +# Create a regular session. Works with functional interface as well. +m = start_session() + + +async def main() -> None: + response: ModelOutputThunk[str] = await m.ainstruct( + "Say 'We're Streaming Now!' and then add a fun fact!", + strategy=None, # Cannot perform lazy compute / top level streaming if using a strategy. + model_options={ + ModelOption.STREAM: True # Set streaming to True for top level streaming. + }, + ) + + # Iterate the thunk to receive each delta as it arrives. `async with` cancels + # the generation if we leave the loop early. + async with response: + async for delta in response: + print(delta) + + +asyncio.run(main()) diff --git a/docs/examples/streaming/README.md b/docs/examples/streaming/README.md index 685764c60e..9a02142c55 100644 --- a/docs/examples/streaming/README.md +++ b/docs/examples/streaming/README.md @@ -15,15 +15,15 @@ ollama serve ### Basic Streaming with Chunking ```bash -uv run streaming_chunking.py +uv run validated_streaming.py ``` Demonstrates: - Streaming token-by-token generation -- Sentence-level chunking via `stream_with_chunking()` +- Sentence-level chunking via `stream()` - Per-chunk validation with custom `stream_validate()` methods +- Accessing stream events via the STREAMING_EVENT hook - Early exit on validation failure -- Accessing stream events (`ChunkEvent`, `QuickCheckEvent`, `FullValidationEvent`) ### Word-Level Chunking @@ -58,6 +58,16 @@ Demonstrates: - Defining custom stream validators - Advanced streaming patterns +### Events Across Multiple Concurrent Streams + +```bash +uv run multi_stream_events.py +``` + +Demonstrates: +- Consuming `STREAMING_EVENT` hook events with a small plugin +- Correlating events across multiple concurrent streams by `streaming_id` + ## Key Concepts **Streaming**: Receive LLM output token-by-token in real-time instead of waiting for complete generation. @@ -66,11 +76,13 @@ Demonstrates: **Stream Validation**: Apply requirements at chunk level for early exit—stop generation when a constraint is violated. -**Stream Events**: Process stream events to monitor generation progress: +**Stream Events**: Process stream events through the `STREAMING_EVENT` hook to monitor generation progress: - `ChunkEvent` — A new chunk of text - `QuickCheckEvent` — Initial validation result - `FullValidationEvent` — Complete validation after full generation - `StreamingDoneEvent` — Generation complete +- `CompletedEvent` — Stream exited (success or not) +- `ErrorEvent` — An exception occurred mid-stream ## See Also diff --git a/docs/examples/streaming/custom_chunking.py b/docs/examples/streaming/custom_chunking.py index 5163e1290b..7e2b99d94a 100644 --- a/docs/examples/streaming/custom_chunking.py +++ b/docs/examples/streaming/custom_chunking.py @@ -3,24 +3,22 @@ """Streaming generation with a custom ChunkingStrategy subclass. Demonstrates: -- Subclassing :class:`~mellea.stdlib.chunking.ChunkingStrategy` to define a - new splitting boundary +- Subclassing `ChunkingStrategy` to define a new splitting boundary - Implementing `split()` (stateless, idempotent) and `flush()` (end-of-stream release of any withheld trailing fragment) -- Using the custom chunker with `stream_with_chunking()` in place of a string alias +- Using the custom chunker with `stream()` in place of a string alias - Validating line-by-line output from a numbered-list prompt -`LineChunker` splits on single newlines (`\\n`), emitting one line per -`stream_validate` call. It sits between :class:`~mellea.stdlib.chunking.WordChunker` -(one word) and :class:`~mellea.stdlib.chunking.SentenceChunker` (one sentence) in -granularity, and is a natural fit for list-formatted model output. +`LineChunking` splits on single newlines (`\\n`), emitting one line per +`stream_validate` call. It sits between `WordChunking` (one word) and +`SentenceChunking` (one sentence) in granularity, and is a natural fit for +list-formatted model output. Extension pattern: 1. Subclass `ChunkingStrategy`. - 2. Implement `split(accumulated_text)` — return all complete chunks found in - the accumulated text so far; withhold any trailing fragment. The method is - called on every new token delta, so it must be stateless and idempotent. - 3. Override `flush(accumulated_text)` to release the withheld trailing fragment + 2. Implement `split(text)` — return all complete chunks in `text`; + withhold any trailing fragment. It must be stateless and idempotent. + 3. Override `flush(text)` to release the withheld trailing fragment when the stream ends naturally. The default base implementation returns `[]` (fragment discarded); override it when the trailing fragment is semantically significant. @@ -28,6 +26,7 @@ import asyncio import re +from typing import Any from mellea.core.backend import Backend from mellea.core.base import Context @@ -36,6 +35,7 @@ Requirement, ValidationResult, ) +from mellea.plugins import hook, register from mellea.stdlib.chunking import ChunkingStrategy from mellea.stdlib.components import Instruction from mellea.stdlib.streaming import ( @@ -44,7 +44,7 @@ FullValidationEvent, QuickCheckEvent, StreamingDoneEvent, - stream_with_chunking, + stream, ) # Matches a leading list marker: "1.", "1)", "1 .", or a bare number followed @@ -52,11 +52,11 @@ _NUMBERED_LINE = re.compile(r"^\s*\d+[\.\)]\s") -class LineChunker(ChunkingStrategy): - """Splits accumulated text on single newlines, emitting one line per chunk. +class LineChunking(ChunkingStrategy): + """Splits text on single newlines, emitting one line per chunk. The line after the last `\\n` is withheld as a trailing fragment until - the stream ends and :meth:`flush` is called. Blank lines are skipped — + the stream ends and `flush()` is called. Blank lines are skipped — they carry no content for a line-level validator. This chunker is a good fit for numbered-list output, code listings, and @@ -64,38 +64,36 @@ class LineChunker(ChunkingStrategy): rather than sentence-ending punctuation or double newlines. """ - def split(self, accumulated_text: str) -> list[str]: + def split(self, text: str) -> list[str]: """Return all complete lines (up to the last newline). Args: - accumulated_text: The full text accumulated so far. + text: The text to split. Returns: Non-empty lines found before the last newline character. The text after the last newline is withheld as a trailing fragment. """ - if "\n" not in accumulated_text: + if "\n" not in text: return [] - last_nl = accumulated_text.rfind("\n") - complete_section = accumulated_text[:last_nl] + last_nl = text.rfind("\n") + complete_section = text[:last_nl] return [line for line in complete_section.split("\n") if line.strip()] - def flush(self, accumulated_text: str) -> list[str]: + def flush(self, text: str) -> list[str]: """Release the trailing line fragment at end of stream. Args: - accumulated_text: The full accumulated text at stream end. + text: The text whose trailing fragment to release. Returns: The text after the last newline as a single-element list (stripped), or an empty list if the text ends with a newline or is empty. """ - if not accumulated_text: + if not text: return [] - last_nl = accumulated_text.rfind("\n") - trailing = ( - accumulated_text if last_nl == -1 else accumulated_text[last_nl + 1 :] - ).strip() + last_nl = text.rfind("\n") + trailing = (text if last_nl == -1 else text[last_nl + 1 :]).strip() return [trailing] if trailing else [] @@ -103,7 +101,7 @@ class NumberedLineReq(Requirement): """Fails the stream if any line does not start with a list number. Each `stream_validate` call receives one complete line (from - :class:`LineChunker`). This requirement enforces that every line follows + `LineChunking`). This requirement enforces that every line follows the `N. item` format, catching unstructured paragraphs or stray headers that sneak into what should be a clean numbered list. """ @@ -142,15 +140,12 @@ async def main() -> None: "List five world capitals, one per line, numbered 1 through 5. " "Use the format: '1. City'. Output only the numbered list, nothing else." ) - chunker = LineChunker() + chunker = LineChunking() req = NumberedLineReq() - result = await stream_with_chunking( - action, backend, ctx, requirements=[req], chunking=chunker - ) - - print("Streaming events as they arrive (one ChunkEvent per line):") - async for event in result.events(): + @hook("streaming_event") + async def print_events(payload: Any, ctx: Any) -> None: + event = payload.event match event: case ChunkEvent(): print(f" LINE[{event.chunk_index}]: {event.text!r}") @@ -170,14 +165,22 @@ async def main() -> None: case _: pass - await result.acomplete() + register(print_events) - print(f"\nCompleted normally: {result.completed}") - if result.streaming_failures: - for _req, pvr in result.streaming_failures: + print("Stream events as they arrive (one per line):") + async with await stream( + action, backend, ctx, requirements=[req], chunking=chunker + ) as streamer: + # Draining the stream fires the events; the hook does the printing. + async for _line in streamer: + pass + + print(f"\nCompleted normally: {streamer.completed_normally}") + if streamer.streaming_failures: + for _req, pvr in streamer.streaming_failures: print(f"Streaming failure: {pvr.reason}") else: - print(f"Full text:\n{result.full_text}") + print(f"Full text:\n{streamer.full_text}") asyncio.run(main()) diff --git a/docs/examples/streaming/multi_stream_events.py b/docs/examples/streaming/multi_stream_events.py new file mode 100644 index 0000000000..3250d9632a --- /dev/null +++ b/docs/examples/streaming/multi_stream_events.py @@ -0,0 +1,123 @@ +# pytest: ollama, e2e + +"""Observing streaming events across multiple concurrent streams via a plugin. + +`stream()` yields validated chunks through `async for`, but its typed lifecycle +events (`ChunkEvent`, `QuickCheckEvent`, `StreamingDoneEvent`, +`FullValidationEvent`, `CompletedEvent`, `ErrorEvent`) are surfaced through the +`STREAMING_EVENT` plugin hook rather than the iterator. + +Demonstrates: +- Registering a `@hook("streaming_event")` function to receive every stream's events +- Running two streams concurrently with `asyncio.gather`, so their events arrive + interleaved +- Using `payload.streaming_id` to tell interleaved events apart — the same key a + real consumer (dashboard, metrics sink, websocket) would demultiplex on +""" + +import asyncio +from typing import Any + +from mellea.core.backend import Backend +from mellea.core.base import Context +from mellea.core.requirement import ( + PartialValidationResult, + Requirement, + ValidationResult, +) +from mellea.plugins import hook, register +from mellea.stdlib.components import Instruction +from mellea.stdlib.context import SimpleContext +from mellea.stdlib.streaming import ( + ChunkEvent, + CompletedEvent, + ErrorEvent, + FullValidationEvent, + QuickCheckEvent, + StreamingDoneEvent, + stream, +) + + +@hook("streaming_event") +async def print_streaming_events(payload: Any, ctx: Any) -> None: + """Print each streaming event live, tagged with its stream's id.""" + event = payload.event + match event: + case ChunkEvent(): + summary = f"CHUNK[{event.chunk_index}]: {event.text!r}" + case QuickCheckEvent(): + summary = f"QUICK_CHECK[{event.chunk_index}]: {'pass' if event.passed else 'FAIL'}" + case StreamingDoneEvent(): + summary = f"STREAMING_DONE: {len(event.full_text)} chars" + case FullValidationEvent(): + summary = f"FULL_VALIDATION: {'PASS' if event.passed else 'FAIL'}" + case CompletedEvent(): + summary = f"COMPLETED: success={event.success}" + case ErrorEvent(): + summary = f"ERROR: {event.exception_type}: {event.detail}" + case _: + summary = type(event).__name__ + print(f" [{payload.streaming_id[:8]}] {summary}") + + +class MaxSentencesReq(Requirement): + """Fails mid-stream once the response exceeds a sentence budget.""" + + def __init__(self, limit: int) -> None: + super().__init__() + self._limit = limit + self._count = 0 + + def format_for_llm(self) -> str: + return f"The response must be at most {self._limit} sentences long." + + async def stream_validate( + self, chunk: str, *, backend: Backend, ctx: Context + ) -> PartialValidationResult: + self._count += sum(chunk.count(p) for p in ".!?") + if self._count > self._limit: + return PartialValidationResult( + "fail", reason=f"Exceeded {self._limit} sentence limit mid-stream" + ) + return PartialValidationResult("unknown") + + async def validate( + self, + backend: Backend, + ctx: Context, + *, + format: type | None = None, + model_options: dict | None = None, + ) -> ValidationResult: + return ValidationResult(result=self._count <= self._limit) + + +async def main() -> None: + from mellea.stdlib.session import start_session + + # One backend serving two independent requests, each with a fresh context. + backend = start_session().backend + + register(print_streaming_events) + + async def run(prompt: str) -> None: + """Drive one validated stream to completion; the hook prints its events.""" + async with await stream( + Instruction(prompt), + backend, + SimpleContext(), + requirements=[MaxSentencesReq(limit=3)], + chunking="sentence", + ) as streamer: + async for _chunk in streamer: + pass + + # gather runs both at once, so events interleave — the id prefix says which stream. + await asyncio.gather( + run("Describe the water cycle in two sentences."), + run("Describe photosynthesis in two sentences."), + ) + + +asyncio.run(main()) diff --git a/docs/examples/streaming/paragraph_chunking.py b/docs/examples/streaming/paragraph_chunking.py index 32734b66c9..88784f4f98 100644 --- a/docs/examples/streaming/paragraph_chunking.py +++ b/docs/examples/streaming/paragraph_chunking.py @@ -1,25 +1,26 @@ # pytest: ollama, e2e -"""Streaming generation with per-paragraph validation using ParagraphChunker. +"""Streaming generation with per-paragraph validation using ParagraphChunking. Demonstrates: - Using the `"paragraph"` chunking alias for coarse-grained, structure-aware validation - A paragraph-length gate that cancels generation if any paragraph is too long -- How ParagraphChunker withholds text until a blank line (`\\n\\n`) is seen, +- How ParagraphChunking withholds text until a blank line (`\\n\\n`) is seen, then emits the entire paragraph as a single chunk -- The latency trade-off vs. SentenceChunker: fewer, larger chunks mean lower +- The latency trade-off vs. SentenceChunking: fewer, larger chunks mean lower validation overhead but later detection -ParagraphChunker splits on two or more consecutive newlines. Unlike -SentenceChunker, it waits for the model to produce a blank line before +ParagraphChunking splits on two or more consecutive newlines. Unlike +SentenceChunking, it waits for the model to produce a blank line before emitting anything — so if the model writes everything as one long paragraph -the stream completes before any chunk is emitted. Use ParagraphChunker when +the stream completes before any chunk is emitted. Use ParagraphChunking when the validation logic requires full paragraph context: topic coherence, heading structure, citation presence, or overall paragraph quality. """ import asyncio +from typing import Any from mellea.core.backend import Backend from mellea.core.base import Context @@ -28,6 +29,7 @@ Requirement, ValidationResult, ) +from mellea.plugins import hook, register from mellea.stdlib.components import Instruction from mellea.stdlib.streaming import ( ChunkEvent, @@ -35,7 +37,7 @@ FullValidationEvent, QuickCheckEvent, StreamingDoneEvent, - stream_with_chunking, + stream, ) _MAX_PARAGRAPH_WORDS = 60 @@ -45,7 +47,7 @@ class ParagraphLengthReq(Requirement): """Fails the stream if any paragraph exceeds a word-count limit. Each `stream_validate` call receives one complete paragraph (from - :class:`~mellea.stdlib.chunking.ParagraphChunker`). The validator counts + `ParagraphChunking`). The validator counts words and immediately fails the stream if the paragraph is too long. This lets you enforce a maximum paragraph length at generation time rather than post-processing. @@ -99,12 +101,9 @@ async def main() -> None: ) req = ParagraphLengthReq(max_words=_MAX_PARAGRAPH_WORDS) - result = await stream_with_chunking( - action, backend, ctx, requirements=[req], chunking="paragraph" - ) - - print("Streaming events as they arrive (one ChunkEvent per paragraph):") - async for event in result.events(): + @hook("streaming_event") + async def print_events(payload: Any, ctx: Any) -> None: + event = payload.event match event: case ChunkEvent(): word_count = len(event.text.split()) @@ -129,14 +128,22 @@ async def main() -> None: case _: pass - await result.acomplete() + register(print_events) - print(f"\nCompleted normally: {result.completed}") - if result.streaming_failures: - for _req, pvr in result.streaming_failures: + print("Stream events as they arrive (one per paragraph):") + async with await stream( + action, backend, ctx, requirements=[req], chunking="paragraph" + ) as streamer: + # Draining the stream fires the events; the hook does the printing. + async for _paragraph in streamer: + pass + + print(f"\nCompleted normally: {streamer.completed_normally}") + if streamer.streaming_failures: + for _req, pvr in streamer.streaming_failures: print(f"Streaming failure: {pvr.reason}") else: - print(f"Full text:\n{result.full_text}") + print(f"Full text:\n{streamer.full_text}") asyncio.run(main()) diff --git a/docs/examples/streaming/streaming_chunking.py b/docs/examples/streaming/validated_streaming.py similarity index 78% rename from docs/examples/streaming/streaming_chunking.py rename to docs/examples/streaming/validated_streaming.py index fc698fe669..cea216875d 100644 --- a/docs/examples/streaming/streaming_chunking.py +++ b/docs/examples/streaming/validated_streaming.py @@ -1,16 +1,18 @@ # pytest: ollama, e2e -"""Streaming generation with per-chunk validation using stream_with_chunking(). +"""Streaming generation with per-chunk validation using stream(). Demonstrates: - Subclassing Requirement to override stream_validate() for early-exit checks -- Calling stream_with_chunking() with sentence-level chunking -- Observing the full event vocabulary via events() as they arrive -- Awaiting full completion with acomplete() to access final_validations and full_text +- Calling stream() with sentence-level chunking +- Observing the typed StreamEvents via the STREAMING_EVENT hook as they arrive +- Driving the stream with `async with` + `async for` for safe cleanup +- Reading terminal state (failed_early, full_text, final_validations) after the loop """ import asyncio import re +from typing import Any from mellea.core.backend import Backend from mellea.core.base import Context @@ -19,6 +21,7 @@ Requirement, ValidationResult, ) +from mellea.plugins import hook, register from mellea.stdlib.components import Instruction from mellea.stdlib.streaming import ( ChunkEvent, @@ -26,7 +29,7 @@ FullValidationEvent, QuickCheckEvent, StreamingDoneEvent, - stream_with_chunking, + stream, ) # Crude sentence-terminator detector. A run of ``.``/``!``/``?`` counts once @@ -92,12 +95,9 @@ async def main() -> None: ) req = MaxSentencesReq(limit=3) - result = await stream_with_chunking( - action, backend, ctx, requirements=[req], chunking="sentence" - ) - - print("Streaming events as they arrive:") - async for event in result.events(): + @hook("streaming_event") + async def print_events(payload: Any, ctx: Any) -> None: + event = payload.event match event: case ChunkEvent(): print(f" CHUNK[{event.chunk_index}]: {event.text!r}") @@ -117,17 +117,25 @@ async def main() -> None: case _: pass # RetryEvent and any future event types - await result.acomplete() + register(print_events) + + print("Stream events as they arrive:") + async with await stream( + action, backend, ctx, requirements=[req], chunking="sentence" + ) as streamer: + # Draining the stream fires the events; the hook does the printing. + async for _chunk in streamer: + pass - print(f"\nCompleted normally: {result.completed}") - print(f"Full text: {result.full_text!r}") + print(f"\nCompleted normally: {streamer.completed_normally}") + print(f"Full text: {streamer.full_text!r}") - if result.streaming_failures: - for _req, pvr in result.streaming_failures: + if streamer.streaming_failures: + for _req, pvr in streamer.streaming_failures: print(f"Streaming failure: {pvr.reason}") - if result.final_validations: - for vr in result.final_validations: + if streamer.final_validations: + for vr in streamer.final_validations: print(f"Final validation: {'PASS' if vr.as_bool() else 'FAIL'}") diff --git a/docs/examples/streaming/word_chunking.py b/docs/examples/streaming/word_chunking.py index 0ca2517c42..4cbfdf37b8 100644 --- a/docs/examples/streaming/word_chunking.py +++ b/docs/examples/streaming/word_chunking.py @@ -1,25 +1,26 @@ # pytest: ollama, e2e -"""Streaming generation with per-word validation using WordChunker. +"""Streaming generation with per-word validation using WordChunking. Demonstrates: - Using the `"word"` chunking alias for the finest-grained validation - Detecting a forbidden word the moment it appears in the stream - Early-exit cancelling generation before the consumer sees the bad word -- How WordChunker compares to SentenceChunker in reaction time +- How WordChunking compares to SentenceChunking in reaction time -WordChunker splits on whitespace, so each `stream_validate` call receives +WordChunking splits on whitespace, so each `stream_validate` call receives exactly one word. This is the highest-sensitivity strategy: validation fires before the model has finished even the current clause, letting you catch prohibited content with minimal output produced. -The trade-off vs. SentenceChunker: validators that need sentence-level context +The trade-off vs. SentenceChunking: validators that need sentence-level context (grammar, coherence) cannot operate correctly at word granularity because each -chunk carries only a single token. Use WordChunker when the check is +chunk carries only a single token. Use WordChunking when the check is token-local — forbidden words, length budgets, numeric thresholds. """ import asyncio +from typing import Any from mellea.core.backend import Backend from mellea.core.base import Context @@ -28,6 +29,7 @@ Requirement, ValidationResult, ) +from mellea.plugins import hook, register from mellea.stdlib.components import Instruction from mellea.stdlib.streaming import ( ChunkEvent, @@ -35,7 +37,7 @@ FullValidationEvent, QuickCheckEvent, StreamingDoneEvent, - stream_with_chunking, + stream, ) # Words that must not appear in the model's response. @@ -45,9 +47,9 @@ class ForbiddenWordReq(Requirement): """Fails the stream immediately if a forbidden word appears. - Each `stream_validate` call receives a single word (from - :class:`~mellea.stdlib.chunking.WordChunker`). The check is O(1) - per word — set membership test — so it adds negligible latency. + Each `stream_validate` call receives a single word (from `WordChunking`). + The check is O(1) per word — set membership test — so it adds negligible + latency. """ def __init__(self, forbidden: set[str]) -> None: @@ -92,13 +94,12 @@ async def main() -> None: ) req = ForbiddenWordReq(forbidden=_FORBIDDEN) - result = await stream_with_chunking( - action, backend, ctx, requirements=[req], chunking="word" - ) - - print("Streaming events as they arrive (one per word):") word_count = 0 - async for event in result.events(): + + @hook("streaming_event") + async def print_events(payload: Any, ctx: Any) -> None: + nonlocal word_count + event = payload.event match event: case ChunkEvent(): word_count += 1 @@ -121,15 +122,23 @@ async def main() -> None: case _: pass - await result.acomplete() + register(print_events) - print(f"\nCompleted normally: {result.completed}") - if result.streaming_failures: - for _req, pvr in result.streaming_failures: + print("Stream events as they arrive (one per word):") + async with await stream( + action, backend, ctx, requirements=[req], chunking="word" + ) as streamer: + # Draining the stream fires the events; the hook does the printing. + async for _word in streamer: + pass + + print(f"\nCompleted normally: {streamer.completed_normally}") + if streamer.streaming_failures: + for _req, pvr in streamer.streaming_failures: print(f"Streaming failure: {pvr.reason}") - print(f"Text at cancellation: {result.full_text!r}") + print(f"Text at cancellation: {streamer.full_text!r}") else: - print(f"Full text: {result.full_text!r}") + print(f"Full text: {streamer.full_text!r}") asyncio.run(main()) diff --git a/mellea/core/base.py b/mellea/core/base.py index 4b6a28e1d7..01c3c597ee 100644 --- a/mellea/core/base.py +++ b/mellea/core/base.py @@ -830,6 +830,8 @@ def __init__( # Set computed to True if a value is passed in. self._computed: bool = True if value is not None else False self._cancelled: bool = False + # Guards `__aiter__` against a second consumer splitting the same stream. + self._aiter_started: bool = False # Additional fields that should be standardized across apis. self.tool_calls = tool_calls @@ -937,53 +939,54 @@ def _drain() -> None: # Drain before awaiting — unblocks any put() the task is stuck on. _drain() - if self._gen.generate is not None: - try: - await self._gen.generate - except asyncio.CancelledError: - # Re-raise if the *outer* task is being cancelled (Python 3.11+ - # task.cancelling() > 0) so we don't silently absorb external - # cancellation. For the inner task's own CancelledError (the - # expected result of .cancel() above), cancelling() is 0. - cur = asyncio.current_task() - if cur is not None and cur.cancelling() > 0: - raise - except Exception: - pass - - if self._gen.generate_extra is not None: - try: - await self._gen.generate_extra - except asyncio.CancelledError: - cur = asyncio.current_task() - if cur is not None and cur.cancelling() > 0: - raise - except Exception: - pass - - # Drain again for any final item the task put before terminating. - _drain() + try: + if self._gen.generate is not None: + try: + await self._gen.generate + except asyncio.CancelledError: + # Re-raise if the *outer* task is being cancelled (Python 3.11+ + # task.cancelling() > 0) so we don't silently absorb external + # cancellation. For the inner task's own CancelledError (the + # expected result of .cancel() above), cancelling() is 0. + cur = asyncio.current_task() + if cur is not None and cur.cancelling() > 0: + raise + except Exception: + pass - if has_plugins(HookType.GENERATION_ERROR): - from ..plugins.hooks.generation import GenerationErrorPayload + if self._gen.generate_extra is not None: + try: + await self._gen.generate_extra + except asyncio.CancelledError: + cur = asyncio.current_task() + if cur is not None and cur.cancelling() > 0: + raise + except Exception: + pass + finally: + # Drain again for any final item the task put before terminating. + _drain() + + if self._underlying_value is None: + self._underlying_value = "" + self._cancelled = True + self._computed = True - recorded: Exception = ( - error if error is not None else RuntimeError("Generation cancelled") - ) - await invoke_hook( - HookType.GENERATION_ERROR, - GenerationErrorPayload( - exception=recorded, - model_output=self, - generation_id=self._call.generation_id, - latency_ms=self._elapsed_ms(), - ), - ) + if has_plugins(HookType.GENERATION_ERROR): + from ..plugins.hooks.generation import GenerationErrorPayload - if self._underlying_value is None: - self._underlying_value = "" - self._cancelled = True - self._computed = True + recorded: Exception = ( + error if error is not None else RuntimeError("Generation cancelled") + ) + await invoke_hook( + HookType.GENERATION_ERROR, + GenerationErrorPayload( + exception=recorded, + model_output=self, + generation_id=self._call.generation_id, + latency_ms=self._elapsed_ms(), + ), + ) @property def cancelled(self) -> bool: @@ -1259,6 +1262,72 @@ 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. + + Single-consumer: the underlying stream has one cursor, so a second + iterator would split the same stream rather than start an independent + one. The first `__aiter__` marks the thunk consumed; a second call + raises. This mirrors the single-consumer constraint on `astream()`. + + Returns: + ModelOutputThunk[S]: This thunk, acting as its own iterator. + + Raises: + RuntimeError: If the thunk is already being iterated by another + consumer. + """ + if self._aiter_started: + raise RuntimeError( + "ModelOutputThunk is already being iterated; async iteration is " + "single-consumer. A second `async for` would split the same stream " + "rather than start an independent one." + ) + self._aiter_started = True + 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() + + async def aclose(self) -> None: + """Cancel an unfinished generation and release its resources. + + Makes `async for delta in mot` safe to abandon: iterating and breaking + out early otherwise leaves the backend generation running, since the + async-iterator protocol has no reliable finalizer. Idempotent — a no-op + once the thunk is computed (including after natural completion), so it is + safe to call after full iteration or more than once. + + Prefer `async with mot:` so this runs automatically on every exit path; + call `aclose()` directly only when not using the context manager. + """ + if not self._computed: + await self.cancel_generation() + + async def __aenter__(self) -> ModelOutputThunk[S]: + """Enter the async context manager, returning this thunk.""" + return self + + async def __aexit__(self, *exc_info: object) -> None: + """Exit the context manager, cancelling any in-flight generation.""" + await self.aclose() + def __str__(self) -> str: """Stringifies the thunk value.""" return self.value if self.value else "" diff --git a/mellea/core/requirement.py b/mellea/core/requirement.py index cb715a1e02..592595b177 100644 --- a/mellea/core/requirement.py +++ b/mellea/core/requirement.py @@ -332,22 +332,21 @@ async def stream_validate( are shared by reference under `copy()`. Reassign rather than mutate in place (`self._buffer = self._buffer + [chunk]`, not `self._buffer.append(chunk)`), or override `__copy__` for proper - isolation. If an override raises, the enclosing - :func:`~mellea.stdlib.streaming.stream_with_chunking` call aborts before - any backend generation starts and the exception propagates unchanged. + isolation. If an override raises, the enclosing `stream()` call aborts + before any backend generation starts and the exception propagates unchanged. Overrides with externally visible side effects (file writes, network calls) should perform them only after any logic that could raise, since the framework cannot roll them back. Implementations must not call `mot.astream()` or otherwise read the - underlying stream; the orchestrator is the single consumer of the MOT + underlying stream; the stream driver is the single consumer of the MOT stream (see `ModelOutputThunk.astream`). Requirements that need access to the text seen so far should accumulate it themselves from the `chunk` values they receive. Args: chunk: A single complete semantic chunk produced by the chunking - strategy (e.g. one sentence for `SentenceChunker`). This is + strategy (e.g. one sentence for `SentenceChunking`). This is the delta since the previous `stream_validate` call for this attempt, not the accumulated output. Requirements that need earlier context should retain it on `self` across calls. diff --git a/mellea/plugins/hooks/streaming.py b/mellea/plugins/hooks/streaming.py index 521fb8f6ed..19d0ba4464 100644 --- a/mellea/plugins/hooks/streaming.py +++ b/mellea/plugins/hooks/streaming.py @@ -11,11 +11,11 @@ class StreamingStartPayload(MelleaBasePayload): - """Payload for `streaming_start` — before a `stream_with_chunking` run starts. + """Payload for `streaming_start` — before a `stream` run starts. Attributes: streaming_id: UUID correlating start/event/end hooks for a single - `stream_with_chunking` run. + `stream` run. has_requirements: `True` when the orchestrator was given at least one `Requirement` to validate against. requirement_count: Number of `Requirement` instances supplied. @@ -45,7 +45,7 @@ class StreamingEventPayload(MelleaBasePayload): class StreamingEndPayload(MelleaBasePayload): - """Payload for `streaming_end` — when `stream_with_chunking` finishes. + """Payload for `streaming_end` — when `stream` finishes. Fires on every completing path: natural completion, validation-fail early-exit, and an unhandled exception. `success` and `exception` @@ -70,26 +70,3 @@ class StreamingEndPayload(MelleaBasePayload): model: str | None = None provider: str | None = None full_text_length: int = 0 - - -class StreamingOrchestrationStartPayload(MelleaBasePayload): - """Payload for `streaming_orchestration_start` — on the orchestration task, before the stream is drained. - - Attributes: - streaming_id: UUID correlating with the matching `streaming_start`. - """ - - streaming_id: str = "" - - -class StreamingOrchestrationEndPayload(MelleaBasePayload): - """Payload for `streaming_orchestration_end` — on the orchestration task, after the stream is drained. - - Fires on the same task as `streaming_orchestration_start`. - - Attributes: - streaming_id: UUID correlating with the matching - `streaming_orchestration_start`. - """ - - streaming_id: str = "" diff --git a/mellea/plugins/types.py b/mellea/plugins/types.py index 343c8bdc81..cddc6b865e 100644 --- a/mellea/plugins/types.py +++ b/mellea/plugins/types.py @@ -79,8 +79,6 @@ class HookType(StrEnum): STREAMING_START = "streaming_start" STREAMING_EVENT = "streaming_event" STREAMING_END = "streaming_end" - STREAMING_ORCHESTRATION_START = "streaming_orchestration_start" - STREAMING_ORCHESTRATION_END = "streaming_orchestration_end" # Lazily populated mapping: hook_type -> (payload_class, result_class). @@ -126,8 +124,6 @@ def _build_hook_registry() -> dict[str, tuple[type, type]]: from mellea.plugins.hooks.streaming import ( StreamingEndPayload, StreamingEventPayload, - StreamingOrchestrationEndPayload, - StreamingOrchestrationStartPayload, StreamingStartPayload, ) from mellea.plugins.hooks.tool import ToolPostInvokePayload, ToolPreInvokePayload @@ -196,14 +192,6 @@ def _build_hook_registry() -> dict[str, tuple[type, type]]: HookType.STREAMING_START.value: (StreamingStartPayload, PluginResult), HookType.STREAMING_EVENT.value: (StreamingEventPayload, PluginResult), HookType.STREAMING_END.value: (StreamingEndPayload, PluginResult), - HookType.STREAMING_ORCHESTRATION_START.value: ( - StreamingOrchestrationStartPayload, - PluginResult, - ), - HookType.STREAMING_ORCHESTRATION_END.value: ( - StreamingOrchestrationEndPayload, - PluginResult, - ), } diff --git a/mellea/stdlib/__init__.py b/mellea/stdlib/__init__.py index b916bd9272..cd35e7df30 100644 --- a/mellea/stdlib/__init__.py +++ b/mellea/stdlib/__init__.py @@ -13,11 +13,10 @@ `mellea.stdlib.session` — for day-to-day use. Streaming chunking strategies (for use with streaming validation) are available at -`mellea.stdlib.chunking` and re-exported here for convenience. The core streaming -orchestration primitive :func:`~mellea.stdlib.streaming.stream_with_chunking` and -its result type :class:`~mellea.stdlib.streaming.StreamChunkingResult` are also -re-exported here, alongside the full :class:`~mellea.stdlib.streaming.StreamEvent` -vocabulary for typed event observation. +`mellea.stdlib.chunking` and re-exported here for convenience, alongside the +`Chunker` that drives them over a stream. The core streaming primitive `stream()` +and its async-iterable handle `Streamer` are also re-exported here, alongside the +full `StreamEvent` vocabulary for typed event observation. Low-level primitives for tool execution are available in `mellea.stdlib.functional`: `call_tools` and `acall_tools` for executing model-requested tool calls with full @@ -26,31 +25,38 @@ the generated tools. These primitives are rarely needed outside custom agentic loops. """ -from .chunking import ChunkingStrategy, ParagraphChunker, SentenceChunker, WordChunker +from .chunking import ( + Chunker, + ChunkingStrategy, + ParagraphChunking, + SentenceChunking, + WordChunking, +) from .streaming import ( ChunkEvent, CompletedEvent, ErrorEvent, FullValidationEvent, QuickCheckEvent, - StreamChunkingResult, + Streamer, StreamEvent, StreamingDoneEvent, - stream_with_chunking, + stream, ) __all__ = [ "ChunkEvent", + "Chunker", "ChunkingStrategy", "CompletedEvent", "ErrorEvent", "FullValidationEvent", - "ParagraphChunker", + "ParagraphChunking", "QuickCheckEvent", - "SentenceChunker", - "StreamChunkingResult", + "SentenceChunking", "StreamEvent", + "Streamer", "StreamingDoneEvent", - "WordChunker", - "stream_with_chunking", + "WordChunking", + "stream", ] diff --git a/mellea/stdlib/chunking.py b/mellea/stdlib/chunking.py index 1343aeae08..f566190a22 100644 --- a/mellea/stdlib/chunking.py +++ b/mellea/stdlib/chunking.py @@ -1,54 +1,55 @@ # Copyright IBM Corp. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -"""ChunkingStrategy ABC and built-in implementations for streaming validation.""" +"""ChunkingStrategy, its built-in implementations, and the Chunker driver. + +A `ChunkingStrategy` is the stateless "how to split": given text, it returns the +complete chunks and withholds any trailing fragment. + +A `Chunker` is the stateful "how far through this stream we are": it wraps a +strategy and drives it incrementally, feeding one delta at a time and holding the +trailing fragment between calls. +""" import re from abc import ABC, abstractmethod -__all__ = ["ChunkingStrategy", "ParagraphChunker", "SentenceChunker", "WordChunker"] +__all__ = [ + "Chunker", + "ChunkingStrategy", + "ParagraphChunking", + "SentenceChunking", + "WordChunking", + "resolve_chunking_strategy", +] class ChunkingStrategy(ABC): """Abstract base class for text chunking strategies used in streaming validation. - A chunking strategy receives the full accumulated text so far and returns a - list of complete chunks ready for downstream validation. Any trailing fragment - that has not yet reached a chunk boundary is withheld — it is not included in - the returned list. Each call is stateless and idempotent given the same input. - - **Performance:** `split()` is called on every streaming delta, re-scanning - the full accumulated text each time (O(n) in total accumulated length per - call). The orchestrator tracks `prev_chunk_count` to extract only the new - chunks. This keeps the chunker stateless and removes the need for `reset()` - or deep-copy support, at the cost of re-scanning text already seen. For - typical model outputs (a few KB) the cost is negligible; for very long - streams, a stateful chunker that only processes the new delta would be more - efficient. - - End-of-stream contract: `split()` always withholds the trailing fragment. - When the stream terminates, callers are responsible for processing any remainder: - take the full accumulated text, identify everything after the last returned - chunk boundary, and handle it appropriately (e.g. pass to a final validator - or discard). + A chunking strategy receives text and returns a list of complete chunks + ready for downstream validation. Any trailing fragment that has not yet reached + a chunk boundary is withheld — it is not included in the returned list. Each + call is stateless and idempotent given the same input. + + End-of-stream contract: `split()` always withholds the trailing fragment; + `flush()` releases it once no more text is coming. Note: this ABC operates on text streams only. Multi-modal output (audio - segments, image regions) is not supported — the `accumulated_text: str` + segments, image regions) is not supported — the `text: str` signatures on `split` and `flush` preclude it. """ @abstractmethod - def split(self, accumulated_text: str) -> list[str]: - """Return complete chunks from accumulated_text, excluding any trailing fragment. + def split(self, text: str) -> list[str]: + """Return complete chunks from text, excluding any trailing fragment. + + Each returned chunk must be a verbatim substring of `text`; the `Chunker` + driver rejects a strategy that mutates chunk text (e.g. normalizes + whitespace inside a chunk). Dropping text between chunks is fine. Args: - accumulated_text: The full text accumulated so far, including all - previously seen tokens and the latest delta. Implementations - that scan this string are O(n) in accumulated length per call. - Stateful implementations that only process the new delta are - possible but must never mutate state on `self` in place — - use reassignment (`self._buf = self._buf + [x]`) so that - `copy()`-based cloning in the orchestrator works correctly. + text: The text to split. Returns: A list of complete chunks. If no chunk boundary has been reached yet, @@ -56,25 +57,25 @@ def split(self, accumulated_text: str) -> list[str]: """ ... - def flush(self, accumulated_text: str) -> list[str]: + def flush(self, text: str) -> list[str]: """Return any trailing fragment that `split` withheld. - Called once by the orchestrator after the stream has ended naturally - (not on early-exit cancellation). Gives the chunker a chance to - release the final fragment that did not reach a terminator. + Called once after the stream has ended naturally (not on early-exit + cancellation). Gives the strategy a chance to release the final fragment + that did not reach a terminator. The default implementation returns an empty list — the trailing fragment is discarded. Built-in chunkers override this to return the withheld fragment as a single-element list when non-empty. Args: - accumulated_text: The full accumulated text at stream end. + text: The text whose trailing fragment to release. Returns: The trailing fragment as `[fragment]` if it should be treated as a final chunk, or an empty list to discard it. """ - _ = accumulated_text + _ = text return [] @@ -84,42 +85,44 @@ def flush(self, accumulated_text: str) -> list[str]: # quotes (U+201D, U+2019), and closing paren. _SENTENCE_BOUNDARY = re.compile("[.!?][\"'\u201d\u2019)]?\\s") -# Whitespace run separator used by WordChunker. +# Whitespace run separator used by WordChunking. _WHITESPACE = re.compile(r"\s+") -# Paragraph boundary patterns used by ParagraphChunker. +# Paragraph boundary patterns used by ParagraphChunking. _PARA_BOUNDARY = re.compile(r"\n{2,}") _PARA_BOUNDARY_END = re.compile(r"\n{2,}$") -class SentenceChunker(ChunkingStrategy): - """Splits accumulated text on sentence boundaries. +class SentenceChunking(ChunkingStrategy): + """Splits text on sentence boundaries. Sentence boundaries are detected by `.`, `!`, or `?`, optionally followed by a closing quote (straight or curly) or parenthesis, then whitespace. The final sentence is only returned once it is followed by whitespace or another sentence — a trailing fragment with no following whitespace is withheld. Abbreviations are a known edge case: they will - be split on (simple regex, not NLP). Inter-sentence whitespace (including - double-space or tab) is discarded and does not appear as leading whitespace - in subsequent chunks. + be split on (simple regex, not NLP). Leading and inter-sentence whitespace + (including double-space or tab) is discarded — no chunk, including the first, + begins with whitespace. """ - def split(self, accumulated_text: str) -> list[str]: - """Return complete sentences from accumulated_text. + def split(self, text: str) -> list[str]: + """Return complete sentences from text. Args: - accumulated_text: The full text accumulated so far. + text: The text to split. Returns: Complete sentences detected so far. The trailing fragment (if any) is withheld. """ - if not accumulated_text: + if not text: return [] chunks: list[str] = [] - remaining = accumulated_text + # lstrip the leading edge so the first chunk, like the rest, never starts + # with separator whitespace carried in from a prior boundary. + remaining = text.lstrip() while True: match = _SENTENCE_BOUNDARY.search(remaining) @@ -136,18 +139,15 @@ def split(self, accumulated_text: str) -> list[str]: return chunks - def flush(self, accumulated_text: str) -> list[str]: + def flush(self, text: str) -> list[str]: """Return the trailing sentence fragment (if any) as a final chunk. - Trailing whitespace on the fragment is non-semantic for sentence - boundaries and is dropped via `rstrip`. Leading whitespace is - already removed by the loop's `lstrip` on each advance, so no - `lstrip` is needed here. The result is the fragment's content - only, consistent with how :meth:`split` returns sentences without - trailing whitespace. + Leading and trailing whitespace on the fragment is non-semantic for + sentence boundaries and is stripped, consistent with how `split` returns + sentences with no surrounding whitespace. Args: - accumulated_text: The full accumulated text at stream end. + text: The text whose trailing fragment to release. Returns: A single-element list containing the trailing sentence fragment @@ -155,9 +155,9 @@ def flush(self, accumulated_text: str) -> list[str]: when there is no fragment (all content ended in a sentence boundary or the input is empty/whitespace-only). """ - if not accumulated_text: + if not text: return [] - remaining = accumulated_text + remaining = text.lstrip() while True: match = _SENTENCE_BOUNDARY.search(remaining) if match is None: @@ -167,29 +167,29 @@ def flush(self, accumulated_text: str) -> list[str]: return [trailing] if trailing else [] -class WordChunker(ChunkingStrategy): - """Splits accumulated text on whitespace boundaries. +class WordChunking(ChunkingStrategy): + """Splits text on whitespace boundaries. Each word is a chunk. Trailing text not yet followed by whitespace is withheld. """ - def split(self, accumulated_text: str) -> list[str]: - """Return complete words from accumulated_text. + def split(self, text: str) -> list[str]: + """Return complete words from text. Args: - accumulated_text: The full text accumulated so far. + text: The text to split. Returns: All whitespace-delimited words except the trailing fragment (if any). An empty list is returned when no whitespace boundary has been seen. """ - if not accumulated_text: + if not text: return [] # Split on runs of whitespace; the last token is a trailing fragment - # unless accumulated_text ends with whitespace. - parts = _WHITESPACE.split(accumulated_text) + # unless text ends with whitespace. + parts = _WHITESPACE.split(text) # re.split on leading whitespace produces an empty first element; strip it. if parts and parts[0] == "": @@ -201,12 +201,12 @@ def split(self, accumulated_text: str) -> list[str]: return [] # If the text does not end with whitespace, the last part is a fragment. - if not accumulated_text[-1].isspace(): + if not text[-1].isspace(): return parts[:-1] return parts - def flush(self, accumulated_text: str) -> list[str]: + def flush(self, text: str) -> list[str]: """Return the trailing word fragment (if any) as a final chunk. The trailing fragment is the text after the last whitespace run when @@ -215,26 +215,26 @@ def flush(self, accumulated_text: str) -> list[str]: released. Args: - accumulated_text: The full accumulated text at stream end. + text: The text whose trailing fragment to release. Returns: A single-element list containing the trailing word fragment, or an empty list when the input ends with whitespace (every word already complete) or is empty. """ - if not accumulated_text: + if not text: return [] - if accumulated_text[-1].isspace(): + if text[-1].isspace(): return [] - parts = _WHITESPACE.split(accumulated_text) + parts = _WHITESPACE.split(text) for part in reversed(parts): if part: return [part] return [] -class ParagraphChunker(ChunkingStrategy): - r"""Splits accumulated text on double-newline paragraph boundaries. +class ParagraphChunking(ChunkingStrategy): + r"""Splits text on double-newline paragraph boundaries. Two or more consecutive newline characters are treated as a paragraph separator. The trailing paragraph fragment (text not yet followed by `\n\n`) @@ -244,33 +244,33 @@ class ParagraphChunker(ChunkingStrategy): (`\r\n\r\n`) paragraph separators are not supported. """ - def split(self, accumulated_text: str) -> list[str]: - """Return complete paragraphs from accumulated_text. + def split(self, text: str) -> list[str]: + """Return complete paragraphs from text. Args: - accumulated_text: The full text accumulated so far. + text: The text to split. Returns: Complete paragraphs (separated by two or more newlines). The trailing incomplete paragraph is withheld. Returns an empty list if no paragraph boundary has been reached. """ - if not accumulated_text: + if not text: return [] - parts = _PARA_BOUNDARY.split(accumulated_text) + parts = _PARA_BOUNDARY.split(text) # If the text does not end with \n\n, the last part is a trailing fragment. - if not _PARA_BOUNDARY_END.search(accumulated_text): + if not _PARA_BOUNDARY_END.search(text): parts = parts[:-1] # _PARA_BOUNDARY.split on leading \n\n produces an empty first element. return [p for p in parts if p] - def flush(self, accumulated_text: str) -> list[str]: + def flush(self, text: str) -> list[str]: r"""Return the trailing paragraph fragment (if any) as a final chunk. - Unlike :class:`SentenceChunker.flush`, the fragment is returned + Unlike `SentenceChunking.flush`, the fragment is returned byte-for-byte without stripping. Internal whitespace — including a trailing single `\n` — can be semantically meaningful inside a paragraph (e.g. a list item or a deliberate line break), and a @@ -278,17 +278,132 @@ def flush(self, accumulated_text: str) -> list[str]: it was withheld. Args: - accumulated_text: The full accumulated text at stream end. + text: The text whose trailing fragment to release. Returns: A single-element list containing the trailing paragraph fragment byte-for-byte, or an empty list when the input ends with a paragraph boundary (`\n\n` or more) or is empty. """ - if not accumulated_text: + if not text: return [] - if _PARA_BOUNDARY_END.search(accumulated_text): + if _PARA_BOUNDARY_END.search(text): return [] - parts = _PARA_BOUNDARY.split(accumulated_text) + parts = _PARA_BOUNDARY.split(text) trailing = parts[-1] if parts else "" return [trailing] if trailing else [] + + +_ALIASES: dict[str, type[ChunkingStrategy]] = { + "sentence": SentenceChunking, + "word": WordChunking, + "paragraph": ParagraphChunking, +} + + +def resolve_chunking_strategy( + chunking: str | ChunkingStrategy | None, +) -> ChunkingStrategy | None: + """Resolve a chunking argument to a `ChunkingStrategy` instance, or `None`. + + Args: + chunking: A `ChunkingStrategy` (returned as-is), a recognized alias string + (instantiated to its strategy), or `None` (passed through, meaning no + chunking). + + Returns: + The resolved strategy, or `None`. + + Raises: + ValueError: If `chunking` is a string that is not a recognized alias. The + message lists the recognized aliases. + """ + if isinstance(chunking, str): + cls = _ALIASES.get(chunking) + if cls is None: + raise ValueError( + f"Unknown chunking alias {chunking!r}. Choose from: {list(_ALIASES)}" + ) + return cls() + return chunking + + +class Chunker: + """Drives a `ChunkingStrategy` incrementally over a stream of deltas. + + The stateful counterpart to a `ChunkingStrategy`: the strategy is the + stateless "how to split," the `Chunker` holds "how far through this stream we + are." Feed it one delta at a time with `feed()`; it returns any newly complete + chunks and holds the trailing fragment until the next call. Call `flush()` + once at stream end to release the final fragment. + + The `Chunker` holds only the pending fragment (text since the last boundary) + — not the full accumulated text. A caller that needs the full raw stream keeps + its own copy. + + Delta-invariant: feeding text in any delta slicing yields the same chunks as a + single `split()` over the whole text. + + Args: + strategy: The stateless chunking strategy that decides boundaries. + """ + + def __init__(self, strategy: ChunkingStrategy) -> None: + """Wrap `strategy` for incremental driving.""" + self._strategy = strategy + self._pending = "" + + def feed(self, delta: str) -> list[str]: + """Add one stream delta and return any newly complete chunks. + + Args: + delta: The new text received since the previous delta. + + Returns: + The chunks completed by this delta, in order. Empty when the delta + did not complete a boundary. The trailing fragment is withheld until + a later `feed()` completes it or `flush()` releases it. + + Raises: + ValueError: If the strategy's `split()` returns an empty chunk or one + that is not a verbatim substring of the buffered text (i.e. it + mutated the text). + """ + self._pending += delta + chunks = self._strategy.split(self._pending) + if not chunks: + return [] + + # Carry forward whatever follows the last emitted chunk. split() may drop + # inter-chunk whitespace, so locate each chunk by position rather than + # string-subtracting, then keep the raw suffix as the new pending fragment. + cursor = 0 + for c in chunks: + if not c: + raise ValueError( + f"{type(self._strategy).__name__}.split() returned an empty " + "chunk; split() must return only non-empty substrings." + ) + pos = self._pending.find(c, cursor) + if pos < 0: + raise ValueError( + f"{type(self._strategy).__name__}.split() returned a chunk that " + "is not a verbatim substring of the buffered text; split() must " + "not mutate chunk text (see ChunkingStrategy.split)." + ) + cursor = pos + len(c) + self._pending = self._pending[cursor:] + return chunks + + def flush(self) -> list[str]: + """Release the trailing fragment withheld after the last boundary. + + Call once when the stream ends naturally. Returns whatever the strategy's + `flush()` makes of the pending fragment (a single-element list, or empty + when nothing remains). + + Returns: + The final chunk as a one-element list, or an empty list when the + pending fragment is empty or the strategy discards it. + """ + return self._strategy.flush(self._pending) diff --git a/mellea/stdlib/streaming.py b/mellea/stdlib/streaming.py index 0a294a0634..0cb9a3947c 100644 --- a/mellea/stdlib/streaming.py +++ b/mellea/stdlib/streaming.py @@ -1,23 +1,26 @@ # Copyright IBM Corp. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -"""Streaming generation with per-chunk validation. +"""Streaming generation: a single-task `async for` primitive. -Provides :func:`stream_with_chunking`, the core orchestration primitive that -consumes a streaming :class:`~mellea.core.base.ModelOutputThunk`, applies a -:class:`~mellea.stdlib.chunking.ChunkingStrategy` to produce semantic chunks, -and runs :meth:`~mellea.core.requirement.Requirement.stream_validate` on each -chunk in parallel. Higher-level streaming APIs build on this function. +`stream()` starts a streaming generation and returns a `Streamer` you consume with +`async for`. It drives token draining, chunking, and (when requirements are given) +per-chunk and final validation. -The orchestrator emits typed :class:`StreamEvent` objects that consumers can -observe via :meth:`StreamChunkingResult.events`. Raw validated chunks remain -available via :meth:`StreamChunkingResult.astream`. +Consume inside `async with` so cleanup always runs: the stream runs on the caller's +task, and leaving the block — normally, on early `break`, or on exception — +cancels the generation and fires the `STREAMING_END` hook. + +Typed `StreamEvent` objects are emitted via the `STREAMING_EVENT` hook; subscribe a +plugin to observe them (see `docs/examples/streaming/`). """ +from __future__ import annotations + import asyncio import time import uuid -from collections.abc import AsyncIterator, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Sequence from copy import copy from dataclasses import dataclass, field from typing import Any @@ -26,16 +29,9 @@ from ..core.backend import Backend from ..core.base import CBlock, Component, Context, ModelOutputThunk from ..core.requirement import PartialValidationResult, Requirement, ValidationResult -from ..core.utils import MelleaLogger from ..plugins.manager import has_plugins, invoke_hook from ..plugins.types import HookType -from .chunking import ChunkingStrategy, ParagraphChunker, SentenceChunker, WordChunker - -_CHUNKING_ALIASES: dict[str, type[ChunkingStrategy]] = { - "sentence": SentenceChunker, - "word": WordChunker, - "paragraph": ParagraphChunker, -} +from .chunking import Chunker, ChunkingStrategy, resolve_chunking_strategy # --------------------------------------------------------------------------- # Streaming event types @@ -44,7 +40,7 @@ @dataclass class StreamEvent: - """Base class for all streaming events emitted by :func:`stream_with_chunking`. + """Base class for all streaming events emitted by `stream`. The `timestamp` field is auto-populated at instantiation time; callers do not set it. Because `timestamp` has `init=False` it is never part @@ -64,13 +60,12 @@ class ChunkEvent(StreamEvent): """Emitted after each validated chunk is delivered to the consumer. Fired after all active requirements' `stream_validate` calls return - non-`"fail"` for this chunk and the chunk has been placed on the - consumer queue. + non-`"fail"` for this chunk and the chunk has been yielded to the consumer. Args: text: The chunk text that was validated and emitted. chunk_index: Zero-based position of this chunk in the stream. - attempt: Sampling attempt number (always `1` in v1). + attempt: Sampling attempt number; currently always `1`. """ text: str @@ -87,12 +82,11 @@ class QuickCheckEvent(StreamEvent): Args: chunk_index: Zero-based position of the chunk that was validated. - attempt: Sampling attempt number (always `1` in v1). + attempt: Sampling attempt number; currently always `1`. passed: `True` if all active requirements returned non-`"fail"` for this chunk. - results: :class:`~mellea.core.requirement.PartialValidationResult` - from each active requirement, in the same order as the active - slice of `requirements`. + results: `PartialValidationResult` from each active requirement, in the + same order as the active slice of `requirements`. """ chunk_index: int @@ -106,12 +100,12 @@ class StreamingDoneEvent(StreamEvent): """Emitted after all chunks have been validated and delivered to the consumer. Fired after the regular token stream and any trailing fragment released by - :meth:`~mellea.stdlib.chunking.ChunkingStrategy.flush` have both been - processed. Only emitted on natural completion — not on early exit (a - requirement returned `"fail"`) or on exception. + the chunker's `flush()` have both been processed. Only emitted on natural + completion — not on early exit (a requirement returned `"fail"`) or on + exception. Args: - attempt: Sampling attempt number (always `1` in v1). + attempt: Sampling attempt number; currently always `1`. full_text: Complete accumulated text at stream end. """ @@ -121,17 +115,15 @@ class StreamingDoneEvent(StreamEvent): @dataclass class FullValidationEvent(StreamEvent): - """Emitted after the final :meth:`~mellea.core.requirement.Requirement.validate` calls complete. + """Emitted after the final `Requirement.validate` calls complete. - Only emitted when at least one requirement did not fail during streaming - and the stream completed naturally. Not emitted on early exit. + Only emitted when the stream completed naturally (no requirement failed + during streaming). Not emitted on early exit. Args: - attempt: Sampling attempt number (always `1` in v1). - passed: `True` if all final - :class:`~mellea.core.requirement.ValidationResult` objects passed. - results: :class:`~mellea.core.requirement.ValidationResult` from each - non-failed requirement, in requirement order. + attempt: Sampling attempt number; currently always `1`. + passed: `True` if all final `ValidationResult` objects passed. + results: `ValidationResult` from each requirement, in requirement order. """ attempt: int @@ -143,10 +135,9 @@ class FullValidationEvent(StreamEvent): class RetryEvent(StreamEvent): """Reserved for future use. - Defined for API completeness — `RetryEvent` is not emitted by the - v1 orchestrator because v1 retry is caller-driven re-invocation of - :func:`stream_with_chunking`. When orchestrator-side retry is added, - this event will fire before each re-attempt. + Defined for API completeness — `RetryEvent` is not currently emitted; today + retry is caller-driven re-invocation of `stream`. If retry is added to + streaming itself, this event will fire before each re-attempt. Args: attempt: Attempt number being started (1-based). @@ -159,17 +150,17 @@ class RetryEvent(StreamEvent): @dataclass class CompletedEvent(StreamEvent): - """Emitted when the orchestrator exits, including early-exit cases. + """Emitted when the stream exits, including early-exit cases. - Always the last event before :meth:`StreamChunkingResult.events` - terminates. `success` reflects :attr:`StreamChunkingResult.completed`. + Always the last `StreamEvent` on every exit path. `success` reflects + whether the stream completed with no `"fail"` result and no exception. Args: success: `True` if the stream completed normally (no `"fail"` result and no unhandled exception); `False` otherwise. - full_text: Complete accumulated text. On early exit or exception, - reflects whatever was accumulated before cancellation. - attempts_used: Number of orchestrator invocations (always `1` in v1). + full_text: Validated-and-emitted output. On early exit or exception, + reflects whatever passed validation before the stop. + attempts_used: Number of stream attempts; currently always `1`. """ success: bool @@ -179,14 +170,12 @@ class CompletedEvent(StreamEvent): @dataclass class ErrorEvent(StreamEvent): - """Emitted when an unhandled exception occurs in the orchestrator. + """Emitted when an unhandled exception occurs while streaming. Args: exception_type: Python class name of the exception (e.g. `"ValueError"`). - detail: String representation of the exception. If - `cancel_generation()` also raised during cleanup, the cleanup - error is appended. + detail: String representation of the exception. """ exception_type: str @@ -194,536 +183,371 @@ class ErrorEvent(StreamEvent): # --------------------------------------------------------------------------- -# Result container +# Streamer handle # --------------------------------------------------------------------------- -class StreamChunkingResult: - """Result of a :func:`stream_with_chunking` operation. - - Provides async iteration over validated text chunks as they complete - (:meth:`astream`), typed :class:`StreamEvent` objects via :meth:`events`, - a blocking :meth:`acomplete` for awaiting the full result including final - validation, and :attr:`as_thunk` for wrapping the output as a - :class:`~mellea.core.base.ModelOutputThunk`. +class Streamer: + """Async-iterable handle for a `stream` call. - Instances are created by :func:`stream_with_chunking`; do not instantiate - directly. + Iterate the returned `Streamer` object with `async for` to receive the output + as validated chunks, ideally inside `async with` so the stream is released on + every exit. Each chunk is a `str` segment of the model output text, sized by the + `chunking` strategy (or a raw model delta when `chunking` is `None`). The + attributes below track progress and outcome. Instances are created by `stream`; + do not instantiate directly. Args: - mot: The :class:`~mellea.core.base.ModelOutputThunk` from the backend - generation call. - ctx: The generation context returned alongside the MOT. - streaming_id: UUID correlating the streaming hooks for this run. + 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. Attributes: - completed: `False` if the stream exited early because a requirement - returned `"fail"` during streaming; `True` otherwise. - full_text: The generated text available after streaming completes. - On natural completion, the full accumulated text. On early exit - (a requirement returned `"fail"`), only the validated and emitted - portion — i.e. what consumers received via :meth:`astream`. - Available after :meth:`acomplete` returns. - final_validations: :class:`~mellea.core.requirement.ValidationResult` - objects from the final :meth:`~mellea.core.requirement.Requirement.validate` - calls on all non-failed requirements. Available after - :meth:`acomplete` returns. - streaming_failures: `(Requirement, PartialValidationResult)` pairs - for every requirement that returned `"fail"` during streaming. + failed_early: `True` if a requirement returned `"fail"` during streaming + and the stream stopped before natural completion. + completed_normally: `True` only if the stream reached its natural end, + prior to final validation. `False` on requirement failure, an early + `break`, or an exception — unlike `not failed_early`, which stays + `True` after an early `break`. + failure_reason: Human-readable reason when `failed_early` is `True`. + streaming_failures: `(Requirement, PartialValidationResult)` pairs for + every requirement that failed the offending chunk. + full_text: Validated-and-emitted output. On natural completion, the full + accumulated text; on early exit, the accumulated text through the last + emitted chunk. + mot: The computed thunk, set on natural completion; `None` otherwise. + final_validations: `ValidationResult` objects from the stream-end + `validate()` calls; empty on early exit. + streaming_id: UUID correlating this stream's START/EVENT/END hooks. """ - def __init__(self, mot: ModelOutputThunk, ctx: Context, streaming_id: str) -> None: - """Initialise with the MOT and context from the backend call.""" - self._mot = mot - self._ctx = ctx - self._streaming_id = streaming_id - self._chunk_queue: asyncio.Queue[str | None | Exception] = asyncio.Queue() - # If no consumer calls events(), events accumulate in this queue until - # the result object is garbage-collected. That is intentional — event - # production is unconditional; consumption is opt-in. - self._event_queue: asyncio.Queue[StreamEvent | None] = asyncio.Queue() - self._orchestration_task: asyncio.Task[None] | None = None - self._done = asyncio.Event() - # Set synchronously at the very top of _orchestrate_streaming, before - # any await, so external coordinators (e.g. cancellation tests) can wait - # until the task is live and suspended at its first I/O point. - # Single-use: not reset between runs; this object is not re-entrant. - self._orchestration_started = asyncio.Event() - # Stashed so acomplete() surfaces orchestrator failures even when the - # consumer never iterates astream(). Cleared once consumed by - # whichever of the two reads it first. - self._orchestration_exception: BaseException | None = None - # Tracks whether the exception has already been surfaced to the caller - # (by astream OR acomplete). A separate flag rather than reusing the - # stash slot avoids the race where acomplete() clears the stash, a - # subsequent astream() dequeues the exception item, sees the stash is - # None, and silently skips it — leaving the caller with zero chunks - # and no error. - self._exception_surfaced: bool = False - self._events_consumed: bool = False - # Outcome for the streaming_end hook. Recorded by the orchestrator; - # kept separate from `_orchestration_exception`, which raise-once clears. - self._streaming_failure_reason: str | None = None - self._streaming_end_exception: Exception | None = None - # streaming_end fires once even if acomplete() is called repeatedly. - self._streaming_end_fired: bool = False - - self.completed: bool = True + def __init__( + self, + mot: ModelOutputThunk, + ctx: Context, + chunking: ChunkingStrategy | None, + requirements: list[Requirement], + validation_backend: Backend, + streaming_id: str, + ) -> None: + """Wrap an in-flight generation; iterating the `Streamer` drives it.""" + self.failed_early: bool = False + self.completed_normally: 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] = [] - self.streaming_failures: list[tuple[Requirement, PartialValidationResult]] = [] + # Correlates this stream's START/EVENT/END hooks; created in `stream()` + # so START can fire before generation opens the backend span. + self.streaming_id: str = streaming_id + # The in-flight thunk, for teardown. Held separately from the public `mot`, + # which is only set once the stream completes. + self._mot = mot + self._finalized: bool = False + self._gen: AsyncGenerator[str, None] = _drive( + self, mot, ctx, chunking, requirements, validation_backend + ) - async def astream(self) -> AsyncIterator[str]: - """Yield validated text chunks as they complete. - - Each yielded string is a chunk that has passed per-chunk streaming - validation (or the stream had no requirements). Iteration ends when - all chunks have been yielded, whether the stream completed normally or - was cancelled early on a `"fail"` result. - - **Single-consumer.** Chunks are delivered via an - :class:`asyncio.Queue` that this method drains; calling - `astream()` a second time on the same result blocks indefinitely - because the queue is empty and the terminating `None` sentinel - has already been consumed. If you need the chunks after - iteration, capture them into a list during the first pass or use - :attr:`full_text` after :meth:`acomplete`. - - Yields: - str: A validated text chunk from the chunking strategy. - - Raises: - Exception: Propagates any error from the background orchestration - task. - - Note: - Draining `astream()` does not finalize the run; `full_text`, - `final_validations`, `completed`, and the `STREAMING_END` hook are - all driven by `acomplete()`. Call `acomplete()` after iterating. - """ - while True: - item = await self._chunk_queue.get() - if item is None: - return - if isinstance(item, Exception): - if self._exception_surfaced: - # Already surfaced by acomplete(); don't raise twice. - continue - self._exception_surfaced = True - self._orchestration_exception = None - raise item - yield item - - async def events(self) -> AsyncIterator[StreamEvent]: - """Yield typed streaming events as they are emitted by the orchestrator. - - Each yielded object is a :class:`StreamEvent` subclass describing a - point in the orchestration lifecycle. Consumers can dispatch on type: - - ```python - async for event in result.events(): - match event: - case ChunkEvent(): - print(f"chunk {event.chunk_index}: {event.text!r}") - case QuickCheckEvent(passed=False): - print(f"chunk {event.chunk_index} failed validation") - case CompletedEvent(): - print(f"done — success={event.success}") - ``` - - Typical event order (natural completion with requirements): - - 1. :class:`QuickCheckEvent` / :class:`ChunkEvent` pairs, one per chunk - (validation fires first; the chunk is released to the consumer only - after passing). Includes any trailing fragment released by the - chunking strategy's `flush()` method. - 2. :class:`StreamingDoneEvent` — all chunks (including flush) delivered. - 3. :class:`FullValidationEvent` — final `validate()` calls returned. - 4. :class:`CompletedEvent` — orchestrator is exiting. - - On early exit: :class:`QuickCheckEvent` (`passed=False`) is the - last validation event, followed by :class:`CompletedEvent`. No - :class:`StreamingDoneEvent` or :class:`FullValidationEvent` is emitted. - - On exception: :class:`ErrorEvent` followed by :class:`CompletedEvent`. - - **Single-consumer.** Events are delivered via a queue that this method - drains; calling `events()` a second time raises :exc:`RuntimeError`. - - Yields: - StreamEvent: A typed event from the orchestrator. - - Raises: - RuntimeError: If called more than once on the same result. - - Note: - `events()` itself never raises from the event stream. If the - orchestrator encounters an unhandled exception, an - :class:`ErrorEvent` is emitted and iteration ends normally. - Exceptions surface to the caller via :meth:`astream` (as a - re-raised exception) or :meth:`acomplete`. - """ - if self._events_consumed: - raise RuntimeError( - "events() is single-consumer; this iterator has already been drained" - ) - self._events_consumed = True - while True: - item = await self._event_queue.get() - if item is None: - return - yield item - - async def acomplete(self) -> None: - """Await full completion, including final validation. - - After this method returns, :attr:`full_text`, :attr:`completed`, - :attr:`final_validations`, and :attr:`streaming_failures` are all - populated. If :meth:`astream` has already been consumed to - exhaustion, this call is effectively a no-op. - - Raises: - Exception: Propagates the orchestrator exception if :meth:`astream` - has not yet consumed it (raise-once — only one of `astream` - or `acomplete` raises, whichever drains the failure marker - first). - asyncio.CancelledError: If the orchestration task was externally - cancelled (e.g. via :func:`asyncio.wait_for` timeout). + def __aiter__(self) -> AsyncIterator[str]: + """Return the generator that drives generation and yields chunks.""" + return self._gen + + async def _finalize( + self, + *, + success: bool = False, + error: Exception | None = None, + full_text_length: int = 0, + ) -> None: + """Cancel the generation and fire the terminal events, at most once. + + Idempotent: the `_finalized` guard makes every call after the first a + no-op, so callers need not coordinate. It is invoked both from the + driver's `finally` and from `aclose()` — either may run first, or only one + may run at all (e.g. a `Streamer` closed without being iterated) — and the + terminal events still fire exactly once. """ - await self._done.wait() - - if not self._streaming_end_fired: - self._streaming_end_fired = True - if has_plugins(HookType.STREAMING_END): - from ..plugins.hooks.streaming import StreamingEndPayload - - await invoke_hook( - HookType.STREAMING_END, - StreamingEndPayload( - streaming_id=self._streaming_id, - success=self.completed, - failure_reason=self._streaming_failure_reason, - exception=self._streaming_end_exception, - model=self._mot.generation.model, - provider=self._mot.generation.provider, - full_text_length=len(self.full_text), + if self._finalized: + return + # Set before teardown: a cancellation mid-dispatch can curtail the remaining + # STREAMING_END subscribers with no retry (unclosed span / unrecorded metrics). + self._finalized = True + + try: + # aclose() is a no-op once the stream is fully drained; cancels otherwise. + await self._mot.aclose() + finally: + try: + await _emit_event( + self.streaming_id, + CompletedEvent( + success=success, full_text=self.full_text, attempts_used=1 ), ) + finally: + if has_plugins(HookType.STREAMING_END): + from ..plugins.hooks.streaming import StreamingEndPayload + + await invoke_hook( + HookType.STREAMING_END, + StreamingEndPayload( + streaming_id=self.streaming_id, + success=success, + failure_reason=self.failure_reason, + exception=error, + model=self._mot.generation.model, + provider=self._mot.generation.provider, + full_text_length=full_text_length, + ), + ) + + async def aclose(self) -> None: + """Release the stream, cancelling generation if it is still in flight. + + Runs the driver's cleanup (cancelling the backend generation and firing + `STREAMING_END`). Safe and idempotent on every path: after natural + completion, after an early exit/break, and on a `Streamer` that was never + iterated — the eager generation is still cancelled in every case. - # Raise-once: if astream() already surfaced the exception, skip. - exc = self._orchestration_exception - if exc is not None and not self._exception_surfaced: - self._exception_surfaced = True - self._orchestration_exception = None - raise exc - if self._orchestration_task is not None and self._orchestration_task.done(): - # Raise-once: a prior call already surfaced the exception. - if self._exception_surfaced: - return - # `task.exception()` raises CancelledError on a cancelled task - # (rather than returning it), so check cancelled status first. - # This branch covers BaseException paths that bypass the - # `except Exception` handler in `_orchestrate_streaming`. - if self._orchestration_task.cancelled(): - self._exception_surfaced = True - raise asyncio.CancelledError() - task_exc = self._orchestration_task.exception() - if task_exc is not None: - self._exception_surfaced = True - raise task_exc - - @property - def as_thunk(self) -> ModelOutputThunk[str]: - """Wrap the output as a computed :class:`~mellea.core.base.ModelOutputThunk`. - - Returns a new thunk with `value` set to :attr:`full_text` and - generation metadata copied from the original MOT. Safe to call on - early-exit results; `value` reflects the validated and emitted - portion (same as :attr:`full_text` — see its docstring). - - Note: - On early exit, `cancel_generation()` forces the MOT into a - computed state without running the backend's - `post_processing()`. `value` and `streaming` are - reliable. `parsed_repr` is set to the raw text (same as - `value`) — consistent with normal completion for plain-text - outputs, but for typed outputs the backend-parsed representation - will not be available. Telemetry fields (`generation.usage`, - `generation.ttfb_ms`, etc.) may be `None` or reflect the - partial state at cancellation time; usage totals are not - recoverable. - - Returns: - ModelOutputThunk[str]: A computed thunk containing the streamed output. - - Raises: - RuntimeError: If called before :meth:`acomplete` has returned. + Prefer consuming with `async with stream(...) as s:` so this runs + automatically on every exit path; call `aclose()` directly only when not + using the context manager. """ - if not self._done.is_set(): - raise RuntimeError( - "as_thunk accessed before acomplete() — await acomplete() first" - ) - thunk = ModelOutputThunk(value=self.full_text) - thunk._cancelled = self._mot.cancelled - thunk.generation = copy(self._mot.generation) - thunk.parsed_repr = thunk.value # type: ignore[assignment] - return thunk + # Closing the generator finalizes via its finally if iteration started; + # the explicit call handles the never-iterated case. + await self._gen.aclose() + await self._finalize() + + async def __aenter__(self) -> Streamer: + """Enter the async context manager, returning this `Streamer`.""" + return self + + async def __aexit__(self, *exc_info: object) -> None: + """Exit the context manager, releasing the stream via `aclose()`.""" + await self.aclose() # --------------------------------------------------------------------------- -# Orchestrator +# Driver # --------------------------------------------------------------------------- -async def _orchestrate_streaming( - result: StreamChunkingResult, - mot: ModelOutputThunk, - ctx: Context, - cloned_reqs: list[Requirement], - chunking: ChunkingStrategy, - val_backend: Backend, +async def _emit_event( + streaming_id: str, ev: StreamEvent, *, requirements: list[Requirement] | None = None ) -> None: - # Signal that the coroutine body is executing before the first suspension. - # External coordinators waiting on _orchestration_started are guaranteed to - # resume only after this task has yielded at its first real await, so a - # subsequent cancel() always lands on a live, non-done task. - result._orchestration_started.set() + """Fire the STREAMING_EVENT hook for `ev`. - accumulated = "" - emitted_end = 0 # byte offset in accumulated after the last emitted chunk - prev_chunk_count = 0 - failed_indices: set[int] = set() - early_exit = False - chunk_index = 0 + For a `QuickCheckEvent`, `requirements` carries the active requirement + instances in result order so a subscriber can attribute each result. - async def _emit_event( - ev: StreamEvent, *, requirements: list[Requirement] | None = None - ) -> None: - """Push *ev* on the consumer queue, then fire its streaming_event hook. + Args: + streaming_id: UUID correlating this stream's events. + ev: The event to emit. + requirements: Active requirements for a `QuickCheckEvent`, in result + order; `None` for other event types. + """ + if has_plugins(HookType.STREAMING_EVENT): + from ..plugins.hooks.streaming import StreamingEventPayload - For a `QuickCheckEvent`, *requirements* carries the active requirement - instances in result order, so a subscriber can attribute each result to - its requirement type. - """ - await result._event_queue.put(ev) - if has_plugins(HookType.STREAMING_EVENT): - from ..plugins.hooks.streaming import StreamingEventPayload + await invoke_hook( + HookType.STREAMING_EVENT, + StreamingEventPayload( + streaming_id=streaming_id, event=ev, requirements=requirements or [] + ), + ) - await invoke_hook( - HookType.STREAMING_EVENT, - StreamingEventPayload( - streaming_id=result._streaming_id, - event=ev, - requirements=requirements or [], - ), - ) - async def _process_chunk(c: str, ci: int) -> bool: - """Validate *c*, emit events, push to consumer queue. +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` if a `"fail"` was recorded (caller should - trigger early exit), `False` if the chunk was validated and - emitted successfully. - """ - active = [ - (i, req) for i, req in enumerate(cloned_reqs) if i not in failed_indices - ] - pvrs: list[PartialValidationResult] = [] - if active: - pvrs = list( - await asyncio.gather( - *[ - req.stream_validate(c, backend=val_backend, ctx=ctx) - for _, req in active - ] - ) - ) - for (idx, req), pvr in zip(active, pvrs): - if pvr.success == "fail": - failed_indices.add(idx) - result.streaming_failures.append((req, pvr)) + 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. - any_fail = any(pvr.success == "fail" for pvr in pvrs) - await _emit_event( - QuickCheckEvent( - chunk_index=ci, attempt=1, passed=not any_fail, results=pvrs - ), - requirements=[req for _, req in active], - ) + Args: + streamer: The handle recording failures for the caller. + chunk: The chunk text to validate. + chunk_index: Zero-based position of this chunk in the stream. + requirements: Requirements to validate against. + validation_backend: Backend used for validation calls. + ctx: The generation context. + on_flush: `True` when validating the trailing flushed fragment. - if failed_indices: - return True + Returns: + `True` if the chunk passed and may be emitted; `False` if it failed. + """ + 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, +) -> AsyncGenerator[str, None]: + """Drive the whole stream from one generator on the caller's task. - await result._chunk_queue.put(c) - await _emit_event(ChunkEvent(text=c, chunk_index=ci, attempt=1)) - return False + 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). - try: - # Inside the try so a cancellation at this await still runs the finally. - if has_plugins(HookType.STREAMING_ORCHESTRATION_START): - from ..plugins.hooks.streaming import StreamingOrchestrationStartPayload + 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"`. - await invoke_hook( - HookType.STREAMING_ORCHESTRATION_START, - StreamingOrchestrationStartPayload(streaming_id=result._streaming_id), - ) + Args: + streamer: The handle recording terminal state for the caller. + mot: The in-flight streaming thunk. + ctx: The generation context. + chunking: Resolved chunking strategy, or `None` for raw deltas. + requirements: Requirements to validate against. + validation_backend: Backend used for validation calls. - while not mot.is_computed(): - try: - delta = await mot.astream() - except RuntimeError: - # Expected race: mot.is_computed() was False at the top of the - # loop but the stream finished before we re-entered astream(). - # Any other RuntimeError is a real bug and must propagate. - if mot.is_computed(): - break - raise + Yields: + str: Each validated chunk, in order. + """ + # `accumulated` is the full raw text across deltas; the Chunker holds only the + # pending fragment. chunking=None means yield raw deltas, no Chunker. + accumulated = "" + chunk_index = 0 + success = False + error: Exception | None = None + chunker = Chunker(chunking) if chunking is not None else None + emitted_end = 0 # offset in `accumulated` just past the last emitted chunk + + def _snapshot_full_text(chunk: str) -> None: + nonlocal emitted_end + pos = accumulated.find(chunk, emitted_end) + if pos >= 0: + emitted_end = pos + len(chunk) + streamer.full_text = accumulated[:emitted_end] + try: + async for delta in mot: accumulated += delta - chunks = chunking.split(accumulated) - new_chunks = chunks[prev_chunk_count:] - prev_chunk_count = len(chunks) + + if chunker is None: + new_chunks = [delta] if delta else [] # raw mode: delta is the chunk + else: + new_chunks = chunker.feed(delta) for c in new_chunks: - failed = await _process_chunk(c, chunk_index) - if failed: - early_exit = True - result.completed = False - await mot.cancel_generation() - reason = result.streaming_failures[-1][1].reason or "" - result._streaming_failure_reason = ( - f"Streaming validation failed: {reason}" - ) - break - pos = accumulated.find(c, emitted_end) - if pos >= 0: - emitted_end = pos + len(c) + if not await _validate_chunk( + streamer, c, chunk_index, requirements, validation_backend, ctx + ): + return + _snapshot_full_text(c) # record before yield; a break skips past it + await _emit_event( + streamer.streaming_id, + ChunkEvent(text=c, chunk_index=chunk_index, attempt=1), + ) + yield c chunk_index += 1 - if early_exit: - break - - # Stream ended naturally: flush any withheld trailing fragment, then - # emit StreamingDoneEvent once all chunks (regular + flush) have been - # validated and delivered. If a flush chunk fails, early_exit is - # set and StreamingDoneEvent is suppressed (same contract as the - # regular early-exit path). Skipped entirely on early exit. - if not early_exit: - for c in chunking.flush(accumulated): - failed = await _process_chunk(c, chunk_index) - if failed: - early_exit = True - result.completed = False - reason = result.streaming_failures[-1][1].reason or "" - result._streaming_failure_reason = ( - f"Streaming validation failed on flush: {reason}" - ) - break - pos = accumulated.find(c, emitted_end) - if pos >= 0: - emitted_end = pos + len(c) + # Flush the trailing fragment the chunker withheld (skipped in raw mode). + if chunker is not None: + for c in chunker.flush(): + if not await _validate_chunk( + streamer, + c, + chunk_index, + requirements, + validation_backend, + ctx, + on_flush=True, + ): + return + _snapshot_full_text(c) + await _emit_event( + streamer.streaming_id, + ChunkEvent(text=c, chunk_index=chunk_index, attempt=1), + ) + yield c chunk_index += 1 - if not early_exit: - await _emit_event(StreamingDoneEvent(attempt=1, full_text=accumulated)) - - # On early exit, full_text is the portion of accumulated that was - # actually validated and emitted to the consumer. On natural - # completion, the full accumulated text is used. - result.full_text = accumulated[:emitted_end] if early_exit else accumulated + streamer.full_text = accumulated + streamer.mot = mot + streamer.completed_normally = True + await _emit_event( + streamer.streaming_id, StreamingDoneEvent(attempt=1, full_text=accumulated) + ) - if not early_exit: - non_failed = [ - req for i, req in enumerate(cloned_reqs) if i not in failed_indices - ] - if non_failed: - vrs: list[ValidationResult] = list( - await asyncio.gather( - *[req.validate(val_backend, ctx) for req in non_failed] - ) - ) - result.final_validations = vrs - all_passed = all(vr.as_bool() for vr in vrs) - await _emit_event( - FullValidationEvent(attempt=1, passed=all_passed, results=list(vrs)) + # 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: - # Stash the exception before any await so acomplete() can always - # surface it even if a subsequent await is interrupted by an - # external CancelledError. - result._orchestration_exception = exc - result._streaming_end_exception = exc - # Mark as failed immediately — before any event is enqueued — so - # that CompletedEvent.success and result.completed are consistent - # if the consumer observes them during ErrorEvent processing. - result.completed = False - result.full_text = accumulated # best-effort partial capture - # Only cancel generation if the stream hasn't already completed - # (e.g. an exception from the final validate() call arrives after - # the token stream ended naturally — cancelling an already-computed - # MOT is a no-op at best and misleading in telemetry). - if not mot.is_computed(): - try: - await mot.cancel_generation(error=exc) - error_detail = str(exc) - except Exception as cleanup_exc: - # Never let cleanup mask the original exception. - error_detail = f"{exc!r} (cancel cleanup raised: {cleanup_exc!r})" - MelleaLogger.get_logger().debug( - "stream_with_chunking: cancel_generation() raised during " - "exception cleanup (original: %r, cleanup: %r)", - exc, - cleanup_exc, - ) - else: - error_detail = str(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( - ErrorEvent(exception_type=type(exc).__name__, detail=error_detail) + streamer.streaming_id, + ErrorEvent(exception_type=type(exc).__name__, detail=str(exc)), ) - await result._chunk_queue.put(exc) + raise finally: - # CancelledError (BaseException, not Exception) bypasses the except - # block above, so cancel_generation() may not have been called. - # Guard here ensures the backend producer is always stopped, even on - # external task cancellation (e.g. asyncio.wait_for timeout). - # Also mark completion as failed for any BaseException path (e.g. - # CancelledError) that bypassed the except block — otherwise - # result.completed stays True and CompletedEvent / metrics lie. - if not mot.is_computed(): - result.completed = False - try: - await mot.cancel_generation() - except BaseException: - pass - - # Shielded so a CancelledError from the hook cannot skip the terminal - # queue bookkeeping and _done.set() below. - if has_plugins(HookType.STREAMING_ORCHESTRATION_END): - from ..plugins.hooks.streaming import StreamingOrchestrationEndPayload - - try: - await invoke_hook( - HookType.STREAMING_ORCHESTRATION_END, - StreamingOrchestrationEndPayload(streaming_id=result._streaming_id), - ) - except BaseException: - pass - # STREAMING_END is not fired here: it pairs with STREAMING_START - # (caller's task, pre-spawn) and a subscriber may hold task-affine - # state, so teardown must run on the caller's task — acomplete(). - - completed_ev = CompletedEvent( - success=result.completed, full_text=result.full_text, attempts_used=1 + # Driver-side teardown on every exit path + await streamer._finalize( + success=success, error=error, full_text_length=len(streamer.full_text) ) - # Use put_nowait for the terminal bookkeeping: both queues are - # unbounded so this can never raise QueueFull, and it eliminates - # the await points that could be interrupted by a pending - # CancelledError before _done.set() runs. - result._event_queue.put_nowait(completed_ev) - result._chunk_queue.put_nowait(None) - result._event_queue.put_nowait(None) - result._done.set() # --------------------------------------------------------------------------- @@ -731,129 +555,59 @@ async def _process_chunk(c: str, ci: int) -> bool: # --------------------------------------------------------------------------- -async def stream_with_chunking( +async def stream( action: Component[Any] | CBlock, backend: Backend, ctx: Context, *, + chunking: str | ChunkingStrategy | None = None, requirements: Sequence[Requirement] | None = None, - chunking: str | ChunkingStrategy = "sentence", validation_backend: Backend | None = None, -) -> StreamChunkingResult: - """Generate a streaming response with per-chunk validation. - - Starts a backend generation with streaming enabled, consumes the - :class:`~mellea.core.base.ModelOutputThunk`'s async stream in a single - background task, splits the accumulated text using *chunking*, and runs - :meth:`~mellea.core.requirement.Requirement.stream_validate` on each new - chunk in parallel across all requirements. - - For each new complete chunk produced by the chunking strategy, - `stream_validate` is called once per active requirement (in parallel - via :func:`asyncio.gather`), receiving that single chunk. Multiple - chunks produced from one `astream()` iteration are validated - sequentially in order, so early exit on a `"fail"` result prevents - later chunks in the same batch from being validated or emitted to the - consumer. - - If any requirement returns `"fail"`, the generation is cancelled - immediately (via - :meth:`~mellea.core.base.ModelOutputThunk.cancel_generation`) and - :attr:`StreamChunkingResult.completed` is set to `False`. The - failing chunk is not emitted to the consumer; use - :attr:`StreamChunkingResult.streaming_failures` to inspect what failed. - - When the stream ends naturally, any trailing fragment withheld by the - chunking strategy (see :meth:`~mellea.stdlib.chunking.ChunkingStrategy.flush`) - is released as a final chunk and run through `stream_validate` on the - same terms as the regular chunks. On early exit, the trailing fragment - is discarded because the generation was cancelled mid-token. - - After the stream ends naturally, `validate()` is called on every - requirement that did not return `"fail"` — both `"pass"` and - `"unknown"` trigger final validation. On early exit, no `validate()` - call is made; :attr:`StreamChunkingResult.final_validations` remains - empty. Requirements are cloned (`copy(req)`) before backend generation - begins, so the originals are never mutated and a raising `__copy__` - cannot leak an in-flight backend task. - - The orchestrator emits typed :class:`StreamEvent` objects throughout - execution. Consume them via :meth:`StreamChunkingResult.events` in - parallel with or instead of :meth:`StreamChunkingResult.astream`. - - Requirements that need context beyond the current chunk should - accumulate it themselves across `stream_validate` calls (e.g. - `self._seen = self._seen + chunk`). They must not read `mot.astream()` - directly — this orchestrator is the single consumer of the MOT stream. - - Note: - Chunks are emitted to the consumer (via - :meth:`StreamChunkingResult.astream`) only after every requirement's - `stream_validate` has returned for that chunk. A slow validator - (for example, one that invokes an LLM) therefore adds latency to - every chunk — the consumer sees a chunk at most as quickly as the - slowest active validator. This trade is deliberate in v1: it - preserves the invariant that the consumer never sees content that - has not been validated, which matters for UIs displaying generated - text live. A future fast-path mode that emits chunks to the - consumer concurrently with validation (at the cost of that - invariant) may be added if a concrete use case calls for it. - - Note: - v1 retry is simple re-invocation of this function. Plugin hooks - (`SAMPLING_LOOP_START`, `SAMPLING_REPAIR`, etc.) do not fire - during streaming — use :meth:`StreamChunkingResult.events` for - observability instead. +) -> Streamer: + """Start a streaming generation. + + Generation begins eagerly, before this call returns. Consume the returned + `Streamer` inside `async with` so the stream is always released. On early + exit or `break`, `async with` cancels the in-flight generation: + + ```python + async with await stream(action, backend, ctx) as s: + async for chunk in s: + ... + ``` + + Each iteration yields a chunk — a unit produced by the `chunking` strategy, + or the raw model delta when `chunking` is `None`. A chunk is delivered 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: The backend used for generation and final validation. + backend: Backend used for generation and, unless `validation_backend` + is set, validation. ctx: The generation context. - requirements: Sequence of requirements to validate against each chunk - during streaming. `None` disables streaming validation (chunks - are still produced; `validate()` is not called at stream end). - chunking: Chunking strategy — either a :class:`~mellea.stdlib.chunking.ChunkingStrategy` - instance or one of the string aliases `"sentence"` (default), - `"word"`, or `"paragraph"`. - validation_backend: Optional alternate backend for both - `stream_validate` and final `validate` calls. When `None`, - *backend* is used for validation. + chunking: A `ChunkingStrategy`, a recognized alias string, or `None` + (default) 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: - StreamChunkingResult: A result object providing :meth:`~StreamChunkingResult.astream` - for incremental chunk consumption, :meth:`~StreamChunkingResult.events` for - typed streaming events, and - :meth:`~StreamChunkingResult.acomplete` for blocking until done. + Streamer: An async-iterable handle over the validated chunks. Raises: - ValueError: If *chunking* is a string that does not match any known - alias (`"sentence"`, `"word"`, `"paragraph"`). - RuntimeError: If the backend returns an already-computed - :class:`~mellea.core.base.ModelOutputThunk` instead of a streaming - one. This indicates the backend is not honouring - `ModelOption.STREAM`. - - Note: - Any exception raised by `copy(req)` on a `requirements` entry - propagates to the caller; no backend generation is started in that - case. See :class:`~mellea.core.Requirement` for the `__copy__` - override contract. + 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 — i.e. 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}. Choose from: {list(_CHUNKING_ALIASES)}" - ) - chunking = cls() - - opts: dict[str, Any] = {ModelOption.STREAM: True} + strategy = resolve_chunking_strategy(chunking) - # Clone requirements before starting backend generation so that a raising - # __copy__ (an advertised extension point on Requirement) cannot leave the - # backend feeder task wedged against a full streaming queue with no consumer. + # 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 [])] - val_backend = validation_backend if validation_backend is not None else backend + resolved_backend = validation_backend if validation_backend is not None else backend streaming_id = str(uuid.uuid4()) if has_plugins(HookType.STREAMING_START): @@ -865,40 +619,21 @@ async def stream_with_chunking( streaming_id=streaming_id, has_requirements=bool(cloned_reqs), requirement_count=len(cloned_reqs), - chunking_strategy=type(chunking).__name__, + chunking_strategy=type(strategy).__name__ if strategy else "none", ), ) - mot: ModelOutputThunk | None = None + mot = None try: mot, gen_ctx = await backend.generate_from_context( - action, ctx, model_options=opts + action, ctx, model_options={ModelOption.STREAM: True} ) if mot.is_computed(): raise RuntimeError( - "stream_with_chunking() requires a streaming backend; the backend " - "returned an already-computed MOT. Ensure the backend honours " - "ModelOption.STREAM." + "stream() requires a streaming backend; the backend returned an " + "already-computed MOT. Ensure the backend honours ModelOption.STREAM." ) - result = StreamChunkingResult(mot, gen_ctx, streaming_id) - coro = _orchestrate_streaming( - result, mot, gen_ctx, cloned_reqs, chunking, val_backend - ) - try: - result._orchestration_task = asyncio.create_task(coro) - except BaseException: - coro.close() # prevent "coroutine was never awaited" RuntimeWarning - raise except BaseException as exc: - if mot is not None: - try: - await mot.cancel_generation() - except Exception as cleanup_exc: - MelleaLogger.get_logger().warning( - "stream_with_chunking: cancel_generation() raised during " - "setup-path cleanup (cleanup: %r)", - cleanup_exc, - ) if has_plugins(HookType.STREAMING_END): from ..plugins.hooks.streaming import StreamingEndPayload @@ -914,4 +649,4 @@ async def stream_with_chunking( ) raise - return result + return Streamer(mot, gen_ctx, strategy, cloned_reqs, resolved_backend, streaming_id) diff --git a/mellea/telemetry/metrics_plugins.py b/mellea/telemetry/metrics_plugins.py index 70f59e2f89..bca95ec6e0 100644 --- a/mellea/telemetry/metrics_plugins.py +++ b/mellea/telemetry/metrics_plugins.py @@ -297,7 +297,7 @@ async def record_batch_error_metrics( async def record_streaming_error_metrics( self, payload: StreamingEndPayload, context: dict[str, Any] ) -> None: - """Record error metrics when `stream_with_chunking` ends with an exception. + """Record error metrics when `stream` ends with an exception. Args: payload: Contains the exception plus the model and provider from @@ -443,15 +443,15 @@ async def record_sampling_outcome( async def record_streaming_outcome( self, payload: StreamingEndPayload, context: dict[str, Any] ) -> None: - """Record the `stream_with_chunking` outcome when the orchestrator finishes. + """Record the `stream` outcome when the stream finishes. Args: - payload: Contains the orchestrator's success flag. + payload: Contains the stream's success flag. context: Plugin context (unused). """ from mellea.telemetry.metrics import record_sampling_outcome - record_sampling_outcome("stream_with_chunking", payload.success) + record_sampling_outcome("stream", payload.success) class RequirementMetricsPlugin(Plugin, name="requirement_metrics", priority=1055): diff --git a/mellea/telemetry/tracing.py b/mellea/telemetry/tracing.py index 0d3ba3ee29..e359d07767 100644 --- a/mellea/telemetry/tracing.py +++ b/mellea/telemetry/tracing.py @@ -33,7 +33,6 @@ from __future__ import annotations -import asyncio import os import warnings from importlib.metadata import version @@ -242,13 +241,7 @@ def get_backend_tracer() -> Any: return _backend_tracer -_in_flight_spans: dict[ - str, tuple[Span, Token[Context] | None, asyncio.Task[Any] | None] -] = {} - -# reattach_span() entries, keyed by correlation key: the OTel context token plus -# the task that attached it. Released by release_reattached_span() on that task. -_reattached_tokens: dict[str, tuple[Token[Context], asyncio.Task[Any] | None]] = {} +_in_flight_spans: dict[str, tuple[Span, Token[Context] | None]] = {} def _attach_span_context(span: Span, *, attach: bool) -> Token[Context] | None: @@ -270,72 +263,15 @@ def _attach_span_context(span: Span, *, attach: bool) -> Token[Context] | None: return otel_context.attach(trace.set_span_in_context(span)) -def _current_task() -> asyncio.Task[Any] | None: - """Return the running asyncio task, or None when no loop is running.""" - try: - return asyncio.current_task() - except RuntimeError: - return None - - -def _safe_detach( - token: Token[Context] | None, attach_task: asyncio.Task[Any] | None -) -> None: - """Detach `token`, suppressing only the cross-task detach we expect and understand. - - OTel context tokens are bound to the `contextvars.Context` of the task that - created them, so detaching from a different task fails — OTel catches the - `ValueError` and logs it at ERROR as "Failed to detach context". Most spans - attach and finish on one task and never hit this. - - A cross-task detach is suppressed (skipped, since it would only fail) and - logged at debug only when the detaching task holds an open `reattach_span` - scope — the marker that this task knowingly opened a span elsewhere and - expects the mismatch. Any other cross-task detach is left to run so OTel - surfaces its ERROR with a traceback to the real origin; a warning is added - first to name the task mismatch OTel's message omits. - - Example: - Under `stream_with_chunking` the backend `chat` span attaches in the - caller task but finishes in the orchestration task that drains the MOT. - To keep that span's `chat` children nesting correctly, the orchestration - task re-attaches the streaming span for the duration of the drain - (`reattach_span` / `release_reattached_span`). The cross-task `chat` - detach that then happens within that scope is the expected, suppressed - case. - - Note: - The reattach scope is an *incomplete* proxy for "expected". It holds for - streaming because that case both needs sibling-nesting protection (so it - reattaches) and has an expected cross-task detach. A future case with the - same open-in-parent / close-in-child shape but no siblings to protect - would not reattach, so its equally-expected cross-task detach falls - through to the warn-and-detach path. That is harmless (the detach only - fails, and the task ends right after, so nothing leaks); the warning is - the signal that the new use needs its own way to mark the detach expected. +def _detach_token(token: Token[Context] | None) -> None: + """Detach an OTel context token, or no-op when `token` is `None`. Args: - token: The OTel context token returned by the matching `attach`, or - `None` when attach was skipped — a no-op. - attach_task: The task that performed the `attach`, or None if it was - attached outside any running task. + token: The token returned by the matching `attach`, or `None` when + attach was skipped. """ if token is None: return - current = _current_task() - if attach_task is not None and current is not attach_task: - from mellea.core.utils import MelleaLogger - - if any(task is current for _, task in _reattached_tokens.values()): - MelleaLogger.get_logger().debug( - "Skipped expected cross-task OTel context detach within a " - "reattached-span scope." - ) - return - MelleaLogger.get_logger().warning( - "Detaching an OTel context token across asyncio tasks; the span's " - "attach and detach ran on different tasks. OTel will log the failure." - ) otel_context.detach(token) @@ -411,7 +347,7 @@ def start_backend_span( set_conversation_id(span) token = _attach_span_context(span, attach=attach_context) - _in_flight_spans[generation_id] = (span, token, _current_task()) + _in_flight_spans[generation_id] = (span, token) return span @@ -439,7 +375,7 @@ def finish_backend_span_success( entry = _in_flight_spans.pop(generation_id, None) if entry is None: return - span, token, attach_task = entry + span, token = entry try: if gen is not None: set_request_attrs(span, gen, operation) @@ -448,7 +384,7 @@ def finish_backend_span_success( if mot is not None: set_mellea_attrs(span, mot) finally: - _safe_detach(token, attach_task) + _detach_token(token) span.end() @@ -471,7 +407,7 @@ def finish_backend_span_error( entry = _in_flight_spans.pop(generation_id, None) if entry is None: return - span, token, attach_task = entry + span, token = entry try: if gen is not None: set_request_attrs(span, gen, operation) @@ -479,7 +415,7 @@ def finish_backend_span_error( span.set_status(trace.Status(trace.StatusCode.ERROR, str(exception))) span.set_attribute("error.type", type(exception).__name__) finally: - _safe_detach(token, attach_task) + _detach_token(token) span.end() @@ -507,7 +443,7 @@ def _start_application_span( set_attribute_safe(span, k, v) token = _attach_span_context(span, attach=attach_context) - _in_flight_spans[key] = (span, token, _current_task()) + _in_flight_spans[key] = (span, token) return span @@ -527,13 +463,13 @@ def _finish_application_span_success( entry = _in_flight_spans.pop(key, None) if entry is None: return - span, token, attach_task = entry + span, token = entry try: if extra_attributes: for k, v in extra_attributes.items(): set_attribute_safe(span, k, v) finally: - _safe_detach(token, attach_task) + _detach_token(token) span.end() @@ -559,7 +495,7 @@ def _finish_application_span_error( entry = _in_flight_spans.pop(key, None) if entry is None: return - span, token, attach_task = entry + span, token = entry try: if extra_attributes: for k, v in extra_attributes.items(): @@ -571,7 +507,7 @@ def _finish_application_span_error( else: span.set_status(trace.Status(trace.StatusCode.ERROR, description or "")) finally: - _safe_detach(token, attach_task) + _detach_token(token) span.end() @@ -842,7 +778,7 @@ def start_streaming_span( chunking_strategy: str | None, attach_context: bool = True, ) -> Span | None: - """Open the `stream_with_chunking` span for one orchestration run. + """Open the `stream` span for one streaming run. Args: streaming_id: UUID correlating this streaming run across hooks. @@ -855,7 +791,7 @@ def start_streaming_span( The span, or `None` if tracing is disabled. """ return _start_application_span( - "stream_with_chunking", + "stream", streaming_id, { "mellea.streaming.has_requirements": has_requirements, @@ -884,40 +820,6 @@ def add_span_event(key: str, *, event_name: str, attributes: dict[str, Any]) -> span.add_event(event_name, filtered) -def reattach_span(key: str) -> None: - """Make the in-flight span `key` the current task's ambient context. - - Spans opened by later work on this task then parent under it. Paired with - `release_reattached_span()`, which must run on the same task. No-op when the - span is not in flight. See `_safe_detach` for how this scope is used to - classify the cross-task detach it enables. - - Args: - key: Correlation key of an in-flight span (the key it was stashed under). - """ - entry = _in_flight_spans.get(key) - if entry is None: - return - span = entry[0] - token = otel_context.attach(trace.set_span_in_context(span)) - _reattached_tokens[key] = (token, _current_task()) - - -def release_reattached_span(key: str) -> None: - """Release a reattached span from a matching `reattach_span()` call. - - Must run on the same task that called `reattach_span()`. No-op when no token - is stored. - - Args: - key: Correlation key from the matching `reattach_span()` call. - """ - entry = _reattached_tokens.pop(key, None) - if entry is not None: - token, _ = entry - otel_context.detach(token) - - def finish_streaming_span( streaming_id: str, *, @@ -928,7 +830,7 @@ def finish_streaming_span( provider: str | None = None, full_text_length: int | None = None, ) -> None: - """End the `stream_with_chunking` span, recording its outcome. + """End the `stream` span, recording its outcome. Sets OK status on success. On failure, marks the span ERROR: with the exception recorded when one is given, otherwise with `failure_reason` and @@ -942,7 +844,7 @@ def finish_streaming_span( exception: The exception raised by the orchestrator, when one was. model: Model identifier, when known. provider: Provider name, when known. - full_text_length: Accumulated text length at orchestrator exit. + full_text_length: Length of the validated-and-emitted text at stream exit. """ extra_attributes = { "mellea.streaming.full_text_length": full_text_length, diff --git a/mellea/telemetry/tracing_plugins.py b/mellea/telemetry/tracing_plugins.py index a0014cb8c8..112a5d6a10 100644 --- a/mellea/telemetry/tracing_plugins.py +++ b/mellea/telemetry/tracing_plugins.py @@ -11,8 +11,8 @@ events from the generation_event hook. - ComponentTracingPlugin: Emits application-level spans tracking component execution. -- StreamingTracingPlugin: Emits an application-level orchestration span and - per-chunk span events for `stream_with_chunking` runs. +- StreamingTracingPlugin: Emits an application-level span and per-chunk span + events for `stream` runs. - ToolTracingPlugin: Emits an `execute_tool` span for every tool invocation. - SamplingTracingPlugin: Emits a `sampling` span per sampling loop, with a span event per iteration and repair. @@ -58,8 +58,6 @@ from mellea.plugins.hooks.streaming import ( StreamingEndPayload, StreamingEventPayload, - StreamingOrchestrationEndPayload, - StreamingOrchestrationStartPayload, StreamingStartPayload, ) from mellea.plugins.hooks.tool import ToolPostInvokePayload, ToolPreInvokePayload @@ -311,25 +309,20 @@ async def on_component_post_error( class StreamingTracingPlugin(Plugin, name="streaming_tracing", priority=1042): - """Emits the `stream_with_chunking` application span. + """Emits the `stream` application span. `streaming_start` opens the span; `streaming_event` records a span event for - each mid-stream `StreamEvent`; `streaming_end` records the `completed` span - event and closes the span. `streaming_orchestration_start` / - `streaming_orchestration_end` re-attach the span on the orchestration task - so mid-stream spans parent under it (see `reattach_span`). - - All hooks run SEQUENTIAL: the OTel context Token attached in start is - detached on the originating task in end, and `streaming_orchestration_end` - releases the reattached span before `streaming_end` closes it - (FIRE_AND_FORGET would reorder these and break span nesting). + each `StreamEvent`; `streaming_end` closes the span. + + All hooks run SEQUENTIAL so the OTel context Token attached in start is + detached on the same task in end. """ @hook("streaming_start") async def on_streaming_start( self, payload: StreamingStartPayload, context: dict[str, Any] ) -> None: - """Open the stream_with_chunking span for this orchestrator invocation.""" + """Open the stream span for this streaming run.""" if not payload.streaming_id: return from mellea.telemetry.tracing import start_streaming_span @@ -342,28 +335,6 @@ async def on_streaming_start( attach_context=_CONTEXT_ATTACH_SUPPORTED, ) - @hook("streaming_orchestration_start") - async def on_streaming_orchestration_start( - self, payload: StreamingOrchestrationStartPayload, context: dict[str, Any] - ) -> None: - """Re-attach the streaming span as the orchestration task's ambient context.""" - if not payload.streaming_id or not _CONTEXT_ATTACH_SUPPORTED: - return - from mellea.telemetry.tracing import reattach_span - - reattach_span(payload.streaming_id) - - @hook("streaming_orchestration_end") - async def on_streaming_orchestration_end( - self, payload: StreamingOrchestrationEndPayload, context: dict[str, Any] - ) -> None: - """Detach the streaming span re-attached on the orchestration task.""" - if not payload.streaming_id: - return - from mellea.telemetry.tracing import release_reattached_span - - release_reattached_span(payload.streaming_id) - @hook("streaming_event") async def on_streaming_event( self, payload: StreamingEventPayload, context: dict[str, Any] @@ -373,6 +344,7 @@ async def on_streaming_event( return from mellea.stdlib.streaming import ( ChunkEvent, + CompletedEvent, ErrorEvent, FullValidationEvent, QuickCheckEvent, @@ -425,24 +397,25 @@ async def on_streaming_event( "mellea.error.detail": ev.detail, }, ) + elif isinstance(ev, CompletedEvent): + add_span_event( + payload.streaming_id, + event_name="completed", + attributes={ + "mellea.streaming.success": ev.success, + "mellea.streaming.full_text_length": len(ev.full_text), + }, + ) @hook("streaming_end") async def on_streaming_end( self, payload: StreamingEndPayload, context: dict[str, Any] ) -> None: - """Record the `completed` span event and close the stream_with_chunking span.""" + """Close the stream span.""" if not payload.streaming_id: return - from mellea.telemetry.tracing import add_span_event, finish_streaming_span + from mellea.telemetry.tracing import finish_streaming_span - add_span_event( - payload.streaming_id, - event_name="completed", - attributes={ - "mellea.streaming.success": payload.success, - "mellea.streaming.full_text_length": payload.full_text_length, - }, - ) finish_streaming_span( payload.streaming_id, success=payload.success, diff --git a/test/core/test_astream_mock.py b/test/core/test_astream_mock.py index 835b12ef45..f3e4db03bf 100644 --- a/test/core/test_astream_mock.py +++ b/test/core/test_astream_mock.py @@ -166,3 +166,71 @@ async def test_astream_final_call_returns_full_value(): # All chunks concatenated equal the full value assert chunk1 + chunk2 + chunk3 == "part1part2part3" assert mot.value == "part1part2part3" + + +@pytest.mark.asyncio +async def test_aiter_yields_streamed_deltas(): + """`async for delta in mot` yields each streamed delta; joined, they equal `value`.""" + mot = create_manual_mock_thunk() + + async def _feed() -> None: + for token in ("alpha ", "beta ", "gamma"): + await mot._gen.queue.put(token) + await asyncio.sleep(0) + await mot._gen.queue.put(None) + + task = asyncio.create_task(_feed()) + deltas = [delta async for delta in mot] + await task + + # Deltas are the incremental text; concatenated they equal the full value. + assert "".join(deltas) == "alpha beta gamma" + assert mot.value == "alpha beta gamma" + + +@pytest.mark.asyncio +async def test_aiter_on_computed_thunk_stops_immediately(): + """Iterating an already-computed thunk yields nothing (StopAsyncIteration).""" + mot = ModelOutputThunk(value="done") + assert [delta async for delta in mot] == [] + + +@pytest.mark.asyncio +async def test_aiter_single_consumer_guard_raises_on_second_iterator(): + """A second `async for` over the same thunk raises rather than splitting the stream. + + `__aiter__` marks the thunk consumed on first iteration; a second consumer + would split the one underlying stream, so it is rejected. + """ + mot = create_manual_mock_thunk() + mot._gen.queue.put_nowait("only") + mot._gen.queue.put_nowait(None) + + first = mot.__aiter__() + assert first is mot + + with pytest.raises(RuntimeError, match="already being iterated"): + mot.__aiter__() + + +@pytest.mark.asyncio +async def test_aiter_guard_rejects_reiteration_after_break() -> None: + """The guard is per-thunk and one-shot: once started, no re-iteration. + + Even after breaking out early, the same thunk cannot be re-iterated — the + stream has a single cursor. + """ + mot = create_manual_mock_thunk() + for token in ("one ", "two ", "three"): + mot._gen.queue.put_nowait(token) + mot._gen.queue.put_nowait(None) + + collected = [] + async for delta in mot: + collected.append(delta) + break + assert collected # at least the first delta was yielded + + with pytest.raises(RuntimeError, match="already being iterated"): + async for _ in mot: + pass diff --git a/test/core/test_base.py b/test/core/test_base.py index f663bf60c3..6038e714c7 100644 --- a/test/core/test_base.py +++ b/test/core/test_base.py @@ -689,6 +689,19 @@ def test_mot_error_carried_by_copy_methods() -> None: assert target.error is err +def test_mot_cancelled_carried_by_copy_methods() -> None: + """`_cancelled` survives `copy`, `deepcopy`, and `_copy_from`.""" + mot = ModelOutputThunk(value="") + mot._cancelled = True + + assert copy.copy(mot).cancelled is True + assert copy.deepcopy(mot).cancelled is True + + target = ModelOutputThunk(value=None) + target._copy_from(mot) + assert target.cancelled is True + + def test_mot_thinking_public_field_round_trip(): mot = ModelOutputThunk(value="x") mot.thinking = "reasoning trace" diff --git a/test/plugins/test_hook_call_sites.py b/test/plugins/test_hook_call_sites.py index 033127914d..b0fafa9159 100644 --- a/test/plugins/test_hook_call_sites.py +++ b/test/plugins/test_hook_call_sites.py @@ -1746,54 +1746,46 @@ async def _generate_from_raw(self, actions, ctx, **kwargs): class TestStreamingHookCallSites: - """STREAMING_START/EVENT/END fire in stream_with_chunking() and acomplete().""" + """STREAMING_START/EVENT/END fire across a stream() run consumed with async for.""" - async def test_streaming_start_fires_once_with_payload(self) -> None: - """STREAMING_START fires once carrying requirement and chunking metadata.""" - from mellea.stdlib.streaming import stream_with_chunking + async def test_start_and_end_fire_once_across_a_drained_stream(self) -> None: + """Draining a stream fires STREAMING_START and STREAMING_END once each, with payloads.""" + from mellea.stdlib.streaming import stream - observed: list[Any] = [] + starts: list[Any] = [] + ends: list[Any] = [] @hook("streaming_start") - async def recorder(payload: Any, ctx: Any) -> Any: - observed.append(payload) + async def on_start(payload: Any, ctx: Any) -> Any: + starts.append(payload) return None - register(recorder) - await stream_with_chunking( - CBlock("prompt"), _StreamingBackend(), SimpleContext(), chunking="sentence" - ) - - assert len(observed) == 1 - assert observed[0].has_requirements is False - assert observed[0].requirement_count == 0 - assert observed[0].chunking_strategy == "SentenceChunker" - - async def test_streaming_end_fires_once_on_completion(self) -> None: - """acomplete() fires STREAMING_END once with success and model metadata.""" - from mellea.stdlib.streaming import stream_with_chunking - - observed: list[Any] = [] - @hook("streaming_end") - async def recorder(payload: Any, ctx: Any) -> Any: - observed.append(payload) + async def on_end(payload: Any, ctx: Any) -> Any: + ends.append(payload) return None - register(recorder) - result = await stream_with_chunking( - CBlock("prompt"), _StreamingBackend(), SimpleContext() - ) - await result.acomplete() + register(on_start) + register(on_end) + async with await stream( + CBlock("prompt"), _StreamingBackend(), SimpleContext(), chunking="sentence" + ) as s: + async for _chunk in s: + pass - assert len(observed) == 1 - assert observed[0].success is True - assert observed[0].model == "stream-mock-model" - assert observed[0].provider == "stream-mock-provider" + assert len(starts) == 1 + assert starts[0].has_requirements is False + assert starts[0].requirement_count == 0 + assert starts[0].chunking_strategy == "SentenceChunking" + + assert len(ends) == 1 + assert ends[0].success is True + assert ends[0].model == "stream-mock-model" + assert ends[0].provider == "stream-mock-provider" - async def test_streaming_end_fires_once_across_repeat_acomplete(self) -> None: - """Repeat acomplete() calls fire STREAMING_END exactly once.""" - from mellea.stdlib.streaming import stream_with_chunking + async def test_streaming_end_fires_once_across_repeat_aclose(self) -> None: + """aclose() fires STREAMING_END once; a repeat aclose() does not re-fire it.""" + from mellea.stdlib.streaming import stream observed: list[Any] = [] @@ -1803,60 +1795,17 @@ async def recorder(payload: Any, ctx: Any) -> Any: return None register(recorder) - result = await stream_with_chunking( - CBlock("prompt"), _StreamingBackend(), SimpleContext() - ) - await result.acomplete() - await result.acomplete() - - assert len(observed) == 1 - - async def test_streaming_orchestration_start_fires_once_on_orch_task(self) -> None: - """STREAMING_ORCHESTRATION_START fires once, after consumption begins.""" - from mellea.stdlib.streaming import stream_with_chunking - - observed: list[Any] = [] - - @hook("streaming_orchestration_start") - async def recorder(payload: Any, ctx: Any) -> Any: - observed.append(payload.streaming_id) - return None - - register(recorder) - result = await stream_with_chunking( - CBlock("prompt"), _StreamingBackend(), SimpleContext() - ) - # The hook fires inside the orchestration task, which only runs once the - # result is consumed — nothing observed until astream()/acomplete(). - assert observed == [] - await result.acomplete() - - assert len(observed) == 1 - assert observed[0] == result._streaming_id - - async def test_streaming_orchestration_end_fires_once_on_completion(self) -> None: - """STREAMING_ORCHESTRATION_END fires once, pairing with the start hook.""" - from mellea.stdlib.streaming import stream_with_chunking - - observed: list[Any] = [] - - @hook("streaming_orchestration_end") - async def recorder(payload: Any, ctx: Any) -> Any: - observed.append(payload.streaming_id) - return None - - register(recorder) - result = await stream_with_chunking( - CBlock("prompt"), _StreamingBackend(), SimpleContext() - ) - await result.acomplete() + # Abnormal path on purpose: the Streamer is never consumed, so aclose() + # fires END rather than a natural drain. + s = await stream(CBlock("prompt"), _StreamingBackend(), SimpleContext()) + await s.aclose() + await s.aclose() assert len(observed) == 1 - assert observed[0] == result._streaming_id async def test_streaming_end_fires_when_generation_raises(self) -> None: - """A backend failure before streaming fires STREAMING_END with no model.""" - from mellea.stdlib.streaming import stream_with_chunking + """A backend failure during setup fires STREAMING_END with no model.""" + from mellea.stdlib.streaming import stream class _RaisingBackend(Backend): _model_id = "x" @@ -1877,9 +1826,7 @@ async def recorder(payload: Any, ctx: Any) -> Any: register(recorder) with pytest.raises(RuntimeError, match="backend down"): - await stream_with_chunking( - CBlock("prompt"), _RaisingBackend(), SimpleContext() - ) + await stream(CBlock("prompt"), _RaisingBackend(), SimpleContext()) assert len(observed) == 1 assert observed[0].success is False diff --git a/test/stdlib/test_chunking.py b/test/stdlib/test_chunking.py index 235f3b380b..f3261450ba 100644 --- a/test/stdlib/test_chunking.py +++ b/test/stdlib/test_chunking.py @@ -1,15 +1,20 @@ # Copyright IBM Corp. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -"""Tests for ChunkingStrategy ABC and built-in chunker implementations.""" +"""Tests for ChunkingStrategy ABC, built-in strategies, and the Chunker driver.""" + +import random +from collections.abc import Callable import pytest from mellea.stdlib.chunking import ( + Chunker, ChunkingStrategy, - ParagraphChunker, - SentenceChunker, - WordChunker, + ParagraphChunking, + SentenceChunking, + WordChunking, + resolve_chunking_strategy, ) @@ -19,58 +24,58 @@ def test_chunking_strategy_is_abstract(): # --------------------------------------------------------------------------- -# SentenceChunker +# SentenceChunking # --------------------------------------------------------------------------- def test_sentence_chunker_empty(): - c = SentenceChunker() + c = SentenceChunking() assert c.split("") == [] def test_sentence_chunker_no_boundary(): - c = SentenceChunker() + c = SentenceChunking() assert c.split("The quick brown") == [] def test_sentence_chunker_one_sentence_no_trailing(): # A sentence with no following whitespace is a trailing fragment — withheld. - c = SentenceChunker() + c = SentenceChunking() assert c.split("The quick brown fox.") == [] def test_sentence_chunker_one_sentence_with_space(): # Sentence followed by a space signals completion. - c = SentenceChunker() + c = SentenceChunking() assert c.split("The quick brown fox. ") == ["The quick brown fox."] def test_sentence_chunker_with_trailing(): - c = SentenceChunker() + c = SentenceChunking() result = c.split("The quick brown fox. He") assert result == ["The quick brown fox."] def test_sentence_chunker_multiple(): - c = SentenceChunker() + c = SentenceChunking() result = c.split("Hello world. Goodbye world. ") assert result == ["Hello world.", "Goodbye world."] def test_sentence_chunker_exclamation(): - c = SentenceChunker() + c = SentenceChunking() result = c.split("Stop! Go. ") assert result == ["Stop!", "Go."] def test_sentence_chunker_question(): - c = SentenceChunker() + c = SentenceChunking() result = c.split("Are you sure? Yes. ") assert result == ["Are you sure?", "Yes."] def test_sentence_chunker_closing_quote(): - c = SentenceChunker() + c = SentenceChunking() result = c.split('He said "hello." She left. ') assert result == ['He said "hello."', "She left."] @@ -78,46 +83,46 @@ def test_sentence_chunker_closing_quote(): def test_sentence_chunker_curly_quotes(): # Verifies U+201D (right double curly quote) and U+2019 (right single curly quote) # are recognised as closing marks after sentence-ending punctuation. - c = SentenceChunker() + c = SentenceChunking() result = c.split("She said \u201cdone.\u201d Next sentence. ") assert result == ["She said \u201cdone.\u201d", "Next sentence."] def test_sentence_chunker_unicode(): - c = SentenceChunker() + c = SentenceChunking() result = c.split("Ça va bien. C'est délicieux. ") assert result == ["Ça va bien.", "C'est délicieux."] def test_sentence_chunker_closing_paren(): - c = SentenceChunker() + c = SentenceChunking() result = c.split("(See note.) Continue here. ") assert result == ["(See note.)", "Continue here."] def test_sentence_chunker_double_space_separator(): # Regression: double-space between sentences must not leak into next chunk. - c = SentenceChunker() + c = SentenceChunking() result = c.split("First. Second. ") assert result == ["First.", "Second."] def test_sentence_chunker_tab_separator(): - c = SentenceChunker() + c = SentenceChunking() result = c.split("First.\tSecond. ") assert result == ["First.", "Second."] def test_sentence_chunker_abbreviation_known_bad(): # Known edge case: abbreviations cause a spurious split (simple regex, not NLP). - c = SentenceChunker() + c = SentenceChunking() result = c.split("Dr. Smith went home. He was tired. ") assert result == ["Dr.", "Smith went home.", "He was tired."] def test_sentence_chunker_incremental_simulation(): # Simulate accumulating text token by token. - c = SentenceChunker() + c = SentenceChunking() assert c.split("The") == [] assert c.split("The quick") == [] assert c.split("The quick brown fox.") == [] @@ -130,51 +135,51 @@ def test_sentence_chunker_incremental_simulation(): # --------------------------------------------------------------------------- -# WordChunker +# WordChunking # --------------------------------------------------------------------------- def test_word_chunker_empty(): - c = WordChunker() + c = WordChunking() assert c.split("") == [] def test_word_chunker_no_boundary(): - c = WordChunker() + c = WordChunking() assert c.split("hello") == [] def test_word_chunker_one_word_with_space(): - c = WordChunker() + c = WordChunking() assert c.split("hello ") == ["hello"] def test_word_chunker_trailing_fragment(): - c = WordChunker() + c = WordChunking() result = c.split("hello world") assert result == ["hello"] def test_word_chunker_multiple_words(): - c = WordChunker() + c = WordChunking() result = c.split("one two three ") assert result == ["one", "two", "three"] def test_word_chunker_multiple_spaces(): - c = WordChunker() + c = WordChunking() result = c.split("one two three ") assert result == ["one", "two", "three"] def test_word_chunker_unicode(): - c = WordChunker() + c = WordChunking() result = c.split("naïve résumé ") assert result == ["naïve", "résumé"] def test_word_chunker_incremental_simulation(): - c = WordChunker() + c = WordChunking() assert c.split("foo") == [] assert c.split("foobar") == [] assert c.split("foobar ") == ["foobar"] @@ -185,58 +190,58 @@ def test_word_chunker_incremental_simulation(): def test_word_chunker_leading_whitespace(): # re.split on " hello world" produces ['', 'hello', 'world'] — empty first # element must be stripped. - c = WordChunker() + c = WordChunking() result = c.split(" hello world ") assert result == ["hello", "world"] # --------------------------------------------------------------------------- -# ParagraphChunker +# ParagraphChunking # --------------------------------------------------------------------------- def test_paragraph_chunker_empty(): - c = ParagraphChunker() + c = ParagraphChunking() assert c.split("") == [] def test_paragraph_chunker_no_boundary(): - c = ParagraphChunker() + c = ParagraphChunking() assert c.split("Just one paragraph with no double newline") == [] def test_paragraph_chunker_one_complete_paragraph(): - c = ParagraphChunker() + c = ParagraphChunking() result = c.split("First paragraph.\n\n") assert result == ["First paragraph."] def test_paragraph_chunker_with_trailing(): - c = ParagraphChunker() + c = ParagraphChunking() result = c.split("First paragraph.\n\nSecond paragraph in progress") assert result == ["First paragraph."] def test_paragraph_chunker_multiple(): - c = ParagraphChunker() + c = ParagraphChunking() result = c.split("Para one.\n\nPara two.\n\n") assert result == ["Para one.", "Para two."] def test_paragraph_chunker_triple_newline(): - c = ParagraphChunker() + c = ParagraphChunking() result = c.split("Para one.\n\n\nPara two.\n\n") assert result == ["Para one.", "Para two."] def test_paragraph_chunker_unicode(): - c = ParagraphChunker() + c = ParagraphChunking() result = c.split("Première partie.\n\nDeuxième partie.\n\n") assert result == ["Première partie.", "Deuxième partie."] def test_paragraph_chunker_incremental_simulation(): - c = ParagraphChunker() + c = ParagraphChunking() assert c.split("First") == [] assert c.split("First paragraph.") == [] assert c.split("First paragraph.\n\n") == ["First paragraph."] @@ -256,8 +261,8 @@ def test_default_flush_returns_empty_list(): """The ABC default discards the trailing fragment.""" class Minimal(ChunkingStrategy): - def split(self, accumulated_text: str) -> list[str]: - _ = accumulated_text + def split(self, text: str) -> list[str]: + _ = text return [] assert Minimal().flush("anything at all") == [] @@ -265,59 +270,248 @@ def split(self, accumulated_text: str) -> list[str]: def test_sentence_chunker_flush_empty(): - assert SentenceChunker().flush("") == [] + assert SentenceChunking().flush("") == [] def test_sentence_chunker_flush_only_complete(): """All text ends in a complete sentence with trailing whitespace → no fragment.""" - assert SentenceChunker().flush("One. Two. ") == [] + assert SentenceChunking().flush("One. Two. ") == [] def test_sentence_chunker_flush_trailing_fragment(): """Final sentence without trailing whitespace is released by flush.""" - assert SentenceChunker().flush("One. Two without period") == ["Two without period"] + assert SentenceChunking().flush("One. Two without period") == ["Two without period"] def test_sentence_chunker_flush_terminated_no_trailing_space(): """Final sentence with terminator but no trailing whitespace is a fragment under split() semantics and gets released by flush().""" - assert SentenceChunker().flush("One. Two.") == ["Two."] + assert SentenceChunking().flush("One. Two.") == ["Two."] def test_sentence_chunker_flush_single_sentence_no_terminator(): - assert SentenceChunker().flush("Incomplete sentence") == ["Incomplete sentence"] + assert SentenceChunking().flush("Incomplete sentence") == ["Incomplete sentence"] def test_word_chunker_flush_empty(): - assert WordChunker().flush("") == [] + assert WordChunking().flush("") == [] def test_word_chunker_flush_trailing_whitespace(): """Trailing whitespace means all words are complete → no fragment.""" - assert WordChunker().flush("one two three ") == [] + assert WordChunking().flush("one two three ") == [] def test_word_chunker_flush_trailing_fragment(): - assert WordChunker().flush("one two three") == ["three"] + assert WordChunking().flush("one two three") == ["three"] def test_word_chunker_flush_single_word(): - assert WordChunker().flush("solo") == ["solo"] + assert WordChunking().flush("solo") == ["solo"] def test_paragraph_chunker_flush_empty(): - assert ParagraphChunker().flush("") == [] + assert ParagraphChunking().flush("") == [] def test_paragraph_chunker_flush_only_complete(): - assert ParagraphChunker().flush("Para one.\n\nPara two.\n\n") == [] + assert ParagraphChunking().flush("Para one.\n\nPara two.\n\n") == [] def test_paragraph_chunker_flush_trailing_fragment(): - assert ParagraphChunker().flush("Para one.\n\nPara two (no sep)") == [ + assert ParagraphChunking().flush("Para one.\n\nPara two (no sep)") == [ "Para two (no sep)" ] def test_paragraph_chunker_flush_single_paragraph_no_separator(): - assert ParagraphChunker().flush("Only paragraph") == ["Only paragraph"] + assert ParagraphChunking().flush("Only paragraph") == ["Only paragraph"] + + +# --------------------------------------------------------------------------- +# resolve_chunking_strategy +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("alias", "expected"), + [ + ("sentence", SentenceChunking), + ("word", WordChunking), + ("paragraph", ParagraphChunking), + ], +) +def test_resolve_chunking_strategy_alias(alias: str, expected: type) -> None: + assert isinstance(resolve_chunking_strategy(alias), expected) + + +def test_resolve_chunking_strategy_none_passes_through() -> None: + assert resolve_chunking_strategy(None) is None + + +def test_resolve_chunking_strategy_instance_passes_through() -> None: + strategy = SentenceChunking() + assert resolve_chunking_strategy(strategy) is strategy + + +def test_resolve_chunking_strategy_unknown_alias_raises() -> None: + with pytest.raises(ValueError, match="Unknown chunking alias"): + resolve_chunking_strategy("nonsense") + + +# --------------------------------------------------------------------------- +# Chunker — stateful incremental driver +# --------------------------------------------------------------------------- + + +def _run_chunker(strategy: ChunkingStrategy, deltas: list[str]) -> list[str]: + """Feed `deltas` through a Chunker and return all chunks including flush.""" + chunker = Chunker(strategy) + out: list[str] = [] + for delta in deltas: + out.extend(chunker.feed(delta)) + out.extend(chunker.flush()) + return out + + +def _reference(strategy: ChunkingStrategy, full: str) -> list[str]: + """The one-shot result: split the whole text, then flush the remainder.""" + return strategy.split(full) + strategy.flush(full) + + +# Texts covering the delta-seam risks called out in #1440: sentence `.`+space, +# word partial-word, and the greedy paragraph `\n{2,}` boundary. +_INVARIANCE_TEXTS = { + "sentence": [ + "Hello world. Goodbye now. Bye.", + "First. Second. Third", + "One. Two. Three. Four", + ' Leading. She said "hi." Done.', + "No boundary here", + "Ends with terminator.", + "", + ], + "word": [ + "foo bar baz qux", + "one two three", + " spaced out ", + "trailing fragment", + "single", + "", + ], + "paragraph": [ + "A\n\n\nB", + "Para one\n\nPara two\n\nPara three", + "x\n\ny\n\n\nz\n\n", + "line\nstill same\n\nnext", + "one\n\n\n\n\ntwo", + "no break at all", + "", + ], +} + +_STRATEGIES: dict[str, Callable[[], ChunkingStrategy]] = { + "sentence": SentenceChunking, + "word": WordChunking, + "paragraph": ParagraphChunking, +} + + +def _all_texts() -> list[tuple[str, str]]: + return [(name, text) for name, texts in _INVARIANCE_TEXTS.items() for text in texts] + + +@pytest.mark.parametrize(("name", "text"), _all_texts()) +def test_chunker_delta_invariance_single_delta(name: str, text: str) -> None: + """Feeding the whole text as one delta matches one-shot split+flush.""" + strategy_cls = _STRATEGIES[name] + assert _run_chunker(strategy_cls(), [text]) == _reference(strategy_cls(), text) + + +@pytest.mark.parametrize(("name", "text"), _all_texts()) +def test_chunker_delta_invariance_per_character(name: str, text: str) -> None: + """Feeding one character per delta matches one-shot split+flush.""" + strategy_cls = _STRATEGIES[name] + assert _run_chunker(strategy_cls(), list(text)) == _reference(strategy_cls(), text) + + +@pytest.mark.parametrize(("name", "text"), _all_texts()) +def test_chunker_delta_invariance_random_splits(name: str, text: str) -> None: + """Any random slicing into deltas matches one-shot split+flush. + + This is the core delta-invariance property: the boundary a chunk sits on may + land anywhere relative to the delta seams, so the result must not depend on + how the stream was chopped up. + """ + strategy_cls = _STRATEGIES[name] + reference = _reference(strategy_cls(), text) + rng = random.Random(f"{name}:{text}") # deterministic per case + for _ in range(50): + if len(text) <= 1: + break + k = rng.randint(1, min(6, len(text) - 1)) + cuts = sorted(set(rng.sample(range(1, len(text)), k))) + deltas, prev = [], 0 + for cut in cuts: + deltas.append(text[prev:cut]) + prev = cut + deltas.append(text[prev:]) + assert _run_chunker(strategy_cls(), deltas) == reference, ( + f"delta-invariance broke for {name} on {text!r} sliced as {deltas!r}" + ) + + +def test_chunker_sentence_seam_period_then_space() -> None: + """The sentence terminator and its following space arriving in separate deltas.""" + assert _run_chunker(SentenceChunking(), ["One two.", " Three four. "]) == [ + "One two.", + "Three four.", + ] + + +def test_chunker_word_seam_partial_word() -> None: + """A word split across deltas is held until its trailing space arrives.""" + chunker = Chunker(WordChunking()) + assert chunker.feed("hel") == [] # codespell:ignore + assert chunker.feed("lo wor") == ["hello"] + assert chunker.feed("ld ") == ["world"] + assert chunker.flush() == [] + + +def test_chunker_paragraph_seam_split_newlines() -> None: + """A `\\n\\n` boundary whose newlines arrive in separate deltas is not split early.""" + chunker = Chunker(ParagraphChunking()) + assert chunker.feed("Para one\n") == [] + # A lone second newline completes the boundary; the paragraph is emitted. + assert chunker.feed("\nPara two") == ["Para one"] + assert chunker.flush() == ["Para two"] + + +def test_chunker_flush_without_feed() -> None: + """flush() on an unfed Chunker returns nothing.""" + assert Chunker(SentenceChunking()).flush() == [] + + +def test_chunker_holds_only_pending_fragment() -> None: + """The Chunker's state is the pending fragment, not the whole stream.""" + chunker = Chunker(WordChunking()) + chunker.feed("alpha beta gamma ") + assert "alpha" not in chunker._pending + chunker.feed("delta") + # Only the un-terminated trailing word is held, never already-emitted content. + assert "delta" in chunker._pending + assert "alpha" not in chunker._pending + assert len(chunker._pending) < len("alpha beta gamma delta") + + +def test_chunker_rejects_mutating_strategy() -> None: + """feed() raises if split() returns a chunk that is not a verbatim substring.""" + + class MutatingChunking(ChunkingStrategy): + def split(self, text: str) -> list[str]: + # Normalizes whitespace inside the chunk, so the returned text is no + # longer a substring of the input. + return [" ".join(text.split())] if text.strip() else [] + + with pytest.raises(ValueError, match="verbatim substring"): + Chunker(MutatingChunking()).feed("one two") diff --git a/test/stdlib/test_streaming.py b/test/stdlib/test_streaming.py index d489e50257..2e9bdd2eb3 100644 --- a/test/stdlib/test_streaming.py +++ b/test/stdlib/test_streaming.py @@ -1,18 +1,20 @@ # Copyright IBM Corp. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -"""Tests for stream_with_chunking() and StreamChunkingResult. +"""Tests for stream() and Streamer. Uses StreamingMockBackend — a deterministic test double that feeds tokens from a fixed response string into a MOT queue without network or LLM calls. -All tests are unit tests (no @pytest.mark.ollama needed). +Terminal state (`failed_early`, `full_text`, `final_validations`, +`streaming_failures`, `mot`) is read from the `Streamer` after iteration. Typed +`StreamEvent`s are observed via the `STREAMING_EVENT` hook. """ import asyncio import time +from contextlib import contextmanager from typing import Any -from unittest.mock import patch import pytest @@ -23,11 +25,13 @@ Requirement, ValidationResult, ) +from mellea.plugins import hook, register, unregister from mellea.plugins.manager import ( disable_background_collection, drain_background_tasks, enable_background_collection, ) +from mellea.plugins.registry import _HAS_PLUGIN_FRAMEWORK from mellea.stdlib.context import SimpleContext from mellea.stdlib.streaming import ( ChunkEvent, @@ -36,11 +40,20 @@ FullValidationEvent, QuickCheckEvent, RetryEvent, + Streamer, StreamEvent, StreamingDoneEvent, - stream_with_chunking, + stream, ) +# Tests that observe the stream via the STREAMING_EVENT / STREAMING_END hooks need +# the optional `hooks` extra (cpex); the core streaming tests below do not. Skip +# the hook-observing tests when it is not installed rather than fail on register(). +_cpex_skip = pytest.mark.skipif( + not _HAS_PLUGIN_FRAMEWORK, reason="cpex not installed — install mellea[hooks]" +) + + # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- @@ -50,9 +63,9 @@ async def _drain_fire_and_forget_tasks(): """Drain FIRE_AND_FORGET plugin tasks so they run before the loop closes. - STREAMING_END (in acomplete()) schedules a background task via the suite-wide - `fandf` acceptance plugin, then returns with no further await — the loop tears - down first and the coroutine is GC'd unawaited ("coroutine never awaited"). + STREAMING_END schedules a background task via the suite-wide `fandf` + acceptance plugin, then returns with no further await — the loop tears down + first and the coroutine is GC'd unawaited ("coroutine never awaited"). Collection is global, so disable it after to leave other test files untouched. """ enable_background_collection() @@ -214,21 +227,25 @@ async def validate( return ValidationResult(result=True) -class MutationDetectorReq(Requirement): - """Tracks how many times stream_validate was called on this instance.""" +class ChunkRecordingReq(Requirement): + """Records every chunk passed to stream_validate.""" def __init__(self) -> None: super().__init__() - self._call_count = 0 + self.seen_chunks: list[str] = [] + + def __copy__(self) -> "ChunkRecordingReq": + clone = ChunkRecordingReq() + clone.seen_chunks = [] # fresh list — do not share with original + return clone def format_for_llm(self) -> str: - return "mutation detector" + return "chunk recorder" async def stream_validate( self, chunk: str, *, backend: Any, ctx: Any ) -> PartialValidationResult: - _ = chunk, backend, ctx - self._call_count += 1 + self.seen_chunks.append(chunk) return PartialValidationResult("unknown") async def validate( @@ -250,218 +267,134 @@ def _action() -> CBlock: return CBlock("prompt") +@contextmanager +def _record_events(): + """Install a temporary `streaming_event` hook that collects emitted events. + + Registers a recorder plugin for the duration of the block and unregisters it + on exit. Yields the list each `StreamEvent` is appended to. + """ + events: list[StreamEvent] = [] + + @hook("streaming_event") + async def _recorder(payload: Any, ctx: Any) -> Any: + events.append(payload.event) + return None + + register(_recorder) + try: + yield events + finally: + unregister(_recorder) + + # --------------------------------------------------------------------------- -# Tests +# Consumption + terminal state # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_normal_completion_calls_validate_at_stream_end() -> None: - """All 'unknown' requirements → validate() called at stream end; completed=True.""" + """All 'unknown' requirements → validate() runs at stream end; no early fail.""" response = "Hello world. How are you. " backend = StreamingMockBackend(response, token_size=3) req = AlwaysUnknownReq() - result = await stream_with_chunking( + async with await stream( _action(), backend, _ctx(), requirements=[req], chunking="sentence" - ) - await result.acomplete() + ) as streamer: + async for _chunk in streamer: + pass - assert result.completed is True - assert result.full_text == response - assert len(result.final_validations) == 1 - assert result.final_validations[0].as_bool() is True - assert result.streaming_failures == [] + assert streamer.failed_early is False + assert streamer.full_text == response + assert len(streamer.final_validations) == 1 + assert streamer.final_validations[0].as_bool() is True + assert streamer.streaming_failures == [] + assert streamer.mot is not None + assert streamer.completed_normally is True @pytest.mark.asyncio -async def test_early_exit_on_fail() -> None: - """Requirement fails mid-stream → completed=False, streaming_failures populated.""" - # 5 words to trigger failure - response = "one two three four five six seven eight. " +async def test_mot_set_on_natural_completion() -> None: + """On natural completion, `mot` holds the computed thunk with the full value.""" + response = "One. Two. " backend = StreamingMockBackend(response, token_size=2) - req = FailAfterWordsReq(threshold=4) - - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=[req], chunking="word" - ) - await result.acomplete() - - assert result.completed is False - assert len(result.streaming_failures) == 1 - _req, pvr = result.streaming_failures[0] - assert pvr.success == "fail" - assert pvr.reason == "too many words" - # final_validations should be empty — final validate() skipped on early exit - assert result.final_validations == [] - - -@pytest.mark.asyncio -async def test_clone_isolation_across_retries() -> None: - """Originals must not be mutated; two invocations are independent.""" - response = "Sentence one. Sentence two. " - req = MutationDetectorReq() - original_reqs = [req] - - backend = StreamingMockBackend(response, token_size=4) - - r1 = await stream_with_chunking( - _action(), backend, _ctx(), requirements=original_reqs, chunking="sentence" - ) - await r1.acomplete() - - r2 = await stream_with_chunking( - _action(), backend, _ctx(), requirements=original_reqs, chunking="sentence" - ) - await r2.acomplete() - - # Original requirement must never have been called — only clones are used - assert req._call_count == 0 - - -@pytest.mark.asyncio -async def test_validation_backend_routing() -> None: - """stream_validate and validate receive validation_backend, not the main backend.""" - response = "One sentence. Two sentences. " - main_backend = StreamingMockBackend(response, token_size=3) - val_backend = StreamingMockBackend("unused", token_size=1) - - req = BackendRecordingReq() - - # Capture the cloned requirement so we can inspect which backends it saw. - captured: list[BackendRecordingReq] = [] - original_copy = BackendRecordingReq.__copy__ - - def _capturing_copy(self: BackendRecordingReq) -> BackendRecordingReq: - clone = original_copy(self) - captured.append(clone) - return clone - BackendRecordingReq.__copy__ = _capturing_copy # type: ignore[method-assign] - try: - result = await stream_with_chunking( - _action(), - main_backend, - _ctx(), - requirements=[req], - chunking="sentence", - validation_backend=val_backend, - ) - await result.acomplete() - finally: - BackendRecordingReq.__copy__ = original_copy # type: ignore[method-assign] + async with await stream( + _action(), backend, _ctx(), chunking="sentence" + ) as streamer: + async for _chunk in streamer: + pass - assert result.completed is True - # The original was never called — only clones are used. - assert req.seen_backends == [] - # The clone must have seen val_backend for every call (stream_validate + validate), - # never main_backend. This is the actual routing assertion. - assert len(captured) == 1 - assert len(captured[0].seen_backends) > 0 - assert all(b is val_backend for b in captured[0].seen_backends) + assert streamer.mot is not None + assert streamer.mot.is_computed() + assert streamer.mot.value == streamer.full_text == response @pytest.mark.asyncio -async def test_early_exit_does_not_deadlock() -> None: - """Early failure with a high-throughput stream must not hang.""" - long_response = "word " * 200 - backend = StreamingMockBackend(long_response, token_size=5) - req = FailAfterWordsReq(threshold=3) +async def test_yields_individual_chunks() -> None: + """Each iteration yields one validated chunk, in order.""" + response = "Alpha one. Beta two. Gamma three. " + backend = StreamingMockBackend(response, token_size=3) - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=[req], chunking="word" - ) - # 5-second timeout — should complete in milliseconds on success - await asyncio.wait_for(result.acomplete(), timeout=5.0) + chunks: list[str] = [] + async with await stream( + _action(), backend, _ctx(), chunking="sentence" + ) as streamer: + async for chunk in streamer: + chunks.append(chunk) - assert result.completed is False + assert chunks == ["Alpha one.", "Beta two.", "Gamma three."] @pytest.mark.asyncio -async def test_as_thunk_correctness() -> None: - """as_thunk is computed, value matches full_text, generation metadata preserved.""" - response = "This is a test sentence. " +async def test_no_requirements_streams_without_validation() -> None: + """With no requirements, chunks stream through and no final validation runs.""" + response = "Sentence one. Sentence two. " backend = StreamingMockBackend(response, token_size=4) - result = await stream_with_chunking(_action(), backend, _ctx(), chunking="sentence") - await result.acomplete() + chunks: list[str] = [] + async with await stream( + _action(), backend, _ctx(), chunking="sentence" + ) as streamer: + async for chunk in streamer: + chunks.append(chunk) - thunk = result.as_thunk - assert thunk.is_computed() - assert thunk.value == result.full_text == response + assert streamer.failed_early is False + assert streamer.full_text == response + assert streamer.final_validations == [] + assert streamer.streaming_failures == [] + assert chunks == ["Sentence one.", "Sentence two."] @pytest.mark.asyncio -async def test_as_thunk_raises_before_acomplete() -> None: - """as_thunk raises RuntimeError if accessed before acomplete().""" - response = "Some text. " +async def test_raw_delta_mode_when_chunking_none() -> None: + """chunking=None yields raw deltas verbatim, without re-chunking on boundaries.""" + # Spaces make this a real test: any chunker that split on whitespace would drop + # the inter-word spaces, so the concatenation would not equal the response. + response = "one two three" backend = StreamingMockBackend(response, token_size=2) - result = await stream_with_chunking(_action(), backend, _ctx(), chunking="sentence") - - with pytest.raises(RuntimeError, match="acomplete"): - _ = result.as_thunk - - -@pytest.mark.asyncio -async def test_astream_yields_individual_chunks() -> None: - """Consumer via astream() receives individual chunks, not accumulated text.""" - response = "First sentence. Second sentence. Third sentence. " - backend = StreamingMockBackend(response, token_size=5) - - result = await stream_with_chunking(_action(), backend, _ctx(), chunking="sentence") - chunks: list[str] = [] - async for chunk in result.astream(): - chunks.append(chunk) + async with await stream(_action(), backend, _ctx(), chunking=None) as streamer: + async for chunk in streamer: + chunks.append(chunk) - await result.acomplete() - - # Each chunk must be a complete sentence (not the accumulated text) - assert len(chunks) == 3 - for chunk in chunks: - assert chunk.endswith(".") - # Chunks don't include inter-sentence spaces; joined with a space they appear in full_text - assert " ".join(chunks) in result.full_text + # Deltas may batch, so the exact split is not pinned — but text (spaces + # included) must survive verbatim, proving no whitespace boundary was applied. + assert "".join(chunks) == response + assert streamer.full_text == response + assert any(" " in c for c in chunks) @pytest.mark.asyncio async def test_stream_validate_receives_individual_chunks() -> None: - """stream_validate is called once per chunk with the chunk itself, not accumulated text.""" - - class ChunkRecordingReq(Requirement): - def __init__(self) -> None: - self.seen_chunks: list[str] = [] - - def __copy__(self) -> "ChunkRecordingReq": - clone = ChunkRecordingReq() - clone.seen_chunks = [] - return clone - - def format_for_llm(self) -> str: - return "chunk recorder" - - async def stream_validate( - self, chunk: str, *, backend: Any, ctx: Any - ) -> PartialValidationResult: - self.seen_chunks.append(chunk) - return PartialValidationResult("unknown") - - async def validate( - self, - backend: Any, - ctx: Any, - *, - format: Any = None, - model_options: Any = None, - ) -> ValidationResult: - return ValidationResult(result=True) - - response = "First sentence. Second sentence. Third sentence. " - backend = StreamingMockBackend(response, token_size=4) + """Each stream_validate call receives exactly one chunk, in order.""" + response = "First one. Second two. Third three. " + backend = StreamingMockBackend(response, token_size=3) req = ChunkRecordingReq() - # Capture the cloned requirement used by the orchestrator via a side channel. captured: list[ChunkRecordingReq] = [] original_copy = ChunkRecordingReq.__copy__ @@ -472,56 +405,22 @@ def _capturing_copy(self: ChunkRecordingReq) -> ChunkRecordingReq: ChunkRecordingReq.__copy__ = _capturing_copy # type: ignore[method-assign] try: - result = await stream_with_chunking( + async with await stream( _action(), backend, _ctx(), requirements=[req], chunking="sentence" - ) - await result.acomplete() + ) as streamer: + async for _chunk in streamer: + pass finally: ChunkRecordingReq.__copy__ = original_copy # type: ignore[method-assign] - assert len(captured) == 1 - seen = captured[0].seen_chunks - # Exact match: three separate calls, one per complete sentence, - # each call receiving that sentence and nothing more. Under the old - # accumulated-text semantics, seen would have been - # ["First sentence.", "First sentence. Second sentence.", ...] — - # exact match against the per-chunk list is the direct regression guard. - assert seen == ["First sentence.", "Second sentence.", "Third sentence."] + assert captured[0].seen_chunks == ["First one.", "Second two.", "Third three."] @pytest.mark.asyncio async def test_trailing_fragment_is_flushed_to_consumer() -> None: - """Response without trailing whitespace: final sentence reaches astream() and stream_validate.""" - - class ChunkRecordingReq(Requirement): - def __init__(self) -> None: - self.seen_chunks: list[str] = [] - - def __copy__(self) -> "ChunkRecordingReq": - clone = ChunkRecordingReq() - clone.seen_chunks = [] - return clone - - def format_for_llm(self) -> str: - return "chunk recorder" - - async def stream_validate( - self, chunk: str, *, backend: Any, ctx: Any - ) -> PartialValidationResult: - self.seen_chunks.append(chunk) - return PartialValidationResult("unknown") - - async def validate( - self, - backend: Any, - ctx: Any, - *, - format: Any = None, - model_options: Any = None, - ) -> ValidationResult: - return ValidationResult(result=True) - - # No trailing whitespace after the final sentence — SentenceChunker withholds it. + """A final sentence with no trailing whitespace still reaches the consumer.""" + # No trailing whitespace after the final sentence — the chunker withholds it + # until flush at stream end. response = "First sentence. Second sentence." backend = StreamingMockBackend(response, token_size=4) req = ChunkRecordingReq() @@ -536,29 +435,55 @@ def _capturing_copy(self: ChunkRecordingReq) -> ChunkRecordingReq: ChunkRecordingReq.__copy__ = _capturing_copy # type: ignore[method-assign] try: - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=[req], chunking="sentence" - ) yielded: list[str] = [] - async for chunk in result.astream(): - yielded.append(chunk) - await result.acomplete() + async with await stream( + _action(), backend, _ctx(), requirements=[req], chunking="sentence" + ) as streamer: + async for chunk in streamer: + yielded.append(chunk) finally: ChunkRecordingReq.__copy__ = original_copy # type: ignore[method-assign] - # Both sentences reach the consumer, including the terminating one without trailing whitespace. + # Both sentences reach the consumer, including the terminating one. assert yielded == ["First sentence.", "Second sentence."] # stream_validate was called on both — the flush path is not a shortcut. assert captured[0].seen_chunks == ["First sentence.", "Second sentence."] - assert result.completed is True + assert streamer.failed_early is False + + +# --------------------------------------------------------------------------- +# Early exit on validation failure +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_early_exit_on_fail() -> None: + """A 'fail' stops the stream early, records the failure, skips final validate().""" + response = "one two three four five six seven eight " + backend = StreamingMockBackend(response, token_size=2) + req = FailAfterWordsReq(threshold=3) + + async with await stream( + _action(), backend, _ctx(), requirements=[req], chunking="word" + ) as streamer: + async for _chunk in streamer: + pass + + assert streamer.failed_early is True + assert len(streamer.streaming_failures) == 1 + _req, pvr = streamer.streaming_failures[0] + assert pvr.success == "fail" + assert pvr.reason == "too many words" + assert streamer.final_validations == [] @pytest.mark.asyncio async def test_early_exit_on_trailing_fragment() -> None: - """A fail on the flushed fragment records a streaming failure and skips final validate().""" + """A fail on the flushed fragment records a failure and skips final validate().""" class FailOnSecondSentence(Requirement): def __init__(self) -> None: + super().__init__() self._count = 0 def format_for_llm(self) -> str: @@ -587,90 +512,48 @@ async def validate( backend = StreamingMockBackend(response, token_size=4) req = FailOnSecondSentence() - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=[req], chunking="sentence" - ) yielded: list[str] = [] - async for chunk in result.astream(): - yielded.append(chunk) - await result.acomplete() + async with await stream( + _action(), backend, _ctx(), requirements=[req], chunking="sentence" + ) as streamer: + async for chunk in streamer: + yielded.append(chunk) - assert result.completed is False - assert len(result.streaming_failures) == 1 - # First sentence was emitted; second (the flushed fragment) failed and wasn't emitted. + assert streamer.failed_early is True + assert len(streamer.streaming_failures) == 1 + # First sentence was emitted; the second (flushed fragment) failed, unemitted. assert yielded == ["First sentence."] - # Early exit on fail skips final validate(). - assert result.final_validations == [] - - -@pytest.mark.asyncio -async def test_no_requirements_streams_without_validation() -> None: - """requirements=None → chunks produced, no validate() called.""" - response = "Chunk one. Chunk two. Chunk three. " - backend = StreamingMockBackend(response, token_size=3) - - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=None, chunking="sentence" - ) - await result.acomplete() - - assert result.completed is True - assert result.full_text == response - assert result.final_validations == [] - assert result.streaming_failures == [] - - -@pytest.mark.asyncio -async def test_no_requirements_events_omits_full_validation_event() -> None: - """With no requirements, events() emits StreamingDoneEvent but - NOT FullValidationEvent — there is nothing to validate at stream end.""" - response = "Chunk one. Chunk two. " - backend = StreamingMockBackend(response, token_size=3) - - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=None, chunking="sentence" - ) - await result.acomplete() - - evts = [e async for e in result.events()] - types = [type(e) for e in evts] - - assert StreamingDoneEvent in types - assert FullValidationEvent not in types - assert isinstance(evts[-1], CompletedEvent) - assert evts[-1].success is True + assert streamer.final_validations == [] @pytest.mark.asyncio async def test_multiple_chunks_in_one_batch_with_mid_batch_fail() -> None: - """When one astream() delta produces several complete chunks and one in - the middle fails, earlier chunks emit, failing chunk is recorded, later - chunks are neither validated nor emitted.""" + """When one delta yields several chunks, a mid-batch fail stops before later ones.""" captured: list[Any] = [] - class FailOnNthChunk(Requirement): - def __init__(self, n: int) -> None: - self._n = n - self._calls = 0 - self.seen: list[str] = [] + class FailOnThird(Requirement): + def __init__(self) -> None: + super().__init__() + self._count = 0 + self.seen_chunks: list[str] = [] - def __copy__(self) -> "FailOnNthChunk": - clone = FailOnNthChunk(self._n) + def __copy__(self) -> "FailOnThird": + clone = FailOnThird() captured.append(clone) return clone def format_for_llm(self) -> str: - return f"fail on chunk {self._n}" + return "fail on third chunk" async def stream_validate( self, chunk: str, *, backend: Any, ctx: Any ) -> PartialValidationResult: _ = backend, ctx - self._calls += 1 - self.seen.append(chunk) - if self._calls == self._n: - return PartialValidationResult("fail", reason=f"n={self._n}") + self._count += 1 + self.seen_chunks.append(chunk) + if self._count == 3: + return PartialValidationResult("fail", reason="third chunk") return PartialValidationResult("unknown") async def validate( @@ -681,110 +564,131 @@ async def validate( format: Any = None, model_options: Any = None, ) -> ValidationResult: - _ = backend, ctx, format, model_options return ValidationResult(result=True) - # token_size larger than the whole response → one astream() delta delivers - # the full text, so chunking.split produces 4 sentences in a single batch. + # Whole response arrives as one delta: split() yields four sentences at once. response = "One. Two. Three. Four. " - backend = StreamingMockBackend(response, token_size=100) - req = FailOnNthChunk(n=2) + backend = StreamingMockBackend(response, token_size=len(response)) + req = FailOnThird() - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=[req], chunking="sentence" - ) yielded: list[str] = [] - async for c in result.astream(): - yielded.append(c) - await result.acomplete() - - assert result.completed is False - assert len(result.streaming_failures) == 1 - # Chunk 1 was validated and emitted; chunk 2 was validated and failed - # (NOT emitted); chunks 3 and 4 were NEITHER validated NOR emitted. - assert yielded == ["One."] - assert len(captured) == 1 - assert captured[0].seen == ["One.", "Two."] - assert captured[0]._calls == 2 + async with await stream( + _action(), backend, _ctx(), requirements=[req], chunking="sentence" + ) as streamer: + async for chunk in streamer: + yielded.append(chunk) + + assert streamer.failed_early is True + assert len(streamer.streaming_failures) == 1 + # First two passed and were emitted; the third failed before emission. + assert yielded == ["One.", "Two."] + # The fourth chunk was neither validated nor emitted: validation stopped at + # the failing third. + assert captured[0].seen_chunks == ["One.", "Two.", "Three."] @pytest.mark.asyncio async def test_cancel_generation_invoked_on_fail() -> None: - """Early exit on 'fail' must call mot.cancel_generation() — the spec reason - is that asyncio.Queue(maxsize=20) will block the producer if the consumer - stops without cancelling.""" + """An early fail cancels the backend generation (mot ends computed+cancelled).""" + response = "one two three four five six seven eight nine ten " + backend = StreamingMockBackend(response, token_size=1) + req = FailAfterWordsReq(threshold=2) - from mellea.core.base import ModelOutputThunk + async with await stream( + _action(), backend, _ctx(), requirements=[req], chunking="word" + ) as streamer: + async for _chunk in streamer: + pass - response = "word " * 50 - backend = StreamingMockBackend(response, token_size=3) + assert streamer.failed_early is True + # On early exit the driving MOT was cancelled; `mot` (natural-completion only) + # stays None. + assert streamer.mot is None + assert streamer.completed_normally is False + # The underlying generation was actually cancelled, not merely abandoned. + assert streamer._mot._cancelled is True + assert streamer._mot.is_computed() is True - class FailOnFirstChunk(Requirement): - def format_for_llm(self) -> str: - return "fail immediately" - async def stream_validate( - self, chunk: str, *, backend: Any, ctx: Any - ) -> PartialValidationResult: - _ = chunk, backend, ctx - return PartialValidationResult("fail", reason="nope") +class _FailOnSecondReq(Requirement): + """Passes the first chunk, fails the second — to drive early exit mid-stream.""" - async def validate( - self, - backend: Any, - ctx: Any, - *, - format: Any = None, - model_options: Any = None, - ) -> ValidationResult: - _ = backend, ctx, format, model_options - return ValidationResult(result=True) + def __init__(self) -> None: + super().__init__() + self._count = 0 - call_count = 0 - real_cancel = ModelOutputThunk.cancel_generation + def format_for_llm(self) -> str: + return "fail on second" - async def spy_cancel( - self: ModelOutputThunk, error: Exception | None = None - ) -> None: - nonlocal call_count - call_count += 1 - await real_cancel(self, error) + async def stream_validate( + self, chunk: str, *, backend: Any, ctx: Any + ) -> PartialValidationResult: + _ = chunk, backend, ctx + self._count += 1 + if self._count == 2: + return PartialValidationResult("fail", reason="second") + return PartialValidationResult("unknown") - ModelOutputThunk.cancel_generation = spy_cancel # type: ignore[method-assign] - try: - result = await stream_with_chunking( - _action(), - backend, - _ctx(), - requirements=[FailOnFirstChunk()], - chunking="word", - ) - await asyncio.wait_for(result.acomplete(), timeout=5.0) - finally: - ModelOutputThunk.cancel_generation = real_cancel # type: ignore[method-assign] + async def validate( + self, backend: Any, ctx: Any, *, format: Any = None, model_options: Any = None + ) -> ValidationResult: + return ValidationResult(result=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("token_size", [1, 6, len("One. Two. Three. Four. ")]) +async def test_full_text_is_chunk_exact_on_early_exit(token_size: int) -> None: + """full_text on early exit is the accumulated text through the last EMITTED chunk. + + Chunk-exact, not delta-granular: only "One." was validated and yielded before + the fail on "Two.", so full_text is exactly "One." regardless of how the + response was split into deltas. + """ + response = "One. Two. Three. Four. " + backend = StreamingMockBackend(response, token_size=token_size) + + async with await stream( + _action(), + backend, + _ctx(), + requirements=[_FailOnSecondReq()], + chunking="sentence", + ) as streamer: + async for _chunk in streamer: + pass - assert result.completed is False - assert call_count >= 1 + assert streamer.failed_early is True + assert streamer.full_text == "One." @pytest.mark.asyncio -async def test_cancelled_flag_reflects_cancellation_state() -> None: - """The `cancelled` property on ModelOutputThunk distinguishes an early-exit - cancellation from a normal completion and propagates through `as_thunk`.""" +@pytest.mark.parametrize("token_size", [1, 6, len("One. Two. Three. Four. ")]) +async def test_full_text_spans_multiple_emitted_chunks_on_early_exit( + token_size: int, +) -> None: + """full_text accumulates every emitted chunk, not just the first. + + Failing on the third chunk emits "One." and "Two." first, so full_text is the + concatenation of both — verifying the emitted-text cursor advances across + chunks, regardless of delta boundaries. + """ - # Early exit → cancelled is True, is_computed True, propagates through as_thunk. - fail_response = "word " * 50 - fail_backend = StreamingMockBackend(fail_response, token_size=3) + class FailOnThirdReq(Requirement): + def __init__(self) -> None: + super().__init__() + self._count = 0 - class FailImmediately(Requirement): def format_for_llm(self) -> str: - return "fail immediately" + return "fail on third" async def stream_validate( self, chunk: str, *, backend: Any, ctx: Any ) -> PartialValidationResult: _ = chunk, backend, ctx - return PartialValidationResult("fail", reason="nope") + self._count += 1 + if self._count == 3: + return PartialValidationResult("fail", reason="third") + return PartialValidationResult("unknown") async def validate( self, @@ -794,640 +698,374 @@ async def validate( format: Any = None, model_options: Any = None, ) -> ValidationResult: - _ = backend, ctx, format, model_options return ValidationResult(result=True) - fail_result = await stream_with_chunking( - _action(), - fail_backend, - _ctx(), - requirements=[FailImmediately()], - chunking="word", - ) - await asyncio.wait_for(fail_result.acomplete(), timeout=5.0) - - assert fail_result.completed is False - assert fail_result.as_thunk.cancelled is True - assert fail_result.as_thunk.is_computed() is True - - # Normal completion → cancelled is False. - ok_response = "Hello world. How are you. " - ok_backend = StreamingMockBackend(ok_response, token_size=3) + response = "One. Two. Three. Four. " + backend = StreamingMockBackend(response, token_size=token_size) - ok_result = await stream_with_chunking( - _action(), - ok_backend, - _ctx(), - requirements=[AlwaysUnknownReq()], - chunking="sentence", - ) - await ok_result.acomplete() + async with await stream( + _action(), backend, _ctx(), requirements=[FailOnThirdReq()], chunking="sentence" + ) as streamer: + async for _chunk in streamer: + pass - assert ok_result.completed is True - assert ok_result.as_thunk.cancelled is False - assert ok_result.as_thunk.is_computed() is True + assert streamer.failed_early is True + assert streamer.full_text == "One. Two." @pytest.mark.asyncio -async def test_unknown_chunking_alias_raises_value_error() -> None: - """An unrecognised chunking alias raises ValueError before any backend call.""" - backend = StreamingMockBackend("hello world") - with pytest.raises(ValueError, match="unknown_alias"): - await stream_with_chunking(_action(), backend, _ctx(), chunking="unknown_alias") +async def test_full_text_through_last_emitted_chunk_on_break() -> None: + """A caller `break` leaves full_text at the last chunk actually delivered.""" + response = "One. Two. Three. " + backend = StreamingMockBackend(response, token_size=3) + async with await stream( + _action(), backend, _ctx(), chunking="sentence" + ) as streamer: + async for _chunk in streamer: + break # take only the first chunk -@pytest.mark.asyncio -async def test_exception_in_stream_validate_cancels_generation() -> None: - """Verifies the orchestrator's exception-path cleanup: if stream_validate - raises, cancel_generation() is called and the exception surfaces to the - consumer via astream()/acomplete() without hanging. - - This covers the cancel-on-exception path and the no-hang guarantee. - It does not directly exercise the worst-case "producer already blocked on - full queue" scenario (here the fail happens on chunk 1 so the queue never - fills); the cancel_generation drain logic is covered by its own tests in - test/core/. - """ + assert streamer.full_text == "One." - from mellea.core.base import ModelOutputThunk - class RaisingReq(Requirement): - def format_for_llm(self) -> str: - return "raises" +# --------------------------------------------------------------------------- +# Cleanup contract (async with / aclose) +# --------------------------------------------------------------------------- - async def stream_validate( - self, chunk: str, *, backend: Any, ctx: Any - ) -> PartialValidationResult: - _ = chunk, backend, ctx - raise ValueError("boom") - async def validate( - self, - backend: Any, - ctx: Any, - *, - format: Any = None, - model_options: Any = None, - ) -> ValidationResult: - _ = backend, ctx, format, model_options - return ValidationResult(result=True) - - response = "word " * 50 # enough to fill maxsize=20 queue without cleanup - backend = StreamingMockBackend(response, token_size=3) +@_cpex_skip +@pytest.mark.asyncio +async def test_break_cancels_generation_and_fires_end() -> None: + """`async with` + early break cancels generation and fires STREAMING_END once.""" + response = "One. Two. Three. Four. Five. " + backend = StreamingMockBackend(response, token_size=2) - call_count = 0 - real_cancel = ModelOutputThunk.cancel_generation + ends: list[Any] = [] - async def spy_cancel( - self: ModelOutputThunk, error: Exception | None = None - ) -> None: - nonlocal call_count - call_count += 1 - await real_cancel(self, error) + @hook("streaming_end") + async def _end(payload: Any, ctx: Any) -> Any: + ends.append(payload) + return None - ModelOutputThunk.cancel_generation = spy_cancel # type: ignore[method-assign] + register(_end) try: - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=[RaisingReq()], chunking="word" - ) - with pytest.raises(ValueError, match="boom"): - async for _chunk in result.astream(): - pass - # acomplete must complete (not hang) even though the orchestration - # task raised, because cancel_generation was called in the except path. - await asyncio.wait_for(result.acomplete(), timeout=5.0) + async with await stream(_action(), backend, _ctx(), chunking="sentence") as s: + async for _chunk in s: + break finally: - ModelOutputThunk.cancel_generation = real_cancel # type: ignore[method-assign] + unregister(_end) - assert result.completed is False - assert call_count >= 1 + assert len(ends) == 1 + assert ends[0].success is False + # Generation was cancelled: it did not reach natural completion, so `mot` + # stays None but the underlying stream is computed/cancelled. + assert s.mot is None + assert s.completed_normally is False +@_cpex_skip @pytest.mark.asyncio -async def test_acomplete_surfaces_exception_without_astream() -> None: - """acomplete() must surface orchestrator exceptions even when the - consumer never iterates astream(). - - The alternative — only delivering the exception through the chunk queue - — silently swallows validator failures for callers who skip astream(). - """ - - class RaisingReq(Requirement): - def format_for_llm(self) -> str: - return "raises" +async def test_acquired_never_iterated_aclose_cancels_and_ends() -> None: + """A Streamer acquired but never iterated still cancels + fires END on aclose(). - async def stream_validate( - self, chunk: str, *, backend: Any, ctx: Any - ) -> PartialValidationResult: - _ = chunk, backend, ctx - raise ValueError("surfaced-without-astream") - - async def validate( - self, - backend: Any, - ctx: Any, - *, - format: Any = None, - model_options: Any = None, - ) -> ValidationResult: - _ = backend, ctx, format, model_options - return ValidationResult(result=True) - - response = "word " * 50 - backend = StreamingMockBackend(response, token_size=3) - - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=[RaisingReq()], chunking="word" - ) - # Deliberately skip astream(). wait_for bounds any hang. - with pytest.raises(ValueError, match="surfaced-without-astream"): - await asyncio.wait_for(result.acomplete(), timeout=5.0) - - assert result.completed is False - # Raise-once: a second acomplete() must not re-raise. - await asyncio.wait_for(result.acomplete(), timeout=5.0) - - -@pytest.mark.asyncio -async def test_external_task_cancellation_releases_consumers() -> None: - """External cancellation of the orchestration task must still set _done. - - If the finally cleanup itself contains an `await` (e.g. awaiting a - terminator put into the chunk queue), CancelledError re-raises at that - await and `_done.set()` never runs — any consumer blocked on - `acomplete()` hangs forever. The cleanup must therefore end with - synchronous operations only. + This is the leak the cleanup contract closes: eager generation starts at + stream(), so an abandoned handle must still be released. """ - response = "word " * 200 # long enough that streaming is still in progress + response = "One. Two. Three. " backend = StreamingMockBackend(response, token_size=2) - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=[AlwaysUnknownReq()], chunking="word" - ) + ends: list[Any] = [] + gen_errors: list[Any] = [] - assert result._orchestration_task is not None - # Wait until the orchestration coroutine has started and hit its first - # suspension point. Using _orchestration_started rather than a wall-clock - # sleep avoids the race where a fast runner drains the whole stream within - # the sleep window, making cancel() a no-op on an already-done task. - # _orchestration_started.wait() must precede cancel() — cancelling before - # the first scheduling means the event is never set. - await asyncio.wait_for(result._orchestration_started.wait(), timeout=2.0) - assert not result._orchestration_task.done(), ( - "orchestrator already done before cancel() — test would vacuously pass" - ) + @hook("streaming_end") + async def _end(payload: Any, ctx: Any) -> Any: + ends.append(payload) + return None - # Same mechanism asyncio.wait_for uses on timeout. - result._orchestration_task.cancel() + @hook("generation_error") + async def _gerr(payload: Any, ctx: Any) -> Any: + gen_errors.append(payload) + return None - # _done must be set by the finally cleanup. A hang would time out here. - await asyncio.wait_for(result._done.wait(), timeout=2.0) - assert result._done.is_set() + register(_end) + register(_gerr) + try: + streamer = await stream(_action(), backend, _ctx(), chunking="sentence") + # Never iterated. + await streamer.aclose() + finally: + unregister(_end) + unregister(_gerr) - # acomplete() surfaces the CancelledError via task.exception() and must - # not hang. - with pytest.raises(asyncio.CancelledError): - await asyncio.wait_for(result.acomplete(), timeout=2.0) + assert len(ends) == 1 + assert ends[0].success is False + # The eager, in-flight generation was cancelled. + assert len(gen_errors) >= 1 +@_cpex_skip @pytest.mark.asyncio -async def test_external_cancellation_acomplete_raise_once() -> None: - """Raise-once contract holds for the task-fallback path on external cancel. - - CancelledError bypasses the orchestrator's `except Exception` handler, - so `_orchestration_exception` is never set. `acomplete()` surfaces the - cancel via `self._orchestration_task.exception()` instead — and that - branch must also flip `_exception_surfaced` so a second `acomplete()` - call does not raise the same exception twice. - """ - response = "word " * 200 +async def test_aclose_after_natural_completion_is_noop() -> None: + """aclose() after full drain does not re-fire STREAMING_END.""" + response = "One. Two. " backend = StreamingMockBackend(response, token_size=2) - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=[AlwaysUnknownReq()], chunking="word" - ) + ends: list[Any] = [] - assert result._orchestration_task is not None - await asyncio.wait_for(result._orchestration_started.wait(), timeout=2.0) - assert not result._orchestration_task.done(), ( - "orchestrator already done before cancel() — test would vacuously pass" - ) - result._orchestration_task.cancel() - await asyncio.wait_for(result._done.wait(), timeout=2.0) + @hook("streaming_end") + async def _end(payload: Any, ctx: Any) -> Any: + ends.append(payload) + return None - with pytest.raises(asyncio.CancelledError): - await asyncio.wait_for(result.acomplete(), timeout=2.0) + register(_end) + try: + streamer = await stream(_action(), backend, _ctx(), chunking="sentence") + async for _chunk in streamer: + pass + await streamer.aclose() + finally: + unregister(_end) - # Second call must NOT re-raise — raise-once contract. - await asyncio.wait_for(result.acomplete(), timeout=2.0) + assert len(ends) == 1 + assert ends[0].success is True +@_cpex_skip @pytest.mark.asyncio -async def test_raise_once_acomplete_then_astream() -> None: - """Regression for the raise-once stash bug: acomplete() first, astream() second. +async def test_double_aclose_fires_end_once() -> None: + """Calling aclose() twice fires STREAMING_END exactly once.""" + response = "One. Two. Three. " + backend = StreamingMockBackend(response, token_size=2) - Prior to the fix, acomplete() cleared _orchestration_exception, so a - subsequent astream() call dequeued the exception item, saw the stash was - None, silently skipped it, and returned zero chunks with no error. - """ + ends: list[Any] = [] - class RaisingReq(Requirement): - def format_for_llm(self) -> str: - return "raises" + @hook("streaming_end") + async def _end(payload: Any, ctx: Any) -> Any: + ends.append(payload) + return None - async def stream_validate( - self, chunk: str, *, backend: Any, ctx: Any - ) -> PartialValidationResult: - raise ValueError("raise-once-regression") + register(_end) + try: + async with await stream(_action(), backend, _ctx(), chunking="sentence") as s: + async for _chunk in s: + break + await s.aclose() # explicit second close after context-manager exit + finally: + unregister(_end) - async def validate( - self, - backend: Any, - ctx: Any, - *, - format: Any = None, - model_options: Any = None, - ) -> ValidationResult: - return ValidationResult(result=True) + assert len(ends) == 1 - response = "word " * 10 - backend = StreamingMockBackend(response, token_size=3) - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=[RaisingReq()], chunking="word" - ) - # acomplete() sees the exception first and raises it. - with pytest.raises(ValueError, match="raise-once-regression"): - await asyncio.wait_for(result.acomplete(), timeout=5.0) - - # astream() must NOT re-raise (raise-once semantics). Because the - # exception fired before any chunk was emitted, the queue contains - # [exc, None]. With the separate _exception_surfaced flag, astream() - # correctly skips the exception item and terminates cleanly. Without - # the flag the behaviour is the same, but the guard conflates - # "already surfaced" with "stash was never set" — the flag makes the - # intent unambiguous. - chunks: list[str] = [] - async for chunk in result.astream(): - chunks.append(chunk) - assert chunks == [] # no partial chunks before the exception +async def _feed_tokens_slowly( + mot: ModelOutputThunk, response: str, token_size: int, delay: float +) -> None: + i = 0 + while i < len(response): + await mot._gen.queue.put(response[i : i + token_size]) + await asyncio.sleep(delay) + i += token_size + await mot._gen.queue.put(None) -@pytest.mark.asyncio -async def test_full_text_contains_only_validated_chunks_on_early_exit() -> None: - """full_text must equal exactly what was emitted to the consumer on early exit. +class SlowStreamingMockBackend(StreamingMockBackend): + """Streams with a real per-token delay and exposes the feed task as `_gen.generate`. - When one astream() delta produces N chunks and chunk K fails, full_text - must contain chunks 0..K-1 only — not the failed chunk or any unvalidated - chunks after it in the same delta. + The delay lets an external timeout land mid-stream, and setting `_gen.generate` + means `cancel_generation()` awaits the in-flight task — the path where an + externally cancelled task re-raises `CancelledError`. """ - class FailOnNthChunkText(Requirement): - def __init__(self, n: int) -> None: - self._n = n - self._calls = 0 - - def __copy__(self) -> "FailOnNthChunkText": - return FailOnNthChunkText(self._n) - - def format_for_llm(self) -> str: - return f"fail on chunk {self._n}" - - async def stream_validate( - self, chunk: str, *, backend: Any, ctx: Any - ) -> PartialValidationResult: - self._calls += 1 - if self._calls == self._n: - return PartialValidationResult("fail") - return PartialValidationResult("unknown") - - async def validate( - self, - backend: Any, - ctx: Any, - *, - format: Any = None, - model_options: Any = None, - ) -> ValidationResult: - return ValidationResult(result=True) - - # token_size > full response → single delta with 4 sentences; fail on chunk 2. - response = "One. Two. Three. Four. " - backend = StreamingMockBackend(response, token_size=100) - req = FailOnNthChunkText(n=2) - - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=[req], chunking="sentence" - ) - yielded: list[str] = [] - async for chunk in result.astream(): - yielded.append(chunk) - await result.acomplete() - - assert result.completed is False - # Consumer received only chunk 1. - assert yielded == ["One."] - # full_text must match what the consumer received — not the raw delta. - assert result.full_text == "One." - # as_thunk.value must agree with full_text. - assert result.as_thunk.value == result.full_text - - # Fail on chunk 3: two chunks emitted before early exit. full_text must - # preserve the original inter-sentence spacing from the token stream, not - # the stripped chunk concatenation ("One.Two." would be wrong). - backend2 = StreamingMockBackend(response, token_size=100) - req2 = FailOnNthChunkText(n=3) - result2 = await stream_with_chunking( - _action(), backend2, _ctx(), requirements=[req2], chunking="sentence" - ) - yielded2: list[str] = [] - async for chunk in result2.astream(): - yielded2.append(chunk) - await result2.acomplete() - - assert result2.completed is False - assert yielded2 == ["One.", "Two."] - assert result2.full_text == "One. Two." - assert result2.as_thunk.value == result2.full_text - - -@pytest.mark.asyncio -async def test_cancelled_flag_propagates_through_copy_methods() -> None: - """_cancelled must survive __copy__, __deepcopy__, and _copy_from.""" - from copy import deepcopy - - mot = ModelOutputThunk(value="result") - mot._cancelled = True - - # __copy__ - shallow = mot.__copy__() - assert shallow._cancelled is True, "__copy__ must propagate _cancelled" - - # __deepcopy__ - deep = deepcopy(mot) - assert deep._cancelled is True, "__deepcopy__ must propagate _cancelled" - - # _copy_from - target = ModelOutputThunk(value="original") - assert target._cancelled is False - target._copy_from(mot) - assert target._cancelled is True, "_copy_from must propagate _cancelled" - - # Sanity: default-constructed MOT has _cancelled=False. - fresh = ModelOutputThunk(value="x") - assert fresh._cancelled is False - - -# --------------------------------------------------------------------------- -# Fix 1 — setup-path backend leak: copy(req) before generate_from_context -# --------------------------------------------------------------------------- - - -class _PlainReq(Requirement): - """Default shallow copy — cannot raise.""" - - def format_for_llm(self) -> str: - return "plain" - - async def stream_validate( - self, chunk: str, *, backend: Any, ctx: Any - ) -> PartialValidationResult: - return PartialValidationResult("unknown") - - async def validate( - self, backend: Any, ctx: Any, *, format: Any = None, model_options: Any = None - ) -> ValidationResult: - return ValidationResult(result=True) - - -class _RaisingCopyReq(Requirement): - """__copy__ raises — simulates a user-defined Requirement with a faulty override.""" - - def __copy__(self) -> "_RaisingCopyReq": - raise ValueError("copy boom") - - def format_for_llm(self) -> str: - return "raising copy" - - async def stream_validate( - self, chunk: str, *, backend: Any, ctx: Any - ) -> PartialValidationResult: - return PartialValidationResult("unknown") - - async def validate( - self, backend: Any, ctx: Any, *, format: Any = None, model_options: Any = None - ) -> ValidationResult: - return ValidationResult(result=True) - - -class _InstrumentedBackend(StreamingMockBackend): - """Counts generate_from_context calls and exposes the last MOT produced.""" - - def __init__(self, response: str, token_size: int = 1) -> None: + def __init__( + self, response: str, token_size: int = 1, delay: float = 0.005 + ) -> None: super().__init__(response, token_size) - self.generate_from_context_call_count = 0 - self.last_mot: ModelOutputThunk | None = None + self._delay = delay async def _generate_from_context( self, action: Any, - ctx: Any, + ctx: Context, *, - format: Any = None, - model_options: dict | None = None, - tool_calls: bool = False, - ) -> tuple[ModelOutputThunk, Any]: - self.generate_from_context_call_count += 1 - mot, new_ctx = await super()._generate_from_context( - action, - ctx, - format=format, - model_options=model_options, - tool_calls=tool_calls, + format=None, + model_options=None, + tool_calls=False, + ) -> tuple[ModelOutputThunk, Context]: + _ = format, model_options, tool_calls + mot = _make_mot() + mot._gen.generate = asyncio.create_task( + _feed_tokens_slowly(mot, self._response, self._token_size, self._delay) ) - self.last_mot = mot + new_ctx = ctx.add(action).add(mot) return mot, new_ctx +@_cpex_skip @pytest.mark.asyncio -@pytest.mark.parametrize( - "req_cls,expect_raise", [(_PlainReq, False), (_RaisingCopyReq, True)] -) -async def test_stream_with_chunking_requirement_copy_contract( - req_cls: type, expect_raise: bool -) -> None: - """Fix 1: copy(req) runs before generate_from_context. +async def test_external_cancellation_mid_stream_still_finalizes() -> None: + """An external timeout mid-stream still fires terminal events and computes the MOT. - On __copy__ failure the backend is never started (call_count == 0). - On success the backend is called exactly once. + Wrapping consumption in `asyncio.wait_for` cancels the consuming task while a + chunk is in flight. The cancellation must still run cleanup: STREAMING_END and + GENERATION_ERROR fire once each, and the MOT ends computed + cancelled rather + than stranded, before the TimeoutError propagates. """ - backend = _InstrumentedBackend("Hello world. ", token_size=2) - req = req_cls() - if expect_raise: - with pytest.raises(ValueError, match="copy boom"): - await stream_with_chunking(_action(), backend, _ctx(), requirements=[req]) - # Hard invariant: reorder ensures backend never starts on copy failure. - assert backend.generate_from_context_call_count == 0 - assert backend.last_mot is None - else: - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=[req] - ) - await result.acomplete() - assert backend.generate_from_context_call_count == 1 - assert backend.last_mot is not None + response = "One. Two. Three. Four. Five. Six. Seven. Eight. " + backend = SlowStreamingMockBackend(response, token_size=2, delay=0.005) + ends: list[Any] = [] + gen_errors: list[Any] = [] -# --------------------------------------------------------------------------- -# Fix 3 — TaskGroup cancels peer validators on first failure -# --------------------------------------------------------------------------- -# Event type construction -# --------------------------------------------------------------------------- + @hook("streaming_end") + async def _end(payload: Any, ctx: Any) -> Any: + ends.append(payload) + return None + @hook("generation_error") + async def _gerr(payload: Any, ctx: Any) -> Any: + gen_errors.append(payload) + return None -def test_stream_event_types_have_auto_timestamp() -> None: - """All seven event types set timestamp automatically; callers do not pass it.""" - before = time.time() - all_events = [ - ChunkEvent(text="hello", chunk_index=0, attempt=1), - QuickCheckEvent( - chunk_index=0, - attempt=1, - passed=True, - results=[PartialValidationResult("unknown")], - ), - StreamingDoneEvent(attempt=1, full_text="hello"), - FullValidationEvent( - attempt=1, passed=True, results=[ValidationResult(result=True)] - ), - RetryEvent(attempt=2, reason="too long"), - CompletedEvent(success=True, full_text="hello", attempts_used=1), - ErrorEvent(exception_type="ValueError", detail="boom"), - ] - after = time.time() - - for ev in all_events: - assert isinstance(ev, StreamEvent) - assert before <= ev.timestamp <= after, ( - f"{type(ev).__name__} timestamp out of range" - ) + register(_end) + register(_gerr) + streamer: Streamer | None = None + try: + async def _run() -> None: + nonlocal streamer + async with await stream( + _action(), backend, _ctx(), chunking="sentence" + ) as s: + streamer = s + async for _chunk in s: + pass + + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(_run(), timeout=0.05) + finally: + unregister(_end) + unregister(_gerr) -# --------------------------------------------------------------------------- -# Event emission — happy path -# --------------------------------------------------------------------------- + assert len(ends) == 1 + assert ends[0].success is False + assert len(gen_errors) == 1 + assert streamer is not None + assert streamer._mot.is_computed() is True + assert streamer._mot.cancelled is True @pytest.mark.asyncio -async def test_event_emission_order_happy_path() -> None: - """Happy path: QuickCheckEvent/ChunkEvent pairs, then StreamingDoneEvent, - FullValidationEvent, CompletedEvent(success=True).""" - response = "First sentence. Second sentence. " - backend = StreamingMockBackend(response, token_size=4) - req = AlwaysUnknownReq() +async def test_early_exit_does_not_deadlock() -> None: + """A high-throughput stream that fails early must not hang. - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=[req], chunking="sentence" - ) - await result.acomplete() + The response is far longer than the MOT queue (maxsize 20), so an early + fail must not leave the producer blocked on a full queue after the consumer + stops. A hang trips the timeout. + """ + response = "word " * 200 + backend = StreamingMockBackend(response, token_size=5) + req = FailAfterWordsReq(threshold=3) - evts: list[StreamEvent] = [e async for e in result.events()] + streamer: Streamer | None = None - assert isinstance(evts[-1], CompletedEvent) - assert evts[-1].success is True - assert evts[-1].attempts_used == 1 + async def _run() -> None: + nonlocal streamer + async with await stream( + _action(), backend, _ctx(), requirements=[req], chunking="word" + ) as s: + streamer = s + async for _chunk in s: + pass - types = [type(e) for e in evts] - assert StreamingDoneEvent in types - assert types.index(StreamingDoneEvent) < types.index(CompletedEvent) - assert FullValidationEvent in types - assert types.index(FullValidationEvent) > types.index(StreamingDoneEvent) + await asyncio.wait_for(_run(), timeout=5.0) + assert streamer is not None + assert streamer.failed_early is True - chunk_events = [e for e in evts if isinstance(e, ChunkEvent)] - qc_events = [e for e in evts if isinstance(e, QuickCheckEvent)] - assert len(chunk_events) == 2 - assert len(qc_events) == 2 - assert [e.chunk_index for e in chunk_events] == [0, 1] - assert [e.chunk_index for e in qc_events] == [0, 1] - assert all(e.passed for e in qc_events) - # QuickCheckEvent fires before ChunkEvent within each pair: validation must - # complete before the chunk is released to the consumer queue. - for ci in range(2): - qc_pos = evts.index(qc_events[ci]) - ch_pos = evts.index(chunk_events[ci]) - assert qc_pos < ch_pos, f"chunk {ci}: QuickCheckEvent must precede ChunkEvent" +# --------------------------------------------------------------------------- +# Requirement cloning / backend routing +# --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_streaming_done_event_carries_full_text() -> None: - """StreamingDoneEvent.full_text matches full_text on the result.""" - response = "One sentence. Two sentences. " - backend = StreamingMockBackend(response, token_size=5) +async def test_clone_isolation_across_runs() -> None: + """Requirement instances are cloned per run; the original is never mutated.""" + req = FailAfterWordsReq(threshold=2) - result = await stream_with_chunking(_action(), backend, _ctx(), chunking="sentence") - await result.acomplete() + backend1 = StreamingMockBackend("one two three ", token_size=2) + async with await stream( + _action(), backend1, _ctx(), requirements=[req], chunking="word" + ) as s1: + async for _chunk in s1: + pass - evts = [e async for e in result.events()] - done_events = [e for e in evts if isinstance(e, StreamingDoneEvent)] - assert len(done_events) == 1 - assert done_events[0].full_text == result.full_text + # The original's running counter was not mutated by the first run. + assert req._word_count == 0 + backend2 = StreamingMockBackend("four five six ", token_size=2) + async with await stream( + _action(), backend2, _ctx(), requirements=[req], chunking="word" + ) as s2: + async for _chunk in s2: + pass -# --------------------------------------------------------------------------- -# Event emission — early exit -# --------------------------------------------------------------------------- + assert s1.failed_early is True + assert s2.failed_early is True @pytest.mark.asyncio -async def test_event_emission_on_early_exit() -> None: - """Early exit: QuickCheckEvent(passed=False) present; no StreamingDoneEvent - or FullValidationEvent; CompletedEvent(success=False).""" - response = "word " * 30 - backend = StreamingMockBackend(response, token_size=3) - req = FailAfterWordsReq(threshold=2) - - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=[req], chunking="word" - ) - await result.acomplete() - - evts = [e async for e in result.events()] - - assert isinstance(evts[-1], CompletedEvent) - assert evts[-1].success is False +async def test_validation_backend_routing() -> None: + """validation_backend, when given, receives stream_validate + validate calls.""" + gen_backend = StreamingMockBackend("Hello world. ", token_size=3) + val_backend = StreamingMockBackend("", token_size=1) + req = BackendRecordingReq() - types = [type(e) for e in evts] - assert FullValidationEvent not in types - assert StreamingDoneEvent not in types + captured: list[BackendRecordingReq] = [] + original_copy = BackendRecordingReq.__copy__ - fail_qc = [e for e in evts if isinstance(e, QuickCheckEvent) and not e.passed] - assert len(fail_qc) >= 1 + def _capturing_copy(self: BackendRecordingReq) -> BackendRecordingReq: + clone = original_copy(self) + captured.append(clone) + return clone + BackendRecordingReq.__copy__ = _capturing_copy # type: ignore[method-assign] + try: + async with await stream( + _action(), + gen_backend, + _ctx(), + requirements=[req], + chunking="sentence", + validation_backend=val_backend, + ) as streamer: + async for _chunk in streamer: + pass + finally: + BackendRecordingReq.__copy__ = original_copy # type: ignore[method-assign] -# --------------------------------------------------------------------------- -# Event emission — exception path -# --------------------------------------------------------------------------- + assert streamer.failed_early is False + # The original requirement was never called — only its per-run clone. + assert req.seen_backends == [] + # Every recorded backend was the validation backend, not the generation one. + assert len(captured) == 1 + assert captured[0].seen_backends + assert all(b is val_backend for b in captured[0].seen_backends) @pytest.mark.asyncio -async def test_stream_with_chunking_cancels_peer_validators() -> None: - """Fix 3: a failing stream_validate causes TaskGroup to cancel peer validators. +async def test_requirement_copy_contract() -> None: + """A raising __copy__ propagates from stream() before generation starts.""" - One requirement raises immediately in stream_validate; the second sleeps - for 5 s and sets a flag on completion. Without TaskGroup the slow sibling - runs detached; with it the cancellation is observed. - """ - reached_final_stage = asyncio.Event() + class RaisingCopyReq(Requirement): + def __copy__(self) -> "RaisingCopyReq": + raise RuntimeError("copy boom") - class _RaisingReq(Requirement): def format_for_llm(self) -> str: - return "raiser" + return "raising copy" async def stream_validate( self, chunk: str, *, backend: Any, ctx: Any ) -> PartialValidationResult: - raise RuntimeError("validator failed") + return PartialValidationResult("unknown") async def validate( self, @@ -1437,21 +1075,50 @@ async def validate( format: Any = None, model_options: Any = None, ) -> ValidationResult: - return ValidationResult(result=False) + return ValidationResult(result=True) - class _SlowReq(Requirement): + class CountingBackend(StreamingMockBackend): + def __init__(self, response: str, token_size: int = 1) -> None: + super().__init__(response, token_size) + self.gen_calls = 0 + + async def _generate_from_context(self, *args: Any, **kwargs: Any): + self.gen_calls += 1 + return await super()._generate_from_context(*args, **kwargs) + + # Failure path: the copy fails, so generation is never started. + fail_backend = CountingBackend("Hello. ", token_size=3) + with pytest.raises(RuntimeError, match="copy boom"): + await stream(_action(), fail_backend, _ctx(), requirements=[RaisingCopyReq()]) + assert fail_backend.gen_calls == 0 + + # Success path: a good copy starts generation exactly once. + ok_backend = CountingBackend("Hello. ", token_size=3) + async with await stream( + _action(), ok_backend, _ctx(), requirements=[AlwaysUnknownReq()] + ) as streamer: + async for _chunk in streamer: + pass + assert ok_backend.gen_calls == 1 + + +# --------------------------------------------------------------------------- +# Error + precomputed-MOT handling +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_exception_in_stream_validate_propagates_and_cancels() -> None: + """An exception in stream_validate propagates from the loop and cancels gen.""" + + class RaisingReq(Requirement): def format_for_llm(self) -> str: - return "slow" + return "raiser" async def stream_validate( self, chunk: str, *, backend: Any, ctx: Any ) -> PartialValidationResult: - try: - await asyncio.sleep(5.0) - reached_final_stage.set() - return PartialValidationResult("pass") - except asyncio.CancelledError: - raise # propagate so TaskGroup knows we were cancelled + raise RuntimeError("validate boom") async def validate( self, @@ -1463,27 +1130,28 @@ async def validate( ) -> ValidationResult: return ValidationResult(result=True) - backend = StreamingMockBackend("Hello world. ", token_size=2) - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=[_RaisingReq(), _SlowReq()] - ) - with pytest.raises(RuntimeError, match="validator failed"): - await result.acomplete() + backend = StreamingMockBackend("Hello world. ", token_size=3) - # Give the loop a tick; the slow sibling must NOT have run to completion. - await asyncio.sleep(0.05) - assert not reached_final_stage.is_set(), ( - "slow sibling was not cancelled by TaskGroup" + streamer = await stream( + _action(), backend, _ctx(), requirements=[RaisingReq()], chunking="sentence" ) + with pytest.raises(RuntimeError, match="validate boom"): + async with streamer: + async for _chunk in streamer: + pass + + # The generation was cancelled during teardown, not merely finished. + assert streamer._mot._cancelled is True + assert streamer._mot.is_computed() is True @pytest.mark.asyncio -async def test_stream_with_chunking_rejects_precomputed_mot() -> None: - """Backend returning an already-computed MOT raises RuntimeError immediately. +async def test_rejects_precomputed_mot() -> None: + """A backend returning an already-computed MOT raises RuntimeError. - stream_with_chunking() requires streaming; a pre-computed MOT would cause - the orchestrator loop to skip entirely, producing empty output and silently - passing all final validators against an empty string. + stream() requires streaming; a pre-computed MOT would skip the loop entirely, + producing empty output and silently passing final validators against an empty + string. """ class PrecomputedBackend(Backend): @@ -1507,21 +1175,30 @@ async def _generate_from_raw( raise NotImplementedError with pytest.raises(RuntimeError, match="already-computed MOT"): - await stream_with_chunking(_action(), PrecomputedBackend(), _ctx()) + await stream(_action(), PrecomputedBackend(), _ctx()) @pytest.mark.asyncio -async def test_error_event_on_stream_validate_exception() -> None: - """When stream_validate raises, ErrorEvent is emitted and CompletedEvent follows.""" +async def test_unknown_chunking_alias_raises_value_error() -> None: + """An unknown chunking alias string raises ValueError.""" + backend = StreamingMockBackend("Hello. ", token_size=3) + with pytest.raises(ValueError, match="Unknown chunking alias"): + await stream(_action(), backend, _ctx(), chunking="unknown_alias") + - class RaisingReq2(Requirement): +@pytest.mark.asyncio +async def test_cancels_peer_validators() -> None: + """A failing stream_validate does not let a slow peer run to completion.""" + reached_final_stage = asyncio.Event() + + class _RaisingReq(Requirement): def format_for_llm(self) -> str: - return "raises" + return "raiser" async def stream_validate( self, chunk: str, *, backend: Any, ctx: Any ) -> PartialValidationResult: - raise RuntimeError("test-error") + raise RuntimeError("validator failed") async def validate( self, @@ -1531,164 +1208,229 @@ async def validate( format: Any = None, model_options: Any = None, ) -> ValidationResult: - return ValidationResult(result=True) + return ValidationResult(result=False) - backend = StreamingMockBackend("hello world", token_size=5) - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=[RaisingReq2()], chunking="word" - ) - with pytest.raises(RuntimeError, match="test-error"): - async for _c in result.astream(): - pass - await asyncio.wait_for(result.acomplete(), timeout=5.0) + class _SlowReq(Requirement): + def format_for_llm(self) -> str: + return "slow" - evts = [e async for e in result.events()] + async def stream_validate( + self, chunk: str, *, backend: Any, ctx: Any + ) -> PartialValidationResult: + await asyncio.sleep(5.0) + reached_final_stage.set() + return PartialValidationResult("pass") - error_events = [e for e in evts if isinstance(e, ErrorEvent)] - assert len(error_events) == 1 - assert error_events[0].exception_type == "RuntimeError" - assert "test-error" in error_events[0].detail + async def validate( + self, + backend: Any, + ctx: Any, + *, + format: Any = None, + model_options: Any = None, + ) -> ValidationResult: + return ValidationResult(result=True) + + backend = StreamingMockBackend("Hello world. ", token_size=2) + streamer = await stream( + _action(), backend, _ctx(), requirements=[_RaisingReq(), _SlowReq()] + ) + with pytest.raises(RuntimeError, match="validator failed"): + async with streamer: + async for _chunk in streamer: + pass - assert isinstance(evts[-1], CompletedEvent) - assert evts[-1].success is False + await asyncio.sleep(0.05) + assert not reached_final_stage.is_set(), "slow sibling ran to completion" # --------------------------------------------------------------------------- -# Concurrent astream() + events() +# Event emission (observed via the STREAMING_EVENT hook) # --------------------------------------------------------------------------- +def test_stream_event_types_have_auto_timestamp() -> None: + """All seven event types set timestamp automatically; callers do not pass it.""" + before = time.time() + all_events = [ + ChunkEvent(text="hello", chunk_index=0, attempt=1), + QuickCheckEvent( + chunk_index=0, + attempt=1, + passed=True, + results=[PartialValidationResult("unknown")], + ), + StreamingDoneEvent(attempt=1, full_text="hello"), + FullValidationEvent( + attempt=1, passed=True, results=[ValidationResult(result=True)] + ), + RetryEvent(attempt=2, reason="too long"), + CompletedEvent(success=True, full_text="hello", attempts_used=1), + ErrorEvent(exception_type="ValueError", detail="boom"), + ] + after = time.time() + + for ev in all_events: + assert isinstance(ev, StreamEvent) + assert before <= ev.timestamp <= after, ( + f"{type(ev).__name__} timestamp out of range" + ) + + +@_cpex_skip @pytest.mark.asyncio -async def test_concurrent_astream_and_events() -> None: - """astream() and events() can be consumed concurrently without interference.""" - response = "Alpha. Beta. Gamma. " - backend = StreamingMockBackend(response, token_size=4) - req = AlwaysUnknownReq() +async def test_event_emission_order_happy_path() -> None: + """Natural completion emits QuickCheck/Chunk pairs, then Done, FullValidation, Completed.""" + response = "One. Two. " + backend = StreamingMockBackend(response, token_size=2) - result = await stream_with_chunking( - _action(), backend, _ctx(), requirements=[req], chunking="sentence" - ) + with _record_events() as events: + async with await stream( + _action(), + backend, + _ctx(), + requirements=[AlwaysUnknownReq()], + chunking="sentence", + ) as streamer: + async for _chunk in streamer: + pass - async def drain_chunks() -> list[str]: - return [c async for c in result.astream()] + types = [type(e) for e in events] + assert types[-1] is CompletedEvent + assert events[-1].success is True + assert events[-1].attempts_used == 1 + assert StreamingDoneEvent in types + assert FullValidationEvent in types + # Done precedes FullValidation precedes Completed. + assert types.index(StreamingDoneEvent) < types.index(FullValidationEvent) + assert types.index(FullValidationEvent) < types.index(CompletedEvent) - async def drain_events() -> list[StreamEvent]: - return [e async for e in result.events()] + # Two sentences → two QuickCheck/Chunk pairs, indexed in order. + chunk_events = [e for e in events if isinstance(e, ChunkEvent)] + qc_events = [e for e in events if isinstance(e, QuickCheckEvent)] + assert len(chunk_events) == 2 + assert len(qc_events) == 2 + assert [e.chunk_index for e in chunk_events] == [0, 1] + assert [e.chunk_index for e in qc_events] == [0, 1] + assert all(e.passed for e in qc_events) - chunks, evts = await asyncio.gather(drain_chunks(), drain_events()) - await result.acomplete() + # QuickCheckEvent precedes ChunkEvent within each pair: a chunk is validated + # before it is emitted. + for ci in range(2): + assert events.index(qc_events[ci]) < events.index(chunk_events[ci]) - assert len(chunks) == 3 - assert isinstance(evts[-1], CompletedEvent) - assert evts[-1].success is True - chunk_evts = [e for e in evts if isinstance(e, ChunkEvent)] - assert [e.chunk_index for e in chunk_evts] == list(range(len(chunks))) +@_cpex_skip +@pytest.mark.asyncio +async def test_streaming_done_event_carries_full_text() -> None: + """StreamingDoneEvent.full_text matches the streamer's full_text.""" + response = "One. Two. Three. " + backend = StreamingMockBackend(response, token_size=3) + with _record_events() as events: + async with await stream( + _action(), backend, _ctx(), chunking="sentence" + ) as streamer: + async for _chunk in streamer: + pass -# --------------------------------------------------------------------------- -# events() single-consumer guard -# --------------------------------------------------------------------------- + done = [e for e in events if isinstance(e, StreamingDoneEvent)] + assert len(done) == 1 + assert done[0].full_text == streamer.full_text +@_cpex_skip @pytest.mark.asyncio -async def test_events_single_consumer_guard_raises_on_second_call() -> None: - """events() raises RuntimeError if called a second time on the same result.""" - response = "One sentence. " - backend = StreamingMockBackend(response, token_size=4) +async def test_event_emission_on_early_exit() -> None: + """Early exit emits a failing QuickCheck then Completed; no Done/FullValidation.""" + response = "one two three four five " + backend = StreamingMockBackend(response, token_size=2) + req = FailAfterWordsReq(threshold=2) - result = await stream_with_chunking(_action(), backend, _ctx(), chunking="sentence") - await result.acomplete() + with _record_events() as events: + async with await stream( + _action(), backend, _ctx(), requirements=[req], chunking="word" + ) as streamer: + async for _chunk in streamer: + pass - # First drain — OK. - async for _ in result.events(): - pass + types = [type(e) for e in events] + assert StreamingDoneEvent not in types + assert FullValidationEvent not in types + assert types[-1] is CompletedEvent + assert events[-1].success is False + # The last quick-check failed. + quick = [e for e in events if isinstance(e, QuickCheckEvent)] + assert quick[-1].passed is False - # Second call must raise immediately. - with pytest.raises(RuntimeError, match="single-consumer"): - async for _ in result.events(): - pass +@_cpex_skip +@pytest.mark.asyncio +async def test_no_requirements_omits_full_validation_event() -> None: + """With no requirements, no QuickCheck or FullValidation events fire.""" + response = "One. Two. " + backend = StreamingMockBackend(response, token_size=2) -# --------------------------------------------------------------------------- -# CancelledError path -# --------------------------------------------------------------------------- + with _record_events() as events: + async with await stream( + _action(), backend, _ctx(), chunking="sentence" + ) as streamer: + async for _chunk in streamer: + pass + types = [type(e) for e in events] + assert QuickCheckEvent not in types + assert FullValidationEvent not in types + assert ChunkEvent in types + assert StreamingDoneEvent in types + assert types[-1] is CompletedEvent + assert events[-1].success is True -@pytest.mark.asyncio -async def test_cancelled_task_sets_completed_false() -> None: - """External task cancellation must leave result.completed=False. - - CancelledError is a BaseException and bypasses except Exception, so - the finally block is responsible for setting result.completed=False. - Regression: without the fix, result.completed stays True and - CompletedEvent / record_sampling_outcome lie to callers. - - Uses a backend whose token feed blocks on an asyncio.Event that is - never set, guaranteeing the orchestrator is suspended at astream() - when the task is cancelled. - - Requires `await asyncio.sleep(0)` before `cancel()` — see inline - comment. Python 3.12's C Task implementation skips the coroutine body - entirely (including finally blocks) when cancelled before the first - `coro.send(None)`. - """ - gate = asyncio.Event() # never set — feed task blocks indefinitely - feed_task: asyncio.Task[None] | None = None - async def _blocking_feed(mot: ModelOutputThunk) -> None: - await gate.wait() +@_cpex_skip +@pytest.mark.asyncio +async def test_error_event_on_stream_validate_exception() -> None: + """An exception in stream_validate emits ErrorEvent then Completed(success=False).""" - class BlockingBackend(Backend): - _model_id: str = "blocking-mock-model" - _provider: str = "blocking-mock-provider" + class RaisingReq(Requirement): + def format_for_llm(self) -> str: + return "raiser" - async def _generate_from_context( - self, action: Any, ctx: Any, **kwargs: Any - ) -> tuple[ModelOutputThunk, Any]: - nonlocal feed_task - mot = _make_mot() - feed_task = asyncio.create_task(_blocking_feed(mot)) - return mot, ctx.add(action).add(mot) + async def stream_validate( + self, chunk: str, *, backend: Any, ctx: Any + ) -> PartialValidationResult: + raise RuntimeError("boom") - async def _generate_from_raw(self, *args: Any, **kwargs: Any) -> Any: - raise NotImplementedError + async def validate( + self, + backend: Any, + ctx: Any, + *, + format: Any = None, + model_options: Any = None, + ) -> ValidationResult: + return ValidationResult(result=True) - result = await stream_with_chunking( - _action(), BlockingBackend(), _ctx(), chunking="word" - ) - assert result._orchestration_task is not None + backend = StreamingMockBackend("Hello world. ", token_size=3) - # Deliberately uses sleep(0) rather than _orchestration_started.wait(): - # BlockingBackend blocks indefinitely, so there is no race with the stream - # completing before cancel(). sleep(0) is sufficient to satisfy the - # coro.send(None) requirement: Python 3.12's C Task implementation skips - # the coroutine body entirely (including finally blocks) when cancelled - # before the first send. - await asyncio.sleep(0) + with _record_events() as events: + streamer = await stream( + _action(), backend, _ctx(), requirements=[RaisingReq()], chunking="sentence" + ) + with pytest.raises(RuntimeError, match="boom"): + async with streamer: + async for _chunk in streamer: + pass + + types = [type(e) for e in events] + assert types[-1] is CompletedEvent + assert events[-1].success is False + # Exactly one ErrorEvent, carrying the exception type and message. + error_events = [e for e in events if isinstance(e, ErrorEvent)] + assert len(error_events) == 1 + assert error_events[0].exception_type == "RuntimeError" + assert "boom" in error_events[0].detail - result._orchestration_task.cancel() - try: - await result._orchestration_task - except BaseException: - pass - - # Primary assertion: completed must be False after external cancellation. - assert result.completed is False - - # The finally block must have run to completion: _done must be set and - # acomplete() must not hang. This is the actual failure mode the fix - # guards against — if _done is never set, acomplete() blocks forever. - # External cancellation surfaces as CancelledError (raise-once contract). - assert result._done.is_set() - with pytest.raises(asyncio.CancelledError): - await asyncio.wait_for(result.acomplete(), timeout=2.0) - - # Clean up the blocking feed task to avoid "Task destroyed while pending". - if feed_task is not None: - feed_task.cancel() - try: - await feed_task - except BaseException: - pass +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/test/telemetry/conftest.py b/test/telemetry/conftest.py index 49560e1aeb..e90093c427 100644 --- a/test/telemetry/conftest.py +++ b/test/telemetry/conftest.py @@ -61,7 +61,6 @@ def reset_tracing_state() -> None: tracing._application_tracer = None tracing._backend_tracer = None tracing._in_flight_spans.clear() - tracing._reattached_tokens.clear() # Re-register: another test's shutdown_plugins() may have emptied the manager. tracing._plugins_registered = False tracing._setup_tracing() diff --git a/test/telemetry/test_metrics_plugins.py b/test/telemetry/test_metrics_plugins.py index 5d315966b3..a952c22070 100644 --- a/test/telemetry/test_metrics_plugins.py +++ b/test/telemetry/test_metrics_plugins.py @@ -963,24 +963,24 @@ async def test_sampling_plugin_skips_outcome_on_exception(sampling_plugin): @pytest.mark.asyncio async def test_sampling_plugin_records_streaming_success_outcome(sampling_plugin): - """streaming_end with success=True records a `stream_with_chunking` success.""" + """streaming_end with success=True records a `stream` success.""" payload = StreamingEndPayload(streaming_id="sid", success=True) with patch("mellea.telemetry.metrics.record_sampling_outcome") as mock_record: await sampling_plugin.record_streaming_outcome(payload, {}) - mock_record.assert_called_once_with("stream_with_chunking", True) + mock_record.assert_called_once_with("stream", True) @pytest.mark.asyncio async def test_sampling_plugin_records_streaming_failure_outcome(sampling_plugin): - """streaming_end with success=False records a `stream_with_chunking` failure.""" + """streaming_end with success=False records a `stream` failure.""" payload = StreamingEndPayload(streaming_id="sid", success=False) with patch("mellea.telemetry.metrics.record_sampling_outcome") as mock_record: await sampling_plugin.record_streaming_outcome(payload, {}) - mock_record.assert_called_once_with("stream_with_chunking", False) + mock_record.assert_called_once_with("stream", False) # RequirementMetricsPlugin tests diff --git a/test/telemetry/test_tracing_backend.py b/test/telemetry/test_tracing_backend.py index 8c324150f2..2e98e710de 100644 --- a/test/telemetry/test_tracing_backend.py +++ b/test/telemetry/test_tracing_backend.py @@ -10,7 +10,6 @@ import pytest from mellea.backends.model_ids import IBM_GRANITE_4_1_3B -from mellea.backends.model_options import ModelOption from mellea.backends.ollama import OllamaModelBackend from mellea.plugins.manager import ( disable_background_collection, @@ -115,20 +114,23 @@ async def fake_chat(*args, **kwargs): @pytest.mark.integration @pytest.mark.asyncio @pytest.mark.parametrize("emit", [True, False], ids=["emit_on", "emit_off_default"]) -async def test_streaming_span_creates_and_closes_span(span_exporter, monkeypatch, emit): - """Streaming backend call creates a chat span that closes after the stream completes. - - Uses a mocked Ollama client so no server is needed. Verifies the core - TracingPlugin invariant: the span must remain open for the full duration of - streaming and close only once all chunks are consumed. `chunk_processed` - events are emitted only when `MELLEA_GENERATION_CHUNK_EVENTS` is on; with the env - unset (the default) the span carries no such events. +async def test_stream_mocked(span_exporter, monkeypatch, emit): + """Mocked streaming run: the `stream` span roots the backend `chat` span. + + Uses a mocked Ollama client so no server is needed; the `stream` span is the + root, the `chat` span nests under it, and stays open for the full stream. + `chunk_processed` events are emitted on `chat` only when + `MELLEA_GENERATION_CHUNK_EVENTS` is on; with the env unset (the default) the span + carries no such events. """ if emit: monkeypatch.setenv("MELLEA_GENERATION_CHUNK_EVENTS", "true") else: monkeypatch.delenv("MELLEA_GENERATION_CHUNK_EVENTS", raising=False) + from mellea.stdlib.streaming import stream + from mellea.telemetry.tracing_plugins import _CONTEXT_ATTACH_SUPPORTED + async def fake_chat_stream(*args, **kwargs): for content in ["1", " 2", " 3"]: await asyncio.sleep(0.05) @@ -160,31 +162,33 @@ async def fake_chat_stream(*args, **kwargs): backend = OllamaModelBackend(model_id="test-model") ctx = SimpleContext().add(Message(role="user", content="Count to 3")) - mot, _ = await backend.generate_from_context( - Message(role="assistant", content=""), - ctx, - model_options={ModelOption.STREAM: True}, - ) - - await mot.astream() - await mot.avalue() + async with await stream( + Message(role="assistant", content=""), backend, ctx + ) as streamer: + async for _ in streamer: + pass await drain_background_tasks() trace.get_tracer_provider().force_flush() spans = span_exporter.get_finished_spans() - backend_span = next((s for s in spans if s.name == "chat test-model"), None) + streaming_span = next(s for s in spans if s.name == "stream") + chat_span = next(s for s in spans if s.name == "chat test-model") - assert backend_span is not None, "Backend span not found" - assert backend_span.end_time > backend_span.start_time, ( - "Span must have nonzero duration" - ) + assert streaming_span.parent is None, "streaming span should be a root" + if _CONTEXT_ATTACH_SUPPORTED: + assert chat_span.parent is not None + assert chat_span.parent.span_id == streaming_span.context.span_id, ( + "chat span should nest under stream" + ) + else: + assert chat_span.parent is None, "chat span should be flat on Python <=3.11" - span_duration_s = (backend_span.end_time - backend_span.start_time) / 1e9 # fake stream is 3 chunks x 50 ms ~= 150 ms; >= 0.1 confirms span survived past first chunk - assert span_duration_s >= 0.1, ( - f"Span closed too early — duration {span_duration_s:.3f}s is shorter than " - "the streaming delay, suggesting the span did not stay open for the full stream" + chat_duration_s = (chat_span.end_time - chat_span.start_time) / 1e9 + assert chat_duration_s >= 0.1, ( + f"chat span closed too early — duration {chat_duration_s:.3f}s is shorter " + "than the streaming delay, suggesting it did not stay open for the full stream" ) chunk_events = [ @@ -192,7 +196,7 @@ async def fake_chat_stream(*args, **kwargs): e.attributes.get("mellea.generation.chunk_index"), e.attributes.get("mellea.generation.chunk_text_length"), ) - for e in backend_span.events + for e in chat_span.events if e.name == "chunk_processed" ] if emit: @@ -590,42 +594,41 @@ async def test_multiple_generations_separate_spans(span_exporter): @pytest.mark.ollama @pytest.mark.slow @pytest.mark.asyncio -async def test_stream_with_chunking_e2e(span_exporter): - """A real `stream_with_chunking` run wraps the backend `chat` span and times it. +async def test_stream_e2e(span_exporter): + """A real `stream` run wraps the backend `chat` span and times it. - Covers the streaming path end to end against a live model: the - `stream_with_chunking` span is the root, the backend `chat` span nests under - it, and that `chat` span stays open across the full stream. + Covers the streaming path end to end against a live model: the `stream` span + is the root, the backend `chat` span nests under it, and that `chat` span + stays open across the full stream. """ - from mellea.stdlib.streaming import stream_with_chunking + from mellea.stdlib.streaming import stream backend = OllamaModelBackend(model_id=IBM_GRANITE_4_1_3B.ollama_name) # type: ignore ctx = SimpleContext().add(Message(role="user", content="Count to 3")) - result = await stream_with_chunking( + async with await stream( Message(role="assistant", content=""), backend, ctx - ) - async for _ in result.astream(): - pass - await result.acomplete() + ) as streamer: + async for _ in streamer: + pass await drain_background_tasks() from mellea.telemetry.tracing_plugins import _CONTEXT_ATTACH_SUPPORTED spans = span_exporter.get_finished_spans() - streaming_span = next(s for s in spans if s.name == "stream_with_chunking") + streaming_span = next(s for s in spans if s.name == "stream") chat_span = next( s for s in spans if s.name == f"chat {IBM_GRANITE_4_1_3B.ollama_name}" ) assert streaming_span.parent is None, "streaming span should be a root" - if not _CONTEXT_ATTACH_SUPPORTED: + if _CONTEXT_ATTACH_SUPPORTED: + assert chat_span.parent is not None + assert chat_span.parent.span_id == streaming_span.context.span_id, ( + "chat span should nest under stream" + ) + else: assert chat_span.parent is None, "chat span should be flat on Python <=3.11" - return - assert chat_span.parent is not None - assert chat_span.parent.span_id == streaming_span.context.span_id, ( - "chat span should nest under stream_with_chunking" - ) chat_duration_s = (chat_span.end_time - chat_span.start_time) / 1e9 assert chat_duration_s >= 0.1, ( diff --git a/test/telemetry/test_tracing_plugins.py b/test/telemetry/test_tracing_plugins.py index f1dbb8886e..2dead0b471 100644 --- a/test/telemetry/test_tracing_plugins.py +++ b/test/telemetry/test_tracing_plugins.py @@ -40,6 +40,7 @@ from mellea.plugins.hooks.tool import ToolPostInvokePayload, ToolPreInvokePayload from mellea.stdlib.streaming import ( ChunkEvent, + CompletedEvent, ErrorEvent, FullValidationEvent, QuickCheckEvent, @@ -943,7 +944,7 @@ async def test_streaming_start_starts_span_and_stashes_by_streaming_id( streaming_id="sid-1", has_requirements=True, requirement_count=2, - chunking_strategy="SentenceChunker", + chunking_strategy="SentenceChunking", ) with patch( @@ -951,21 +952,19 @@ async def test_streaming_start_starts_span_and_stashes_by_streaming_id( ): await streaming_plugin.on_streaming_start(payload, {}) - fake_tracer.start_span.assert_called_once_with("stream_with_chunking") + fake_tracer.start_span.assert_called_once_with("stream") assert "sid-1" in tracing._in_flight_spans fake_span.end.assert_not_called() attrs = _attrs(fake_span) assert attrs["mellea.streaming.has_requirements"] is True assert attrs["mellea.streaming.requirement_count"] == 2 - assert attrs["mellea.streaming.chunking_strategy"] == "SentenceChunker" + assert attrs["mellea.streaming.chunking_strategy"] == "SentenceChunking" # The correlation id is the in-flight key, not a span attribute. assert "mellea.streaming_id" not in attrs @pytest.mark.asyncio -async def test_streaming_end_records_completed_event_then_closes_span( - streaming_plugin, enabled_tracing -): +async def test_streaming_end_success_closes_span(streaming_plugin, enabled_tracing): fake_span = MagicMock() fake_tracer = MagicMock() fake_tracer.start_span.return_value = fake_span @@ -981,11 +980,6 @@ async def test_streaming_end_records_completed_event_then_closes_span( await streaming_plugin.on_streaming_end(end, {}) fake_span.end.assert_called_once() - events = _events(fake_span) - assert any(name == "completed" for name, _ in events) - completed_attrs = next(attrs for name, attrs in events if name == "completed") - assert completed_attrs["mellea.streaming.success"] is True - assert completed_attrs["mellea.streaming.full_text_length"] == 11 attrs = _attrs(fake_span) assert attrs["mellea.streaming.full_text_length"] == 11 @@ -1055,21 +1049,31 @@ async def test_streaming_event_records_mid_stream_events( full_val = FullValidationEvent( attempt=1, passed=True, results=[ValidationResult(result=True)] ) + completed = CompletedEvent(success=True, full_text="hello world", attempts_used=1) - for ev in (qc, chunk, done, full_val): + for ev in (qc, chunk, done, full_val, completed): await streaming_plugin.on_streaming_event( StreamingEventPayload(streaming_id="sid-ev", event=ev), {} ) events = _events(fake_span) names = [name for name, _ in events] - assert names == ["quick_check", "chunk", "streaming_done", "full_validation"] + assert names == [ + "quick_check", + "chunk", + "streaming_done", + "full_validation", + "completed", + ] qc_attrs = events[0][1] assert qc_attrs["mellea.streaming.chunk_index"] == 0 assert qc_attrs["mellea.validation.passed"] is True assert qc_attrs["mellea.validation.requirement_count"] == 1 chunk_attrs = events[1][1] assert chunk_attrs["mellea.streaming.chunk_text_length"] == 5 + completed_attrs = events[4][1] + assert completed_attrs["mellea.streaming.success"] is True + assert completed_attrs["mellea.streaming.full_text_length"] == 11 # streaming_event never closes the span. fake_span.end.assert_not_called() diff --git a/test/telemetry/test_tracing_streaming.py b/test/telemetry/test_tracing_streaming.py index 7735b4fa47..9ed102d1ec 100644 --- a/test/telemetry/test_tracing_streaming.py +++ b/test/telemetry/test_tracing_streaming.py @@ -1,7 +1,7 @@ # Copyright IBM Corp. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -"""Tests for the `stream_with_chunking` tracing span.""" +"""Tests for the `stream` tracing span.""" import asyncio from unittest.mock import MagicMock, patch @@ -61,15 +61,15 @@ def test_start_streaming_span_stamps_attrs_and_stashes_under_id(enabled_tracing) "sid-1", has_requirements=True, requirement_count=2, - chunking_strategy="SentenceChunker", + chunking_strategy="SentenceChunking", ) - fake_tracer.start_span.assert_called_once_with("stream_with_chunking") + fake_tracer.start_span.assert_called_once_with("stream") assert "sid-1" in tracing._in_flight_spans attrs = _attrs(fake_span) assert attrs["mellea.streaming.has_requirements"] is True assert attrs["mellea.streaming.requirement_count"] == 2 - assert attrs["mellea.streaming.chunking_strategy"] == "SentenceChunker" + assert attrs["mellea.streaming.chunking_strategy"] == "SentenceChunking" # The correlation id is the in-flight key, not a span attribute. assert "mellea.streaming_id" not in attrs @@ -142,203 +142,6 @@ def test_streaming_span_helpers_silent_when_tracing_disabled(disabled_tracing): finish_streaming_span("sid-d", success=True) # should not raise -def test_reattach_span_attaches_and_releases(enabled_tracing): - from mellea.telemetry.tracing import reattach_span, release_reattached_span - - _, fake_tracer = _patch_app_tracer() - with patch( - "mellea.telemetry.tracing.get_application_tracer", return_value=fake_tracer - ): - start_streaming_span( - "sid-r", has_requirements=False, requirement_count=0, chunking_strategy="x" - ) - reattach_span("sid-r") - # A token is held while the span is re-attached. - assert "sid-r" in tracing._reattached_tokens - release_reattached_span("sid-r") - # The token is gone after release. - assert "sid-r" not in tracing._reattached_tokens - - -def test_reattach_span_noop_when_not_in_flight(enabled_tracing): - from mellea.telemetry.tracing import reattach_span, release_reattached_span - - # No matching in-flight span — both calls must be silent no-ops. - reattach_span("missing") - release_reattached_span("missing") - assert not tracing._reattached_tokens - - -@pytest.mark.asyncio -async def test_cross_task_detach_outside_reattached_scope_warns_and_runs( - enabled_tracing, caplog -): - """A cross-task detach with no reattached scope warns and still runs the detach. - - Without a reattach scope on the current task the mismatch is unexpected, so - the detach is left to run (OTel logs its own ERROR) after a mellea warning. - """ - import logging - - from mellea.telemetry.tracing import finish_backend_span_success, start_backend_span - - fake_span = MagicMock() - fake_tracer = MagicMock() - fake_tracer.start_span.return_value = fake_span - - with ( - patch("mellea.telemetry.tracing.get_backend_tracer", return_value=fake_tracer), - patch("mellea.telemetry.tracing.otel_context.detach") as detach, - ): - # Attach on this (caller) task. - start_backend_span("chat", "gid-x", model="m", provider="p") - - async def _finish_on_other_task() -> None: - finish_backend_span_success( - "gid-x", operation="chat", usage=None, mot=None, gen=None - ) - - with caplog.at_level(logging.WARNING, logger="mellea"): - await asyncio.create_task(_finish_on_other_task()) - - detach.assert_called_once() - fake_span.end.assert_called_once() - assert "gid-x" not in tracing._in_flight_spans - warnings = [r for r in caplog.records if r.levelno == logging.WARNING] - assert any("across asyncio tasks" in r.getMessage() for r in warnings), ( - "expected a warning naming the cross-task detach" - ) - - -@pytest.mark.asyncio -async def test_cross_task_detach_inside_reattached_scope_is_debug( - enabled_tracing, caplog -): - """A cross-task detach inside a reattached scope is skipped quietly at debug. - - The `stream_with_chunking` case: the orchestration task re-attaches the - streaming span, then finishes the caller-attached `chat` span. The chat - token's doomed detach is skipped (only the reattach token is detached, by - release) and no warning is logged. - """ - import logging - - from mellea.telemetry.tracing import ( - finish_backend_span_success, - reattach_span, - release_reattached_span, - start_backend_span, - start_streaming_span, - ) - - fake_span = MagicMock() - fake_tracer = MagicMock() - fake_tracer.start_span.return_value = fake_span - - with ( - patch("mellea.telemetry.tracing.get_backend_tracer", return_value=fake_tracer), - patch( - "mellea.telemetry.tracing.get_application_tracer", return_value=fake_tracer - ), - ): - start_streaming_span( - "sid-x", has_requirements=False, requirement_count=0, chunking_strategy="x" - ) - # Attach the chat span on this (caller) task; capture its doomed token. - start_backend_span("chat", "gid-x", model="m", provider="p") - chat_token = tracing._in_flight_spans["gid-x"][1] - - async def _finish_on_other_task() -> None: - reattach_span("sid-x") - try: - with patch("mellea.telemetry.tracing.otel_context.detach") as detach: - finish_backend_span_success( - "gid-x", operation="chat", usage=None, mot=None, gen=None - ) - # The chat token's cross-task detach was skipped entirely. - assert chat_token not in [c.args[0] for c in detach.call_args_list] - finally: - release_reattached_span("sid-x") - - with caplog.at_level(logging.DEBUG, logger="mellea"): - await asyncio.create_task(_finish_on_other_task()) - - warnings = [r for r in caplog.records if r.levelno == logging.WARNING] - assert not warnings, f"unexpected warning(s): {[r.getMessage() for r in warnings]}" - debug = [r for r in caplog.records if r.levelno == logging.DEBUG] - assert any("reattached-span scope" in r.getMessage() for r in debug), ( - "expected a debug line for the expected cross-task detach" - ) - - -@pytest.mark.asyncio -async def test_cross_task_detach_warns_when_scope_belongs_to_another_task( - enabled_tracing, caplog -): - """A reattach scope on a different task does not mark this task's detach expected. - - Guards the per-task classifier: a concurrent run's open scope (task A) must - not silence an unrelated cross-task detach finishing on task B. - """ - import logging - - from mellea.telemetry.tracing import ( - finish_backend_span_success, - reattach_span, - release_reattached_span, - start_backend_span, - start_streaming_span, - ) - - _, fake_tracer = _patch_app_tracer() - release_a = asyncio.Event() - - with ( - patch("mellea.telemetry.tracing.get_backend_tracer", return_value=fake_tracer), - patch( - "mellea.telemetry.tracing.get_application_tracer", return_value=fake_tracer - ), - ): - start_streaming_span( - "sid-a", has_requirements=False, requirement_count=0, chunking_strategy="x" - ) - start_backend_span("chat", "gid-b", model="m", provider="p") - - async def _task_a_holds_scope() -> None: - reattach_span("sid-a") - try: - await release_a.wait() - finally: - release_reattached_span("sid-a") - - async def _task_b_finishes_span() -> None: - finish_backend_span_success( - "gid-b", operation="chat", usage=None, mot=None, gen=None - ) - - task_a = asyncio.create_task(_task_a_holds_scope()) - await asyncio.sleep(0) # let task A open its scope - with caplog.at_level(logging.WARNING, logger="mellea"): - await asyncio.create_task(_task_b_finishes_span()) - release_a.set() - await task_a - - warnings = [r for r in caplog.records if r.levelno == logging.WARNING] - assert any("across asyncio tasks" in r.getMessage() for r in warnings), ( - "task B's detach should warn despite task A holding a scope" - ) - - -def test_safe_detach_runs_when_no_attach_task(enabled_tracing): - """With no recorded attach task the detach runs unconditionally (no task check).""" - from mellea.telemetry.tracing import _safe_detach - - with patch("mellea.telemetry.tracing.otel_context.detach") as detach: - _safe_detach(MagicMock(), None) - - detach.assert_called_once() - - @pytest.fixture def span_exporter(enabled_tracing): """Attach an in-memory span exporter to the active tracer provider.""" @@ -363,7 +166,7 @@ def _finished_spans(exporter): def _streaming_backend(chunks, *, judge_reply="yes"): """Build an OllamaModelBackend whose AsyncClient is mocked to stream `chunks`. - The mocked `chat` serves both call shapes `stream_with_chunking` triggers: + The mocked `chat` serves both call shapes `stream()` triggers: a streaming generation (`stream=True` → async iterator of deltas) and a non-streaming LLM-as-a-judge `validate()` call (`stream=False` → a single awaited `ChatResponse` carrying `judge_reply`). @@ -418,22 +221,21 @@ def chat(*args, stream=False, **kwargs): async def _run_streaming(backend, *, requirements=None): from mellea.stdlib.components import Message from mellea.stdlib.context import SimpleContext - from mellea.stdlib.streaming import stream_with_chunking + from mellea.stdlib.streaming import stream ctx = SimpleContext().add(Message(role="user", content="Count to three.")) - result = await stream_with_chunking( + async with await stream( Message(role="assistant", content=""), backend, ctx, requirements=requirements - ) - async for _ in result.astream(): - pass - await result.acomplete() - return result + ) as streamer: + async for _ in streamer: + pass + return streamer @pytest.mark.integration @pytest.mark.asyncio -async def test_stream_with_chunking_emits_span_with_lifecycle_events(span_exporter): - """A `stream_with_chunking` call emits one span carrying its events.""" +async def test_stream_emits_span_with_lifecycle_events(span_exporter): + """A `stream` call emits one span carrying its events.""" gen = _streaming_backend(["One.", " Two.", " Three."]) backend = next(gen) try: @@ -442,8 +244,8 @@ async def test_stream_with_chunking_emits_span_with_lifecycle_events(span_export gen.close() spans = _finished_spans(span_exporter) - streaming_span = next((s for s in spans if s.name == "stream_with_chunking"), None) - assert streaming_span is not None, "stream_with_chunking span not emitted" + streaming_span = next((s for s in spans if s.name == "stream"), None) + assert streaming_span is not None, "stream span not emitted" event_names = [e.name for e in streaming_span.events] assert "chunk" in event_names @@ -453,8 +255,8 @@ async def test_stream_with_chunking_emits_span_with_lifecycle_events(span_export @pytest.mark.integration @pytest.mark.asyncio -async def test_stream_with_chunking_chat_span_nests_under_streaming_span(span_exporter): - """The backend `chat` span nests under the `stream_with_chunking` span.""" +async def test_stream_chat_span_nests_under_streaming_span(span_exporter): + """The backend `chat` span nests under the `stream` span.""" gen = _streaming_backend(["One.", " Two."]) backend = next(gen) try: @@ -463,8 +265,8 @@ async def test_stream_with_chunking_chat_span_nests_under_streaming_span(span_ex gen.close() spans = _finished_spans(span_exporter) - streaming_span = next((s for s in spans if s.name == "stream_with_chunking"), None) - assert streaming_span is not None, "stream_with_chunking span not emitted" + streaming_span = next((s for s in spans if s.name == "stream"), None) + assert streaming_span is not None, "stream span not emitted" chat_span = next((s for s in spans if s.name == "chat test-model"), None) assert chat_span is not None, "chat span not emitted" @@ -472,7 +274,7 @@ async def test_stream_with_chunking_chat_span_nests_under_streaming_span(span_ex if _CONTEXT_ATTACH_SUPPORTED: assert chat_span.parent is not None assert chat_span.parent.span_id == streaming_span.context.span_id, ( - "chat span should nest under stream_with_chunking" + "chat span should nest under stream" ) else: assert chat_span.parent is None, "chat span should be flat on Python <=3.11" @@ -480,10 +282,8 @@ async def test_stream_with_chunking_chat_span_nests_under_streaming_span(span_ex @pytest.mark.integration @pytest.mark.asyncio -async def test_stream_with_chunking_validation_chat_span_is_sibling_of_generation( - span_exporter, -): - """Both `chat` spans parent under `stream_with_chunking`, which owns the events.""" +async def test_stream_validation_chat_span_is_sibling_of_generation(span_exporter): + """Both `chat` spans parent under `stream`, which owns the events.""" from mellea.core.requirement import Requirement gen = _streaming_backend(["A full sentence here."]) @@ -494,8 +294,8 @@ async def test_stream_with_chunking_validation_chat_span_is_sibling_of_generatio gen.close() spans = _finished_spans(span_exporter) - streaming_span = next((s for s in spans if s.name == "stream_with_chunking"), None) - assert streaming_span is not None, "stream_with_chunking span not emitted" + streaming_span = next((s for s in spans if s.name == "stream"), None) + assert streaming_span is not None, "stream span not emitted" chat_spans = [s for s in spans if s.name == "chat test-model"] assert len(chat_spans) == 2, f"expected 2 chat spans, got {len(chat_spans)}" @@ -504,7 +304,7 @@ async def test_stream_with_chunking_validation_chat_span_is_sibling_of_generatio assert all( s.parent is not None and s.parent.span_id == streaming_id for s in chat_spans - ), "both chat spans should nest directly under stream_with_chunking" + ), "both chat spans should nest directly under stream" else: assert all(s.parent is None for s in chat_spans), ( "chat spans should be flat on Python <=3.11" @@ -526,12 +326,8 @@ async def test_stream_with_chunking_validation_chat_span_is_sibling_of_generatio @pytest.mark.integration @pytest.mark.asyncio -async def test_stream_with_chunking_span_ends_error_on_early_exit( - span_exporter, caplog -): - """A mid-stream validation fail still closes the streaming span (ERROR), quietly.""" - import logging - +async def test_stream_span_ends_error_on_early_exit(span_exporter): + """A mid-stream validation fail still closes the streaming span (ERROR).""" from opentelemetry.trace import StatusCode from mellea.core.requirement import PartialValidationResult, Requirement @@ -543,16 +339,12 @@ async def stream_validate(self, chunk, *, backend, ctx): gen = _streaming_backend(["A full sentence here."]) backend = next(gen) try: - with caplog.at_level(logging.WARNING, logger="mellea"): - await _run_streaming(backend, requirements=[_FailingReq()]) + await _run_streaming(backend, requirements=[_FailingReq()]) finally: gen.close() streaming_span = next( - (s for s in _finished_spans(span_exporter) if s.name == "stream_with_chunking"), - None, + (s for s in _finished_spans(span_exporter) if s.name == "stream"), None ) - assert streaming_span is not None, "stream_with_chunking span not emitted" + assert streaming_span is not None, "stream span not emitted" assert streaming_span.status.status_code == StatusCode.ERROR - cross_task = [r for r in caplog.records if "across asyncio tasks" in r.getMessage()] - assert not cross_task, "early exit should not warn about cross-task detach"