From 9071c6066d1993c9bd527d9ebab26796ddc72e47 Mon Sep 17 00:00:00 2001 From: jaygupta17 Date: Wed, 26 Aug 2026 12:30:15 +0530 Subject: [PATCH] feat(stream): PartialOutput events for streaming structured output With output_type set, streams only exposed text deltas; the parsed object appeared solely after StreamEnd via Stream.output. Consumers rendering progressive results (live forms, tables) had no way to see the object take shape. Streams now emit a synthetic PartialOutput event after each text delta that changes the parse of the accumulated JSON, plus expose the latest snapshot as Stream.partial_output. Parsing uses json-repair with stream_stable=True so truncated values stay verbatim strings instead of being creatively repaired mid-stream; non-dict parses (dangling keys) are skipped. Partials are plain dicts and deliberately unvalidated -- Stream.output remains the validated contract. Agent runs forward model stream events unchanged, so AgentStream consumers get partials for free; the UI adapter's event match ignores unknown kinds. --- .../content/docs/basics/streaming.mdx | 23 +++ .../content/docs/reference/events.mdx | 7 + pyproject.toml | 1 + src/ai/models/core/api.py | 66 ++++++-- src/ai/types/events.py | 16 ++ tests/models/core/test_api.py | 157 +++++++++++++++++- uv.lock | 13 +- 7 files changed, 268 insertions(+), 15 deletions(-) diff --git a/docs/ai-python/content/docs/basics/streaming.mdx b/docs/ai-python/content/docs/basics/streaming.mdx index 09df7c9f..50d65423 100644 --- a/docs/ai-python/content/docs/basics/streaming.mdx +++ b/docs/ai-python/content/docs/basics/streaming.mdx @@ -111,6 +111,29 @@ if __name__ == "__main__": `stream.output` returns text by default. When you pass `output_type`, it returns an instance of that Pydantic model after the stream finishes. +### Stream partial objects + +While streaming with `output_type`, the stream also emits +`ai.events.PartialOutput` events: a best-effort parse of the JSON generated so +far, updated as deltas arrive. The value is a plain dict and does not validate +against the model -- use it for progressive rendering: + +```python +async with ai.stream(model, messages, output_type=UprisingForecast) as stream: + async for event in stream: + match event: + case ai.events.PartialOutput(value=partial): + print(partial.get("eta", "...")) + case ai.events.TextDelta(): + pass + + forecast = stream.output # validated instance, available after the stream +``` + +The latest snapshot is also available as `stream.partial_output`. Partials are +not validated; `stream.output` remains the validated result once the stream has +ended. + ## Pass tools `ai.stream` accepts a list of `ai.Tool`, however, it does not execute function diff --git a/docs/ai-python/content/docs/reference/events.mdx b/docs/ai-python/content/docs/reference/events.mdx index b0f1d994..29679561 100644 --- a/docs/ai-python/content/docs/reference/events.mdx +++ b/docs/ai-python/content/docs/reference/events.mdx @@ -26,6 +26,13 @@ Text events: - `TextEnd` - `StreamEnd` +Structured-output events: + +- `PartialOutput`: emitted while streaming with `output_type` set. `value` is + the best-effort parse of the JSON generated so far as a plain dict; it does + not validate against `output_type` (the validated instance is + `Stream.output` after the stream ends). + `StreamEnd` carries the final response metadata. New in 0.4.0: - `finish_reason`: Why the model stopped: `stop`, `length`, `content_filter`, diff --git a/pyproject.toml b/pyproject.toml index 2dc71059..3e5d733d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ classifiers = [ requires-python = ">=3.12" dependencies = [ "httpx>=0.28.1", + "json-repair==0.*", "modelsdotdev==0.*", "pydantic>=2.13", "typing-extensions>=4.15.0", diff --git a/src/ai/models/core/api.py b/src/ai/models/core/api.py index 206dd38b..e24c9523 100644 --- a/src/ai/models/core/api.py +++ b/src/ai/models/core/api.py @@ -2,6 +2,7 @@ import contextlib import dataclasses +from collections import deque from contextlib import AbstractAsyncContextManager from typing import ( TYPE_CHECKING, @@ -14,6 +15,7 @@ runtime_checkable, ) +import json_repair import pydantic # ``typing.TypeVar`` lacks the ``default=`` kwarg on Python <3.13. @@ -125,6 +127,11 @@ def __init__( # (``Stream(gen)``, ``Stream.replay_message``). self._span: telemetry.Span[telemetry.AiStreamSpanData] | None = None self._first_output_seen = False + # Synthetic PartialOutput events waiting to be yielded ahead of the + # next provider event, plus the latest parsed snapshot behind + # ``Stream.partial_output``. Only active when ``output_type`` is set. + self._pending_events: deque[types.events.Event] = deque() + self._last_partial: dict[str, Any] | None = None @classmethod def replay_message( @@ -192,19 +199,41 @@ def __aiter__(self) -> Self: return self async def __anext__(self: Self) -> types.events.Event: - try: - event = await self._gen.__anext__() - except StopAsyncIteration: - if not self._hydrator.ended: - raise errors.ProviderIncompleteResponseError( - "provider stream ended without a finish event; " - "the response is incomplete", - # Premature termination is a transient transport or - # provider failure: worth retrying. - is_retryable=True, - ) from None - raise - event = self._hydrator.feed(event) + if self._pending_events: + # Synthetic events (e.g. PartialOutput) bypass the hydrator, + # so stamp the live in-progress message ourselves — every + # yielded ModelEvent must carry it, never the dummy default. + event = self._pending_events.popleft().model_copy( + update={"message": self.message} + ) + else: + try: + event = await self._gen.__anext__() + except StopAsyncIteration: + if not self._hydrator.ended: + raise errors.ProviderIncompleteResponseError( + "provider stream ended without a finish event; " + "the response is incomplete", + # Premature termination is a transient transport or + # provider failure: worth retrying. + is_retryable=True, + ) from None + raise + event = self._hydrator.feed(event) + # Structured-output streaming: after each text delta, re-parse the + # accumulated JSON and queue a best-effort snapshot ahead of the + # next provider event. ``stream_stable`` keeps truncated values as + # verbatim strings instead of creatively repairing them, so + # successive snapshots only grow toward the final object. + if self._output_type is not None and isinstance( + event, types.events.TextDelta + ): + parsed = json_repair.loads(self.message.text, stream_stable=True) + if isinstance(parsed, dict) and parsed != self._last_partial: + self._last_partial = parsed + self._pending_events.append( + types.events.PartialOutput(value=parsed) + ) # Milestones on the live span: replayed work gets no synthetic # timings (span.replay covers the replay branch of ``stream()``, # event.replay covers individual synthetic events). @@ -273,6 +302,17 @@ def output(self) -> StreamOutputT: """ return cast("StreamOutputT", self.message.get_output(self._output_type)) + @property + def partial_output(self) -> dict[str, Any] | None: + """Latest best-effort parse of the structured output so far. + + Only populated when ``output_type`` was set and at least one + text delta has arrived. The value is a plain dict and does not + validate against ``output_type``; use :attr:`output` for the + validated instance once the stream has ended. + """ + return self._last_partial + async def _replay_tool_calls( msg: types.messages.Message, diff --git a/src/ai/types/events.py b/src/ai/types/events.py index 80083766..caf2c490 100644 --- a/src/ai/types/events.py +++ b/src/ai/types/events.py @@ -92,6 +92,21 @@ class TextEnd(ModelEvent): kind: Literal["text_end"] = "text_end" +class PartialOutput(ModelEvent): + """Best-effort snapshot of structured output generated so far. + + Emitted while streaming with ``output_type`` set, after each text + delta that changes the parse. ``value`` is the partial parse of the + JSON generated so far as a plain dict -- it does not validate + against ``output_type`` and grows toward the final object, which is + available validated via ``Stream.output`` when the stream ends. + """ + + value: dict[str, Any] + + kind: Literal["partial_output"] = "partial_output" + + class ReasoningStart(ModelEvent): block_id: str = "" @@ -179,6 +194,7 @@ class FileEvent(ModelEvent): | TextStart | TextDelta | TextEnd + | PartialOutput | ReasoningStart | ReasoningDelta | ReasoningEnd diff --git a/tests/models/core/test_api.py b/tests/models/core/test_api.py index e628cf62..9ddb2ab0 100644 --- a/tests/models/core/test_api.py +++ b/tests/models/core/test_api.py @@ -1,7 +1,7 @@ from __future__ import annotations import asyncio -from collections.abc import AsyncGenerator, Sequence +from collections.abc import AsyncGenerator, Callable, Sequence from typing import Any, Literal, cast import pydantic @@ -822,3 +822,158 @@ async def test_replayed_turn_gets_replay_span(recorder: Recorder) -> None: assert isinstance(call.data, ai.experimental_telemetry.AiStreamSpanData) assert call.data.message is not None assert call.data.message.text == "prior turn" + + +# -- Streaming partial output ---------------------------------------------- + + +class _Answer(pydantic.BaseModel): + a: str + b: int = 0 + + +class _PinsOutput(pydantic.BaseModel): + a: int + b: int = 0 + c: str = "" + + +async def _scripted_deltas(chunks: list[str]) -> AsyncGenerator[events_.Event]: + yield events_.StreamStart() + yield events_.TextStart(block_id="t1") + for chunk in chunks: + yield events_.TextDelta(block_id="t1", chunk=chunk) + yield events_.TextEnd(block_id="t1") + yield events_.StreamEnd() + + +def _scripted_impl( + chunks: list[str], +) -> Callable[ + [models.Model, list[messages_.Message]], AsyncGenerator[events_.Event] +]: + """Provider stream seam replaying hand-built text deltas.""" + + async def impl( + model: models.Model, + messages: list[messages_.Message], + **kwargs: Any, + ) -> AsyncGenerator[events_.Event]: + async for event in _scripted_deltas(chunks): + yield event + + return impl + + +async def test_stream_emits_partial_output_snapshots() -> None: + MOCK_PROVIDER._stream_impl = _scripted_impl(['{"a": "hel', 'lo"}']) + + seen: list[dict[str, Any]] = [] + order: list[str] = [] + async with models.stream( + MOCK_MODEL, [ai.user_message("go")], output_type=_Answer + ) as stream: + async for event in stream: + if isinstance(event, events_.PartialOutput): + seen.append(event.value) + order.append("partial") + elif isinstance(event, events_.TextDelta): + order.append("delta") + + # Every yielded ModelEvent carries the live in-progress + # message, never the dummy default — synthetic PartialOutput + # events bypass the hydrator so Stream stamps them itself. + assert event.message.id != "" + + assert stream.partial_output is None or isinstance( + stream.partial_output, dict + ) + + # One snapshot per parse-changing delta: truncated string stays + # stable under stream_stable, then the completed object. + assert seen == [{"a": "hel"}, {"a": "hello"}] + assert order == ["delta", "partial", "delta", "partial"] + assert stream.output == _Answer(a="hello") + + +async def test_stream_partial_output_property_tracks_latest() -> None: + MOCK_PROVIDER._stream_impl = _scripted_impl(['{"a": "x', '"}']) + + latest: dict[str, Any] | None = None + async with models.stream( + MOCK_MODEL, [ai.user_message("go")], output_type=_Answer + ) as stream: + assert stream.partial_output is None + async for _ in stream: + if stream.partial_output is not None: + latest = stream.partial_output + + assert latest == {"a": "x"} + + +async def test_stream_without_output_type_never_emits_partials() -> None: + MOCK_PROVIDER._stream_impl = _scripted_impl(["plain text"]) + + partials: list[events_.PartialOutput] = [] + async with models.stream(MOCK_MODEL, [ai.user_message("go")]) as stream: + async for event in stream: + if isinstance(event, events_.PartialOutput): + partials.append(event) + elif isinstance(event, events_.TextDelta): + assert stream.partial_output is None + + assert partials == [] + + +async def test_stream_partial_output_pins_json_repair_behaviors() -> None: + # Pins observed json_repair stream_stable behaviors so a minor bump + # cannot silently change what consumers see mid-stream. + MOCK_PROVIDER._stream_impl = _scripted_impl( + [ + '{"a": 1, "b"', + ': 2, "c": "trunc', + 'ated"}', + ] + ) + + snapshots: list[dict[str, Any]] = [] + async with models.stream( + MOCK_MODEL, + [ai.user_message("go")], + output_type=_PinsOutput, + ) as stream: + async for event in stream: + if isinstance(event, events_.PartialOutput): + snapshots.append(event.value) + + assert snapshots == [ + {"a": 1}, # dangling key without a value is dropped + {"a": 1, "b": 2, "c": "trunc"}, # truncated string kept verbatim + {"a": 1, "b": 2, "c": "truncated"}, + ] + assert stream.output == _PinsOutput(a=1, b=2, c="truncated") + + +async def test_agent_stream_forwards_partial_output() -> None: + # Two text parts -> two deltas (emit_events_for_messages emits one + # delta per part). + answer = ai.messages.Message( + id="msg-1", + role="assistant", + parts=[ + ai.messages.TextPart(text='{"a": "hel'), + ai.messages.TextPart(text='lo"}'), + ], + ) + mock_llm([[answer]]) + + my_agent = ai.Agent() + partials: list[dict[str, Any]] = [] + async with my_agent.run( + MOCK_MODEL, [ai.user_message("go")], output_type=_Answer + ) as stream: + async for event in stream: + if isinstance(event, events_.PartialOutput): + partials.append(event.value) + + assert partials == [{"a": "hel"}, {"a": "hello"}] diff --git a/uv.lock b/uv.lock index 7637b4d4..c4c93653 100644 --- a/uv.lock +++ b/uv.lock @@ -8,7 +8,7 @@ resolution-markers = [ ] [options] -exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. +exclude-newer = "2026-08-24T06:32:46.067216Z" exclude-newer-span = "P2D" [[package]] @@ -16,6 +16,7 @@ name = "ai" source = { editable = "." } dependencies = [ { name = "httpx" }, + { name = "json-repair" }, { name = "modelsdotdev" }, { name = "pydantic" }, { name = "typing-extensions" }, @@ -60,6 +61,7 @@ dev = [ requires-dist = [ { name = "anthropic", marker = "extra == 'anthropic'", specifier = ">=0.83.0,<2.0.0" }, { name = "httpx", specifier = ">=0.28.1" }, + { name = "json-repair", specifier = "==0.*" }, { name = "mcp", marker = "extra == 'mcp'", specifier = ">=1.18.0" }, { name = "modelsdotdev", specifier = "==0.*" }, { name = "openai", marker = "extra == 'openai'", specifier = ">=2.34.0,<4.0.0" }, @@ -595,6 +597,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2f/9c/6753e6522b8d0ef07d3a3d239426669e984fb0eba15a315cdbc1253904e4/jiter-0.12.0-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c24e864cb30ab82311c6425655b0cdab0a98c5d973b065c66a3f020740c2324c", size = 346110, upload-time = "2025-11-09T20:49:21.817Z" }, ] +[[package]] +name = "json-repair" +version = "0.63.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b5/fe/708f83b5d42f0dc687fc720b0eba1217d9632539a5ccac432c4cfc4794ee/json_repair-0.63.3.tar.gz", hash = "sha256:87747d2d136124961a56392687c84558f5609244301df5ce641b84dac52f3b92", size = 52739, upload-time = "2026-08-19T17:51:27.829Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e4/63/7d171e5fb75ec70e714ec119df1052af5a7224fc19c5668f0e99faf50e6a/json_repair-0.63.3-py3-none-any.whl", hash = "sha256:3681cc25e44f17749d6d440ae69d7c4275609a75092bcde03c6bacc3f3c8475e", size = 51170, upload-time = "2026-08-19T17:51:26.362Z" }, +] + [[package]] name = "jsonschema" version = "4.26.0"