From 0d4111e3dbea5c13a29db7ca58816750bfdf9e07 Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Mon, 11 May 2026 21:01:54 -0700 Subject: [PATCH 1/5] Support output_type on stream For now, it is dropped from `message` and just added to `Stream`, and we drop StructuredMessagePart. We could add something like that back if we really wanted, but I'm not sure it would be worth it. The parsing would still probably be located in Stream, since that is where Message construction happens now. --- examples/samples/structured_output.py | 20 +++++++-- src/ai/models/core/api.py | 25 +++++++++-- src/ai/types/builders.py | 2 - src/ai/types/events.py | 4 +- src/ai/types/messages.py | 60 --------------------------- tests/conftest.py | 2 - tests/types/test_integrity.py | 2 +- tests/types/test_messages.py | 59 -------------------------- 8 files changed, 41 insertions(+), 133 deletions(-) diff --git a/examples/samples/structured_output.py b/examples/samples/structured_output.py index 406dc5bb..eeaa8a3a 100644 --- a/examples/samples/structured_output.py +++ b/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. + recipe = s.output + assert isinstance(recipe, Recipe) + 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__": diff --git a/src/ai/models/core/api.py b/src/ai/models/core/api.py index bc672913..5b93e318 100644 --- a/src/ai/models/core/api.py +++ b/src/ai/models/core/api.py @@ -84,6 +84,7 @@ def __init__( gen: AsyncGenerator[types.events.Event], *, seed_message: types.messages.Message | None = None, + output_type: type[pydantic.BaseModel] | None = None, ) -> None: """Wrap an event generator. @@ -93,12 +94,18 @@ 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] = {} + self._output_type = output_type async def aclose(self) -> None: await self._gen.aclose() @@ -146,7 +153,15 @@ def tool_calls(self) -> list[types.messages.ToolCallPart]: @property def output(self) -> Any: - return self._message.output + """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. + """ + if self._output_type is None: + return self._message.text + return self._output_type.model_validate_json(self._message.text) def _aggregate_event(self, event: types.events.Event) -> dict[str, Any]: updates: dict[str, Any] = {} @@ -393,7 +408,11 @@ async def _stream( ) -> AsyncIterator[Stream]: 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 +422,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: diff --git a/src/ai/types/builders.py b/src/ai/types/builders.py index 9b80e14b..174ca10d 100644 --- a/src/ai/types/builders.py +++ b/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, ) diff --git a/src/ai/types/events.py b/src/ai/types/events.py index 14e6cff3..e5916752 100644 --- a/src/ai/types/events.py +++ b/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: diff --git a/src/ai/types/messages.py b/src/ai/types/messages.py index 66bfa1d1..56b1da8c 100644 --- a/src/ai/types/messages.py +++ b/src/ai/types/messages.py @@ -1,4 +1,3 @@ -import importlib import uuid from typing import Annotated, Any, Literal, Self @@ -110,56 +109,6 @@ class HookPart[T](pydantic.BaseModel): model_config = pydantic.ConfigDict(frozen=True) -# todo redo this structured output situation and simplify it -def _resolve_class(fully_qualified_name: str) -> type[pydantic.BaseModel]: - """Import and return a class from its fully qualified name. - - E.g. ``"myapp.models.WeatherForecast"`` → the ``WeatherForecast`` class. - """ - module_path, _, class_name = fully_qualified_name.rpartition(".") - if not module_path: - raise ImportError( - f"Cannot resolve '{fully_qualified_name}': " - "expected a fully qualified name like 'mypackage.module.ClassName'" - ) - module = importlib.import_module(module_path) - cls = getattr(module, class_name, None) - if cls is None: - raise ImportError(f"Module '{module_path}' has no attribute '{class_name}'") - if not (isinstance(cls, type) and issubclass(cls, pydantic.BaseModel)): - raise TypeError( - f"'{fully_qualified_name}' is not a pydantic.BaseModel subclass" - ) - return cls - - -class StructuredOutputPart(pydantic.BaseModel): - """Part containing a validated structured output from the LLM. - - ``data`` stores the parsed JSON dict (always serializable). - ``output_type_name`` stores the fully qualified class name so the typed - Pydantic model can be lazily rehydrated via the ``value`` property. - """ - - model_config = pydantic.ConfigDict(frozen=True) - - id: str = pydantic.Field(default_factory=generate_id) - data: dict[str, Any] - output_type_name: str - kind: Literal["structured_output"] = "structured_output" - provider_metadata: dict[str, Any] | None = None - - _hydrated: Any = pydantic.PrivateAttr(default=None) - - @property - def value(self) -> Any: - """Lazily resolve the output type class and validate ``data`` into it.""" - if self._hydrated is None: - cls = _resolve_class(self.output_type_name) - self._hydrated = cls.model_validate(self.data) - return self._hydrated - - class FilePart(pydantic.BaseModel): """File, image, or audio content part. @@ -229,7 +178,6 @@ def from_bytes( | BuiltinToolReturnPart | ReasoningPart | HookPart[Any] - | StructuredOutputPart | FilePart, pydantic.Field(discriminator="kind"), ] @@ -288,11 +236,3 @@ def images(self) -> list[FilePart]: @property def videos(self) -> list[FilePart]: return [p for p in self.files if p.media_type.startswith("video/")] - - @property - def output(self) -> Any: - """Parsed structured output from the first structured-output part.""" - for part in self.parts: - if isinstance(part, StructuredOutputPart): - return part.value - return None diff --git a/tests/conftest.py b/tests/conftest.py index e76275af..f52490a4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -139,8 +139,6 @@ async def emit_events_for_messages( media_type=part.media_type, data=part.data if isinstance(part.data, str) else "", ) - # StructuredOutputPart is not a streamed part; tests that need it - # construct a tailored adapter directly. yield events_.StreamEnd(usage=usage) diff --git a/tests/types/test_integrity.py b/tests/types/test_integrity.py index 85fe680e..cf8ac5d9 100644 --- a/tests/types/test_integrity.py +++ b/tests/types/test_integrity.py @@ -120,7 +120,7 @@ def test_internal_strict_raises() -> None: # --------------------------------------------------------------------------- -# Internal parts (HookPart, StructuredOutputPart) +# Internal parts (HookPart) # --------------------------------------------------------------------------- diff --git a/tests/types/test_messages.py b/tests/types/test_messages.py index d4d1580d..8da5e148 100644 --- a/tests/types/test_messages.py +++ b/tests/types/test_messages.py @@ -2,70 +2,11 @@ from __future__ import annotations -import pydantic import pytest from ai.types import messages, usage -class _Weather(pydantic.BaseModel): - city: str - temperature: float - - -_WEATHER_DATA = {"city": "SF", "temperature": 62.0} -_WEATHER_TYPE_NAME = f"{_Weather.__module__}.{_Weather.__qualname__}" - - -def test_structured_output_part_value() -> None: - part = messages.StructuredOutputPart( - data=_WEATHER_DATA, output_type_name=_WEATHER_TYPE_NAME - ) - val = part.value - assert isinstance(val, _Weather) - assert val.city == "SF" - assert part.value is val - - -def test_structured_output_part_bad_class_name() -> None: - part = messages.StructuredOutputPart( - data=_WEATHER_DATA, output_type_name="nonexistent.module.Cls" - ) - with pytest.raises(ImportError): - _ = part.value - - -def test_message_output_from_part() -> None: - m = messages.Message( - id="m1", - role="assistant", - parts=[ - messages.TextPart(text="{}"), - messages.StructuredOutputPart( - data=_WEATHER_DATA, output_type_name=_WEATHER_TYPE_NAME - ), - ], - ) - assert isinstance(m.output, _Weather) - assert m.output.city == "SF" - - -def test_structured_output_round_trip() -> None: - m = messages.Message( - id="m1", - role="assistant", - parts=[ - messages.TextPart(text="{}"), - messages.StructuredOutputPart( - data=_WEATHER_DATA, output_type_name=_WEATHER_TYPE_NAME - ), - ], - ) - restored = messages.Message.model_validate(m.model_dump()) - assert isinstance(restored.output, _Weather) - assert restored.output.city == "SF" - - def test_usage_add_merges_optional_fields() -> None: a = usage.Usage( input_tokens=100, From fc577cff1575b3567982614b8af4699198b375fc Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Mon, 11 May 2026 21:18:08 -0700 Subject: [PATCH 2/5] Give Stream a type parameter so it works nicer --- examples/samples/structured_output.py | 4 +- pyproject.toml | 1 + src/ai/models/core/api.py | 58 ++++++++++++++++++++------- uv.lock | 2 + 4 files changed, 49 insertions(+), 16 deletions(-) diff --git a/examples/samples/structured_output.py b/examples/samples/structured_output.py index eeaa8a3a..ae23fd91 100644 --- a/examples/samples/structured_output.py +++ b/examples/samples/structured_output.py @@ -29,9 +29,9 @@ async def main() -> None: print() # After iteration, s.output validates the streamed JSON against - # the output_type the stream was opened with. + # the output_type the stream was opened with. ``Stream[Recipe]`` + # types ``s.output`` as ``Recipe``. recipe = s.output - assert isinstance(recipe, Recipe) print(f"\n\nParsed recipe: {recipe.name}") print(f" Ingredients: {', '.join(recipe.ingredients)}") print(f" Prep time: {recipe.prep_time_minutes} min") diff --git a/pyproject.toml b/pyproject.toml index b5577ed6..61197f36 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/src/ai/models/core/api.py b/src/ai/models/core/api.py index 5b93e318..af19bc1a 100644 --- a/src/ai/models/core/api.py +++ b/src/ai/models/core/api.py @@ -2,10 +2,14 @@ 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 @@ -13,6 +17,11 @@ 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,7 +85,7 @@ 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__( @@ -84,7 +93,7 @@ def __init__( gen: AsyncGenerator[types.events.Event], *, seed_message: types.messages.Message | None = None, - output_type: type[pydantic.BaseModel] | None = None, + output_type: type[StreamOutputT] | None = None, ) -> None: """Wrap an event generator. @@ -105,7 +114,10 @@ def __init__( role="assistant", parts=[] ) self._parts: dict[str, types.messages.Part] = {} - self._output_type = output_type + # ``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() @@ -152,7 +164,7 @@ def tool_calls(self) -> list[types.messages.ToolCallPart]: return self._message.tool_calls @property - def output(self) -> Any: + def output(self) -> StreamOutputT: """Return the streamed output as the ``output_type`` passed in. Defaults to the concatenated message text. When a Pydantic @@ -160,8 +172,10 @@ def output(self) -> Any: it and returns the parsed instance. """ if self._output_type is None: - return self._message.text - return self._output_type.model_validate_json(self._message.text) + return cast(StreamOutputT, self._message.text) + return cast( + StreamOutputT, self._output_type.model_validate_json(self._message.text) + ) def _aggregate_event(self, event: types.events.Event) -> dict[str, Any]: updates: dict[str, Any] = {} @@ -338,20 +352,36 @@ def tools(self) -> list[types.tools.Tool]: ... 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, @@ -361,7 +391,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`. @@ -405,7 +435,7 @@ 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( diff --git a/uv.lock b/uv.lock index 994b964b..0704fae4 100644 --- a/uv.lock +++ b/uv.lock @@ -1036,6 +1036,7 @@ dependencies = [ { name = "mcp" }, { name = "openai" }, { name = "pydantic" }, + { name = "typing-extensions" }, { name = "vercel" }, ] @@ -1058,6 +1059,7 @@ requires-dist = [ { name = "mcp", specifier = ">=1.18.0" }, { name = "openai", specifier = ">=2.14.0" }, { name = "pydantic", specifier = ">=2.12.5" }, + { name = "typing-extensions", specifier = ">=4.15.0" }, { name = "vercel", specifier = ">=0.3.8" }, ] From bfc6e64bd3c81d52490ad7b316124d4e5c9daab4 Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Tue, 12 May 2026 11:28:09 -0700 Subject: [PATCH 3/5] put a get_output on Message --- src/ai/models/core/api.py | 4 +--- src/ai/types/messages.py | 4 ++++ 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/ai/models/core/api.py b/src/ai/models/core/api.py index af19bc1a..bcbab4af 100644 --- a/src/ai/models/core/api.py +++ b/src/ai/models/core/api.py @@ -173,9 +173,7 @@ def output(self) -> StreamOutputT: """ if self._output_type is None: return cast(StreamOutputT, self._message.text) - return cast( - StreamOutputT, self._output_type.model_validate_json(self._message.text) - ) + 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] = {} diff --git a/src/ai/types/messages.py b/src/ai/types/messages.py index 56b1da8c..6ba89304 100644 --- a/src/ai/types/messages.py +++ b/src/ai/types/messages.py @@ -236,3 +236,7 @@ def images(self) -> list[FilePart]: @property def videos(self) -> list[FilePart]: return [p for p in self.files if p.media_type.startswith("video/")] + + def get_output[T: pydantic.BaseModel](self, output_type: type[T]) -> T: + """Parse the message's text content as a JSON instance of ``output_type``.""" + return output_type.model_validate_json(self.text) From e7911cd195c6c4ebbdc22f1b886ec5773a4421e6 Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Tue, 12 May 2026 11:55:05 -0700 Subject: [PATCH 4/5] Implement it for AgentStream also --- src/ai/agents/agent.py | 82 ++++++++++++++++++++++++++++++++++++--- src/ai/models/core/api.py | 6 ++- 2 files changed, 82 insertions(+), 6 deletions(-) diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index 5e773176..296a6bce 100644 --- a/src/ai/agents/agent.py +++ b/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,21 @@ 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] + output_type = self._context.output_type + if output_type is None: + return cast(AgentOutputT, last.text) + return cast(AgentOutputT, last.get_output(output_type)) + def tool_result( *items: types.messages.Message @@ -894,14 +932,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 +970,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 +981,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) diff --git a/src/ai/models/core/api.py b/src/ai/models/core/api.py index bcbab4af..354dc3f2 100644 --- a/src/ai/models/core/api.py +++ b/src/ai/models/core/api.py @@ -332,7 +332,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. @@ -344,6 +344,8 @@ 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 @@ -411,6 +413,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=") From c0bb30fd571e7d862496e23976d605754846fc78 Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Tue, 12 May 2026 13:26:11 -0700 Subject: [PATCH 5/5] raise when the message is not finished, move the logic to Message --- src/ai/agents/agent.py | 5 +---- src/ai/models/core/api.py | 2 -- src/ai/types/messages.py | 26 +++++++++++++++++++++++--- 3 files changed, 24 insertions(+), 9 deletions(-) diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index 296a6bce..01b2f63d 100644 --- a/src/ai/agents/agent.py +++ b/src/ai/agents/agent.py @@ -753,10 +753,7 @@ def output(self) -> AgentOutputT: is returned. """ last = self._context.messages[-1] - output_type = self._context.output_type - if output_type is None: - return cast(AgentOutputT, last.text) - return cast(AgentOutputT, last.get_output(output_type)) + return cast(AgentOutputT, last.get_output(self._context.output_type)) def tool_result( diff --git a/src/ai/models/core/api.py b/src/ai/models/core/api.py index 354dc3f2..022168c5 100644 --- a/src/ai/models/core/api.py +++ b/src/ai/models/core/api.py @@ -171,8 +171,6 @@ def output(self) -> StreamOutputT: model subclass was passed, validates the streamed JSON against it and returns the parsed instance. """ - if self._output_type is None: - return cast(StreamOutputT, self._message.text) return cast(StreamOutputT, self._message.get_output(self._output_type)) def _aggregate_event(self, event: types.events.Event) -> dict[str, Any]: diff --git a/src/ai/types/messages.py b/src/ai/types/messages.py index 6ba89304..5f4271ec 100644 --- a/src/ai/types/messages.py +++ b/src/ai/types/messages.py @@ -1,5 +1,5 @@ import uuid -from typing import Annotated, Any, Literal, Self +from typing import Annotated, Any, Literal, Self, overload import pydantic @@ -237,6 +237,26 @@ def images(self) -> list[FilePart]: def videos(self) -> list[FilePart]: return [p for p in self.files if p.media_type.startswith("video/")] - def get_output[T: pydantic.BaseModel](self, output_type: type[T]) -> T: - """Parse the message's text content as a JSON instance of ``output_type``.""" + @overload + def get_output(self, output_type: None = None) -> str: ... + @overload + def get_output[T: pydantic.BaseModel](self, output_type: type[T]) -> T: ... + def get_output(self, output_type: type[pydantic.BaseModel] | None = None) -> Any: + """Return the final output of this assistant turn. + + With no ``output_type``, returns the concatenated text content. + With a Pydantic model class, validates the text as JSON against + it and returns the parsed instance. + + Raises :class:`ValueError` unless the message is a *final* + assistant message: role ``"assistant"`` with no pending tool calls. + """ + if self.role != "assistant" or self.tool_calls: + raise ValueError( + "get_output() requires a final assistant message " + "(role='assistant' with no tool calls); " + f"got role={self.role!r} with {len(self.tool_calls)} tool call(s)" + ) + if output_type is None: + return self.text return output_type.model_validate_json(self.text)