feat(agents): [1/n] return typed answers from session streams

helpers.md

@@ -604,3 +604,32 @@ lookup = pydantic_function_tool(
 ```
 
 Callbacks can be async when used with `AsyncOpenAI`. Existing dictionary handlers still work.
+
+### Typed Agents output (beta)
+
+Pass a Pydantic model (or a Pydantic v2 dataclass) to generate the Agents output
+schema and parse the completed answer. Schemas use the same normalization as
+Responses; the API validates which schema features it supports.
+
+```python
+from pydantic import BaseModel
+
+class Report(BaseModel):
+    summary: str
+    findings: list[str]
+
+with client.beta.agents.sessions.create(
+    agent={"model": MODEL}, environment={"type": "none"},
+    input="Summarize the findings.", stream=True, output_type=Report,
+) as stream:
+    result = stream.get_final_result()
+print(result.output_parsed)
+```
+
+For a session already configured with that schema, use
+`sessions.stream(session_id, input="Update the report.", output_type=Report)`.
+This only selects the local parser; it does not change the session's schema.
+`output_parsed` exposes the first parsed final text part; every final text part is validated.
+`result.parse(Report)` parses an existing raw result. `AgentOutputParseError.result`
+retains the completed raw answer if validation fails. With `AsyncOpenAI`, await
+creation and the result getter, and use `async with`.

src/openai/lib/beta/agents/__init__.py

@@ -5,7 +5,12 @@
     function_tool as function_tool,
     pydantic_function_tool as pydantic_function_tool,
 )
-from ._result import AgentTurnResult as AgentTurnResult, AgentTurnResultError as AgentTurnResultError
+from ._output import agent_text_format as agent_text_format
+from ._result import (
+    AgentTurnResult as AgentTurnResult,
+    AgentTurnResultError as AgentTurnResultError,
+    AgentOutputParseError as AgentOutputParseError,
+)
 from ._stream import (
     AgentSessionEventStream as AgentSessionEventStream,
     AsyncAgentSessionEventStream as AsyncAgentSessionEventStream,

src/openai/lib/beta/agents/_output.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from typing import Any, cast
+from typing_extensions import TypeVar
+
+import pydantic
+
+from ._schema import model_schema
+from ...._types import Omit
+from ...._compat import PYDANTIC_V1
+from ..._pydantic import is_basemodel_type, is_dataclass_like_type
+from ....types.beta.agent_text_param import AgentTextParam
+from ....types.beta.text_format_param import TextFormatParamJSONSchema
+from ....types.beta.agents.session_create_params import Agent
+
+
+def validate_output_type(output_type: type[Any]) -> None:
+    if not is_basemodel_type(output_type) and not (is_dataclass_like_type(output_type) and not PYDANTIC_V1):
+        raise TypeError("Agents output_type must be a Pydantic model or a Pydantic v2 dataclass")
+
+
+def agent_text_format(output_type: type[Any]) -> TextFormatParamJSONSchema:
+    """Beta: build the Agents JSON-schema format for a Pydantic model."""
+    validate_output_type(output_type)
+    if is_basemodel_type(output_type):
+        schema = model_schema(output_type, strict=True)
+    else:
+        schema = model_schema(pydantic.TypeAdapter(output_type), strict=True)
+    return {"type": "json_schema", "schema": schema}
+
+
+def with_output_schema(agent: Agent | Omit, output_type: type[Any] | None) -> Agent | Omit:
+    if output_type is None:
+        return agent
+    config = cast(Agent, {} if isinstance(agent, Omit) else dict(agent))
+    text = cast(AgentTextParam, dict(config.get("text") or {}))
+    if text.get("format") is not None:
+        raise ValueError("Pass output_type or agent.text.format, not both")
+    text["format"] = agent_text_format(output_type)
+    config["text"] = text
+    return config
+
+
+ResponseT = TypeVar("ResponseT")
+
+
+def bind_output_type(response: ResponseT, output_type: type[Any] | None) -> ResponseT:
+    from ._stream import AgentSessionEventStream, AsyncAgentSessionEventStream
+
+    if isinstance(response, (AgentSessionEventStream, AsyncAgentSessionEventStream)):
+        stream = cast("AgentSessionEventStream[Any] | AsyncAgentSessionEventStream[Any]", response)
+        stream._collection.output_type = output_type
+    return cast(ResponseT, response)

src/openai/lib/beta/agents/_result.py

@@ -1,23 +1,50 @@
 from __future__ import annotations
 
 from copy import deepcopy
-from typing import Iterable
+from typing import Any, Generic, Iterable
 from dataclasses import dataclass
-from typing_extensions import Literal
+from typing_extensions import Literal, TypeVar
 
+from ._output import validate_output_type
 from ...._exceptions import OpenAIError
+from ..._parsing._completions import _parse_content
 from ....types.beta.agent_session import RequiredAction
 from ....types.beta.agent_session_event import AgentSessionEvent
 from ....types.beta.agents.sessions.turn import Turn
 from ....types.beta.agent_session_message import AgentSessionMessage
 
+OutputT = TypeVar("OutputT", default=Any)
+ParseT = TypeVar("ParseT")
+
 
 @dataclass(frozen=True)
-class AgentTurnResult:
+class AgentTurnResult(Generic[OutputT]):
     """Beta: the completed final assistant messages from one successful root turn."""
 
     turn: Turn
     messages: list[AgentSessionMessage]
+    output_parsed: OutputT | None = None
+
+    def parse(self, output_type: type[ParseT]) -> AgentTurnResult[ParseT]:
+        """Beta: parse a completed answer without changing its session configuration."""
+        try:
+            validate_output_type(output_type)
+            parsed: ParseT | None = None
+            found = False
+            for message in self.messages:
+                for content in message.content:
+                    if content.type != "output_text":
+                        continue
+                    value = _parse_content(output_type, content.text)
+                    if not found:
+                        parsed = value
+                        found = True
+            if not found:
+                raise ValueError("No final output text to parse")
+        except Exception:
+            # Pydantic errors may include response text in their rendered message.
+            raise AgentOutputParseError(self) from None
+        return AgentTurnResult(turn=self.turn, messages=self.messages, output_parsed=parsed)
 
     @property
     def session_id(self) -> str:
@@ -33,6 +60,14 @@ def output_text(self) -> str:
         return "".join(message.output_text for message in self.messages)
 
 
+class AgentOutputParseError(OpenAIError):
+    """Beta: output parsing failed after a successful hosted turn."""
+
+    def __init__(self, result: AgentTurnResult) -> None:
+        super().__init__("Could not parse the completed agent output")
+        self.result = result
+
+
 ResultErrorReason = Literal["failed", "cancelled", "requires_action", "incomplete", "observation_failed"]
 
 
@@ -160,10 +195,14 @@ def result(self) -> AgentTurnResult:
         return self._result
 
 
-class AgentTurnResultCollection:
+class AgentTurnResultCollection(Generic[OutputT]):
     """Keep ordinary event iteration incremental until collection is requested."""
 
-    def __init__(self, session_id: str | None = None) -> None:
+    def __init__(self, session_id: str | None = None, output_type: type[OutputT] | None = None) -> None:
+        if output_type is not None:
+            validate_output_type(output_type)
+        self.output_type = output_type
+        self._parsed_result: AgentTurnResult[OutputT] | None = None
         self.collector: AgentTurnResultCollector | None = None
         self._session_id = session_id
         self._started = False
@@ -185,3 +224,11 @@ def accept(self, event: AgentSessionEvent) -> None:
     def record_error(self, error: Exception) -> None:
         if self.collector is not None and not self.collector.is_done():
             self.collector.cause = error
+
+    def result(self) -> AgentTurnResult[OutputT]:
+        result = self.enable().result()
+        if self.output_type is None:
+            return result
+        if self._parsed_result is None:
+            self._parsed_result = result.parse(self.output_type)
+        return self._parsed_result

src/openai/lib/beta/agents/_schema.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from copy import deepcopy
+from typing import Any, cast
+
+import pydantic
+
+from ...._compat import PYDANTIC_V1, model_json_schema
+from ..._pydantic import _ensure_strict_json_schema
+
+
+def model_schema(
+    model: type[pydantic.BaseModel] | pydantic.TypeAdapter[Any], *, strict: bool = False
+) -> dict[str, Any]:
+    # Pydantic v1 caches schema dictionaries. Input and output policies must not
+    # mutate each other's schemas or the model's cache.
+    schema = deepcopy(
+        model.json_schema()
+        if not PYDANTIC_V1 and isinstance(model, pydantic.TypeAdapter)
+        else model_json_schema(cast(type[pydantic.BaseModel], model))
+    )
+    return _ensure_strict_json_schema(schema, path=(), root=schema) if strict else schema

src/openai/lib/beta/agents/_stream.py

@@ -1,15 +1,15 @@
 from __future__ import annotations
 
-from typing import Iterator, AsyncIterator
+from typing import Generic, Iterator, AsyncIterator
 from typing_extensions import Self, override
 
-from ._result import AgentTurnResult, AgentTurnResultError, AgentTurnResultCollection
+from ._result import OutputT, AgentTurnResult, AgentTurnResultError, AgentOutputParseError, AgentTurnResultCollection
 from ...._compat import cached_property
 from ...._streaming import Stream, AsyncStream
 from ....types.beta.agent_session_event import AgentSessionEvent
 
 
-class AgentSessionEventStream(Stream[AgentSessionEvent]):
+class AgentSessionEventStream(Stream[AgentSessionEvent], Generic[OutputT]):
     """Beta: a creation event stream that can collect its initial turn's result.
 
     Event iteration and response access behave like ``Stream``. Collection does
@@ -18,7 +18,7 @@ class AgentSessionEventStream(Stream[AgentSessionEvent]):
     """
 
     @cached_property
-    def _collection(self) -> AgentTurnResultCollection:
+    def _collection(self) -> AgentTurnResultCollection[OutputT]:
         return AgentTurnResultCollection()
 
     @override
@@ -36,7 +36,7 @@ def with_result_collection(self) -> Self:
         self._collection.enable()
         return self
 
-    def get_final_result(self) -> AgentTurnResult:
+    def get_final_result(self) -> AgentTurnResult[OutputT]:
         """Consume remaining events and return the successful initial turn result.
 
         Raises AgentTurnResultError for an unsuccessful or unobserved outcome.
@@ -52,8 +52,8 @@ def get_final_result(self) -> AgentTurnResult:
                     collector.check_outcome()
                     if collector.is_done():
                         break
-            return collector.result()
-        except AgentTurnResultError:
+            return self._collection.result()
+        except (AgentTurnResultError, AgentOutputParseError):
             raise
         except Exception as error:
             collector.cause = error
@@ -62,11 +62,11 @@ def get_final_result(self) -> AgentTurnResult:
             self.close()
 
 
-class AsyncAgentSessionEventStream(AsyncStream[AgentSessionEvent]):
+class AsyncAgentSessionEventStream(AsyncStream[AgentSessionEvent], Generic[OutputT]):
     """Beta: asynchronous counterpart of AgentSessionEventStream."""
 
     @cached_property
-    def _collection(self) -> AgentTurnResultCollection:
+    def _collection(self) -> AgentTurnResultCollection[OutputT]:
         return AgentTurnResultCollection()
 
     @override
@@ -84,7 +84,7 @@ def with_result_collection(self) -> Self:
         self._collection.enable()
         return self
 
-    async def get_final_result(self) -> AgentTurnResult:
+    async def get_final_result(self) -> AgentTurnResult[OutputT]:
         """Consume remaining events and return the successful initial turn result.
 
         Enables collection on a fresh stream. To iterate first, call
@@ -98,8 +98,8 @@ async def get_final_result(self) -> AgentTurnResult:
                     collector.check_outcome()
                     if collector.is_done():
                         break
-            return collector.result()
-        except AgentTurnResultError:
+            return self._collection.result()
+        except (AgentTurnResultError, AgentOutputParseError):
             raise
         except Exception as error:
             collector.cause = error

src/openai/lib/beta/agents/_tools.py

@@ -11,8 +11,9 @@
 
 import pydantic
 
+from ._schema import model_schema
 from ...._utils import is_dict
-from ...._compat import PYDANTIC_V1, model_dump, model_json, model_parse, model_json_schema
+from ...._compat import PYDANTIC_V1, model_dump, model_json, model_parse
 from ..._pydantic import resolve_ref
 from ...streaming.agents._types import ToolOutput
 from ....types.beta.agent_tool_param import AgentToolConfigParamFunction
@@ -39,7 +40,7 @@ def __init__(
     ) -> None:
         if not name:
             raise ValueError("Tool name must not be empty")
-        parameters = model_json_schema(model)
+        parameters = model_schema(model)
         root = parameters
         seen: set[str] = set()
         while isinstance(ref := root.get("$ref"), str) and ref not in seen: