feat: streaming support for structured output - #142
Open
Meguazy wants to merge 7 commits into
Open
Conversation
Add stream_structured_response / a_stream_structured_response to the base Client with non-abstract provider hooks (default NotImplementedError). Intermediate chunks carry raw JSON fragments in delta, plus the accumulated snapshot and a lazily parsed partial dict; the final chunk carries the validated StructuredBlock. Agents with stream=True and output_cls now stream structured output and fall back transparently to the blocking structured call when the client does not implement the new hooks.
Stream via generate_content_stream with response_mime_type=application/json and response_schema, accumulate the JSON text deltas and validate the final buffer into the output model. Verified end-to-end against the Gemini API.
Stream via chat.completions.create(stream=True) with a json_schema response_format built through type_to_response_format_param (same mechanism as beta.chat.completions.parse), reusing the existing stream delta accumulation and strict-tools handling. Accumulate the JSON text deltas and validate the final buffer into the output model. Verified end-to-end against a local Ollama server (qwen3.5:9b).
Stream via chat.stream / chat.stream_async with a json_schema response_format built through response_format_from_pydantic_model, accumulate the JSON text deltas and validate the final buffer into the output model. Bump mistralai floor to >=1.5.0: mistralai.extra and chat.parse (already used by the existing structured path) do not exist in earlier releases. Verified end-to-end against the Mistral API (mistral-small-latest).
Stream via messages.create(stream=True) with output_config json_schema (transform_schema, the same transform used by messages.parse), reusing the existing raw-event stream state machine. Text deltas carry the JSON fragments; the final chunk validates the accumulated buffer into the output model. Verified end-to-end against the Anthropic API (claude-sonnet-4-5).
Stream via the Responses API responses.stream(text_format=...) context manager; delta events carry the JSON fragments and the completed event flows through _response_to_client_response, which already extracts the parsed StructuredBlock. The json_object sentinel is handled by parsing output_text directly. Verified end-to-end against the OpenAI API (gpt-4o-mini).
## Summary Adds streaming support for structured output across the framework: the base `Client`, the Agent runner, and the provider clients (Google, OpenAI, OpenAI-like, Anthropic, Mistral). ### Core (`datapizza-ai-core`) - New `stream_structured_response` / `a_stream_structured_response` on the base `Client`, with non-abstract provider hooks defaulting to `NotImplementedError` (existing subclasses keep working unchanged). - Chunk contract: intermediate chunks carry the raw JSON fragments in `chunk.delta` (`content=[]`); the final chunk carries the validated Pydantic instance as a `StructuredBlock` plus `usage` and `stop_reason`. Invalid JSON at end of stream raises `pydantic.ValidationError`. - Progressive consumption built in: every chunk exposes `chunk.snapshot` (JSON text accumulated so far) and `chunk.partial` — a lazy, best-effort parsed dict (via `pydantic_core.from_json(allow_partial=...)`), so apps can render fields as they arrive. - Agents with `stream=True` + `output_cls` now stream the JSON deltas before the final `StepResult`; clients without streaming-structured support fall back transparently to the blocking structured call. ### Providers | Provider | Mechanism | |---|---| | Google | `generate_content_stream` + `response_mime_type=application/json` + `response_schema` | | OpenAI | Responses API `responses.stream(text_format=...)` | | OpenAI-like | `chat.completions.create(stream=True)` + json_schema `response_format` (no beta helpers, works with vLLM/Ollama/any compatible server) | | Anthropic | `messages.create(stream=True)` + `output_config` json_schema, reusing the existing raw-event state machine | | Mistral | `chat.stream` / `chat.stream_async` + `response_format_from_pydantic_model` (mistralai floor bumped to `>=1.5.0`, already required de facto by the existing `chat.parse` usage) | Bedrock and WatsonX are intentionally untouched: they inherit the base default and raise `NotImplementedError`, with the Agent fallback keeping behavior identical to before. ## Testing - ~90 new mocked unit tests across core and provider packages (sync + async, invalid-JSON, request-params assertions, `json_object` sentinel), following the existing test conventions. All suites green, `ruff check` clean. - Verified end-to-end against real APIs: Google (Gemini), OpenAI (`gpt-4o-mini`, Responses API), Anthropic (`claude-sonnet-4-5`), Mistral (`mistral-small-latest`), and OpenAI-like against a local Ollama server — streamed deltas, progressive `partial` dicts, and validated final objects on both sync and async paths.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds streaming support for structured output across the framework: the base
Client, the Agent runner, and the provider clients (Google, OpenAI, OpenAI-like, Anthropic, Mistral).Core (
datapizza-ai-core)stream_structured_response/a_stream_structured_responseon the baseClient, with non-abstract provider hooks defaulting toNotImplementedError(existing subclasses keep working unchanged).chunk.delta(content=[]); the final chunk carries the validated Pydantic instance as aStructuredBlockplususageandstop_reason. Invalid JSON at end of stream raisespydantic.ValidationError.chunk.snapshot(JSON text accumulated so far) andchunk.partial— a lazy, best-effort parsed dict (viapydantic_core.from_json(allow_partial=...)), so apps can render fields as they arrive.stream=True+output_clsnow stream the JSON deltas before the finalStepResult; clients without streaming-structured support fall back transparently to the blocking structured call.Providers
generate_content_stream+response_mime_type=application/json+response_schemaresponses.stream(text_format=...)chat.completions.create(stream=True)+ json_schemaresponse_format(no beta helpers, works with vLLM/Ollama/any compatible server)messages.create(stream=True)+output_configjson_schema, reusing the existing raw-event state machinechat.stream/chat.stream_async+response_format_from_pydantic_model(mistralai floor bumped to>=1.5.0, already required de facto by the existingchat.parseusage)Bedrock and WatsonX are intentionally untouched: they inherit the base default and raise
NotImplementedError, with the Agent fallback keeping behavior identical to before.Testing
json_objectsentinel), following the existing test conventions. All suites green,ruff checkclean.gpt-4o-mini, Responses API), Anthropic (claude-sonnet-4-5), Mistral (mistral-small-latest), and OpenAI-like against a local Ollama server — streamed deltas, progressivepartialdicts, and validated final objects on both sync and async paths.