Support output_type on stream

examples/samples/structured_output.py

@@ -6,7 +6,7 @@
 
 import ai
 
-model = ai.ai_gateway("anthropic/claude-sonnet-4")
+model = ai.ai_gateway("anthropic/claude-sonnet-4.6")
 
 
 class Recipe(pydantic.BaseModel):
@@ -20,9 +20,21 @@ class Recipe(pydantic.BaseModel):
 
 
 async def main() -> None:
-    # Broken for now: stream(output_type=...) requests JSON/schema mode, but
-    # the stream wrapper does not yet validate final text into s.output.
-    raise RuntimeError("structured output aggregation needs to be implemented")
+    # Stream with structured output — watch JSON arrive, get validated at the end.
+    async with ai.stream(model, messages, output_type=Recipe) as s:
+        async for event in s:
+            if isinstance(event, ai.events.TextDelta):
+                print(event.chunk, end="", flush=True)
+            elif isinstance(event, ai.events.StreamEnd):
+                print()
+
+    # After iteration, s.output validates the streamed JSON against
+    # the output_type the stream was opened with.  ``Stream[Recipe]``
+    # types ``s.output`` as ``Recipe``.
+    recipe = s.output
+    print(f"\n\nParsed recipe: {recipe.name}")
+    print(f"  Ingredients: {', '.join(recipe.ingredients)}")
+    print(f"  Prep time: {recipe.prep_time_minutes} min")
 
 
 if __name__ == "__main__":

pyproject.toml

@@ -13,6 +13,7 @@ dependencies = [
     "mcp>=1.18.0",
     "openai>=2.14.0",
     "pydantic>=2.12.5",
+    "typing-extensions>=4.15.0",
     "vercel>=0.3.8",
 ]
 

src/ai/agents/agent.py

@@ -16,10 +16,24 @@
     Coroutine,
     Sequence,
 )
-from typing import Annotated, Any, Protocol, Self, get_type_hints, overload
+from contextlib import AbstractAsyncContextManager
+from typing import (
+    Annotated,
+    Any,
+    Generic,
+    Protocol,
+    Self,
+    cast,
+    get_type_hints,
+    overload,
+)
 
 import pydantic
 
+# ``typing.TypeVar`` lacks the ``default=`` kwarg on Python <3.13.
+# Use the typing_extensions backport so this works on 3.12 too.
+from typing_extensions import TypeVar  # noqa: UP035
+
 from .. import models, types, util
 from ..types import builders
 from ..types import events as events_
@@ -606,6 +620,9 @@ class Context(pydantic.BaseModel):
     model: models.Model
     messages: list[types.messages.Message]
     tools: list[Tool]
+    output_type: type[pydantic.BaseModel] | None = pydantic.Field(
+        default=None, exclude=True, repr=False
+    )
 
     _agent_tools_by_name: dict[str, AgentTool] = pydantic.PrivateAttr(
         default_factory=dict
@@ -672,7 +689,13 @@ def add(
             self.messages.append(msg)
 
 
-class AgentStream:
+# Agent run output type.  Defaults to ``str``: when ``Agent.run`` was
+# called without an ``output_type``, ``AgentStream.output`` returns the
+# final assistant message's concatenated text.
+AgentOutputT = TypeVar("AgentOutputT", default=str)
+
+
+class AgentStream(Generic[AgentOutputT]):
     """Async-iterable wrapper around an agent run's event stream.
 
     Exposes the run's :class:`Context` via :attr:`context` so callers can
@@ -720,6 +743,18 @@ def context(self) -> Context:
     def messages(self) -> list[types.messages.Message]:
         return self._context.messages
 
+    @property
+    def output(self) -> AgentOutputT:
+        """Return the run's output, parsed as the ``output_type`` given to ``run``.
+
+        Defaults to the final assistant message's concatenated text.
+        When an ``output_type`` was passed, the assistant message's text
+        is validated as JSON against that type and the parsed instance
+        is returned.
+        """
+        last = self._context.messages[-1]
+        return cast(AgentOutputT, last.get_output(self._context.output_type))
+
 
 def tool_result(
     *items: types.messages.Message
@@ -894,14 +929,31 @@ async def loop(self, context: Context) -> AsyncGenerator[events_.AgentEvent]:
                 # the loop.
                 context.add(tr.get_tool_message())
 
-    @contextlib.asynccontextmanager
-    async def run(
+    @overload
+    def run(
         self,
         model: models.Model,
         messages: list[types.messages.Message],
         *,
         middleware: list[middleware_.Middleware] | None = None,
-    ) -> AsyncIterator[AgentStream]:
+    ) -> AbstractAsyncContextManager[AgentStream[str]]: ...
+    @overload
+    def run[T: pydantic.BaseModel](
+        self,
+        model: models.Model,
+        messages: list[types.messages.Message],
+        *,
+        output_type: type[T],
+        middleware: list[middleware_.Middleware] | None = None,
+    ) -> AbstractAsyncContextManager[AgentStream[T]]: ...
+    def run(
+        self,
+        model: models.Model,
+        messages: list[types.messages.Message],
+        *,
+        output_type: type[pydantic.BaseModel] | None = None,
+        middleware: list[middleware_.Middleware] | None = None,
+    ) -> AbstractAsyncContextManager[AgentStream[Any]]:
         """Run the agent loop, yielding events to the consumer.
 
         Used as an async context manager whose value the event stream,
@@ -915,6 +967,9 @@ async def run(
         Args:
             model: The model to use for LLM calls.
             messages: Initial conversation messages.
+            output_type: Optional Pydantic model the model's output must
+                conform to.  When set, ``stream.output`` validates the
+                final assistant message's text against it.
             middleware: Optional list of middleware to apply to this run.
                 First in the list = outermost.  Middleware wraps model
                 calls, tool calls, hooks, and the run itself.
@@ -923,10 +978,24 @@ async def run(
         ``yield_from(..., label=...)`` — the label flows via
         ``PartialToolCallResult`` rather than on individual messages.
         """
+        return self._run(
+            model, messages, output_type=output_type, middleware=middleware
+        )
+
+    @contextlib.asynccontextmanager
+    async def _run(
+        self,
+        model: models.Model,
+        messages: list[types.messages.Message],
+        *,
+        output_type: type[pydantic.BaseModel] | None,
+        middleware: list[middleware_.Middleware] | None,
+    ) -> AsyncIterator[AgentStream[Any]]:
         context = Context(
             model=model,
             messages=list(messages),
             tools=[t.tool for t in self._tools],
+            output_type=output_type,
         )
         context._agent_tools_by_name = {t.name: t for t in self._tools}
         _process_interrupted_hooks(context.messages)

src/ai/models/core/api.py

@@ -2,17 +2,26 @@
 import dataclasses
 from collections.abc import AsyncGenerator, AsyncIterator, Sequence
 from contextlib import AbstractAsyncContextManager
-from typing import Any, Protocol, Self, overload, runtime_checkable
+from typing import Any, Generic, Protocol, Self, cast, overload, runtime_checkable
 
 import pydantic
 
+# ``typing.TypeVar`` lacks the ``default=`` kwarg on Python <3.13.
+# Use the typing_extensions backport so this works on 3.12 too.
+from typing_extensions import TypeVar  # noqa: UP035
+
 from ... import types
 from ...types import integrity
 from . import adapters
 from . import client as client_
 from . import model as model_
 from . import params as params_
 
+# Stream output type.  Defaults to ``str``: when the stream was opened
+# without an ``output_type``, ``Stream.output`` returns the concatenated
+# message text.
+StreamOutputT = TypeVar("StreamOutputT", default=str)
+
 
 @dataclasses.dataclass(frozen=True)
 class StreamRequest:
@@ -76,14 +85,15 @@ async def _do_generate(self, request: GenerateRequest) -> types.messages.Message
 _default_executor = Executor()
 
 
-class Stream:
+class Stream(Generic[StreamOutputT]):
     """Async-iterable wrapper around an adapter's event stream."""
 
     def __init__(
         self,
         gen: AsyncGenerator[types.events.Event],
         *,
         seed_message: types.messages.Message | None = None,
+        output_type: type[StreamOutputT] | None = None,
     ) -> None:
         """Wrap an event generator.
 
@@ -93,12 +103,21 @@ def __init__(
         being rebuilt from synthetic events.  When ``None`` (default),
         an empty assistant message is created and rebuilt from the
         incoming events.
+
+        ``output_type`` is the Pydantic model the request was constrained
+        to.  When set, ``Stream.output`` validates the streamed JSON text
+        against it.  When ``None`` (default), ``Stream.output`` returns
+        the concatenated text content unchanged.
         """
         self._gen = gen
         self._message: types.messages.Message = seed_message or types.messages.Message(
             role="assistant", parts=[]
         )
         self._parts: dict[str, types.messages.Part] = {}
+        # ``output_type`` is typed against the public ``StreamOutputT`` type
+        # param for ergonomics; internally we know it's a Pydantic model
+        # subclass (or None for the text-default case).
+        self._output_type = cast(type[pydantic.BaseModel] | None, output_type)
 
     async def aclose(self) -> None:
         await self._gen.aclose()
@@ -145,8 +164,14 @@ def tool_calls(self) -> list[types.messages.ToolCallPart]:
         return self._message.tool_calls
 
     @property
-    def output(self) -> Any:
-        return self._message.output
+    def output(self) -> StreamOutputT:
+        """Return the streamed output as the ``output_type`` passed in.
+
+        Defaults to the concatenated message text.  When a Pydantic
+        model subclass was passed, validates the streamed JSON against
+        it and returns the parsed instance.
+        """
+        return cast(StreamOutputT, self._message.get_output(self._output_type))
 
     def _aggregate_event(self, event: types.events.Event) -> dict[str, Any]:
         updates: dict[str, Any] = {}
@@ -305,7 +330,7 @@ async def _replay_tool_calls(
 
 @runtime_checkable
 class StreamContext(Protocol):
-    """Anything that exposes ``model``/``messages``/``tools``.
+    """Anything that exposes ``model``/``messages``/``tools``/``output_type``.
 
     Used to let callers pass an ``agents.Context`` to :func:`stream`
     without an import-time circular dependency.
@@ -317,26 +342,44 @@ def model(self) -> model_.Model: ...
     def messages(self) -> list[types.messages.Message]: ...
     @property
     def tools(self) -> list[types.tools.Tool]: ...
+    @property
+    def output_type(self) -> type[pydantic.BaseModel] | None: ...
 
 
 @overload
 def stream(
     *,
     context: StreamContext,
-    output_type: type[pydantic.BaseModel] | None = None,
     params: Any = None,
     executor: StreamExecutor = _default_executor,
-) -> AbstractAsyncContextManager[Stream]: ...
+) -> AbstractAsyncContextManager[Stream[str]]: ...
+@overload
+def stream[T: pydantic.BaseModel](
+    *,
+    context: StreamContext,
+    output_type: type[T],
+    params: Any = None,
+    executor: StreamExecutor = _default_executor,
+) -> AbstractAsyncContextManager[Stream[T]]: ...
 @overload
-def stream[ProviderParamsT: pydantic.BaseModel](
+def stream(
     model: model_.Model,
     messages: list[types.messages.Message],
     *,
     tools: Sequence[types.tools.Tool] | None = None,
-    output_type: type[pydantic.BaseModel] | None = None,
     params: Any = None,
     executor: StreamExecutor = _default_executor,
-) -> AbstractAsyncContextManager[Stream]: ...
+) -> AbstractAsyncContextManager[Stream[str]]: ...
+@overload
+def stream[T: pydantic.BaseModel](
+    model: model_.Model,
+    messages: list[types.messages.Message],
+    *,
+    tools: Sequence[types.tools.Tool] | None = None,
+    output_type: type[T],
+    params: Any = None,
+    executor: StreamExecutor = _default_executor,
+) -> AbstractAsyncContextManager[Stream[T]]: ...
 def stream(
     model: model_.Model | None = None,
     messages: list[types.messages.Message] | None = None,
@@ -346,7 +389,7 @@ def stream(
     output_type: type[pydantic.BaseModel] | None = None,
     params: Any = None,
     executor: StreamExecutor = _default_executor,
-) -> AbstractAsyncContextManager[Stream]:
+) -> AbstractAsyncContextManager[Stream[Any]]:
     """Stream an LLM response.
 
     Used as an async context manager whose value is the :class:`Stream`.
@@ -368,6 +411,8 @@ def stream(
         model = context.model
         messages = context.messages
         tools = context.tools
+        if output_type is None:
+            output_type = context.output_type
     elif model is None or messages is None:
         raise TypeError("stream() requires either model and messages or context=")
 
@@ -390,10 +435,14 @@ async def _stream(
     output_type: type[pydantic.BaseModel] | None,
     params: Any,
     executor: StreamExecutor,
-) -> AsyncIterator[Stream]:
+) -> AsyncIterator[Stream[Any]]:
     if messages and messages[-1].replay:
         last = messages[-1]
-        s = Stream(_replay_tool_calls(last), seed_message=last.model_copy(deep=True))
+        s = Stream(
+            _replay_tool_calls(last),
+            seed_message=last.model_copy(deep=True),
+            output_type=output_type,
+        )
     else:
         prepared = integrity.prepare_messages(messages)
         request = StreamRequest(
@@ -403,7 +452,7 @@ async def _stream(
             output_type,
             params,
         )
-        s = Stream(executor._do_stream(request))
+        s = Stream(executor._do_stream(request), output_type=output_type)
     try:
         yield s
     finally:

src/ai/types/builders.py

@@ -18,7 +18,6 @@
     Message,
     Part,
     ReasoningPart,
-    StructuredOutputPart,
     TextPart,
     ToolCallPart,
     ToolResultPart,
@@ -30,7 +29,6 @@
     ToolResultPart,
     ReasoningPart,
     HookPart,
-    StructuredOutputPart,
     FilePart,
 )
 

src/ai/types/events.py

@@ -206,8 +206,8 @@ async def replay_message_events(
     not recoverable from a complete message.
 
     Parts with no model-layer event analog — ``ToolResultPart``,
-    ``HookPart``, ``StructuredOutputPart`` — are skipped silently; they
-    are agent-layer concerns and never appear on the model stream.
+    ``HookPart`` — are skipped silently; they are agent-layer concerns
+    and never appear on the model stream.
     """
     yield StreamStart()
     for part in msg.parts: