WIP: With OutputPart

Author: DouweM

pydantic_ai_slim/pydantic_ai/__init__.py

@@ -12,7 +12,7 @@
 )
 from .format_prompt import format_as_xml
 from .messages import AudioUrl, BinaryContent, DocumentUrl, ImageUrl, VideoUrl
-from .result import ToolOutput
+from .result import JSONSchemaOutput, ToolOutput
 from .tools import RunContext, Tool
 
 __all__ = (
@@ -43,6 +43,7 @@
     'RunContext',
     # result
     'ToolOutput',
+    'JSONSchemaOutput',
     # format_prompt
     'format_as_xml',
 )

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -24,7 +24,7 @@
     result,
     usage as _usage,
 )
-from .result import OutputDataT, ToolOutput
+from .result import OutputDataT
 from .settings import ModelSettings, merge_model_settings
 from .tools import RunContext, Tool, ToolDefinition
 
@@ -240,10 +240,29 @@ async def add_mcp_server_tools(server: MCPServer) -> None:
     )
 
     output_schema = ctx.deps.output_schema
+    model = ctx.deps.model
+
+    output_mode = None
+    output_object = None
+    output_tools = []
+    allow_text_output = _output.allow_text_output(output_schema)
+    if output_schema:
+        output_mode = output_schema.forced_mode or model.default_output_mode
+        output_object = output_schema.object_schema.definition
+        output_tools = output_schema.tool_defs()
+        if output_mode != 'tool':
+            allow_text_output = False
+
+        supported_modes = model.supported_output_modes
+        if output_mode not in supported_modes:
+            raise exceptions.UserError(f"Output mode '{output_mode}' is not among supported modes: {supported_modes}")
+
     return models.ModelRequestParameters(
         function_tools=function_tool_defs,
-        allow_text_output=allow_text_output(output_schema),
-        output_tools=output_schema.tool_defs() if output_schema is not None else [],
+        output_mode=output_mode,
+        output_object=output_object,
+        output_tools=output_tools,
+        allow_text_output=allow_text_output,
     )
 
 
@@ -394,20 +413,24 @@ async def stream(
         async for _event in stream:
             pass
 
-    async def _run_stream(
+    async def _run_stream(  # noqa: C901
         self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]
     ) -> AsyncIterator[_messages.HandleResponseEvent]:
         if self._events_iterator is None:
             # Ensure that the stream is only run once
 
             async def _run_stream() -> AsyncIterator[_messages.HandleResponseEvent]:
                 texts: list[str] = []
+                outputs: list[str] = []
                 tool_calls: list[_messages.ToolCallPart] = []
                 for part in self.model_response.parts:
                     if isinstance(part, _messages.TextPart):
                         # ignore empty content for text parts, see #437
                         if part.content:
                             texts.append(part.content)
+                    elif isinstance(part, _messages.OutputPart):
+                        if part.content:
+                            outputs.append(part.content)
                     elif isinstance(part, _messages.ToolCallPart):
                         tool_calls.append(part)
                     else:
@@ -420,6 +443,9 @@ async def _run_stream() -> AsyncIterator[_messages.HandleResponseEvent]:
                 if tool_calls:
                     async for event in self._handle_tool_calls(ctx, tool_calls):
                         yield event
+                elif outputs:  # TODO: Can we have tool calls and structured output? Should we handle both?
+                    # No events are emitted during the handling of structured outputs, so we don't need to yield anything
+                    self._next_node = await self._handle_outputs(ctx, outputs)
                 elif texts:
                     # No events are emitted during the handling of text responses, so we don't need to yield anything
                     self._next_node = await self._handle_text_response(ctx, texts)
@@ -428,7 +454,7 @@ async def _run_stream() -> AsyncIterator[_messages.HandleResponseEvent]:
                     # when the model has already returned text along side tool calls
                     # in this scenario, if text responses are allowed, we return text from the most recent model
                     # response, if any
-                    if allow_text_output(ctx.deps.output_schema):
+                    if _output.allow_text_output(ctx.deps.output_schema):
                         for message in reversed(ctx.state.message_history):
                             if isinstance(message, _messages.ModelResponse):
                                 last_texts = [p.content for p in message.parts if isinstance(p, _messages.TextPart)]
@@ -511,7 +537,7 @@ async def _handle_text_response(
         output_schema = ctx.deps.output_schema
 
         text = '\n\n'.join(texts)
-        if allow_text_output(output_schema):
+        if _output.allow_text_output(output_schema):
             # The following cast is safe because we know `str` is an allowed result type
             result_data_input = cast(NodeRunEndT, text)
             try:
@@ -533,6 +559,27 @@ async def _handle_text_response(
                 )
             )
 
+    async def _handle_outputs(
+        self,
+        ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
+        outputs: list[str],
+    ) -> ModelRequestNode[DepsT, NodeRunEndT] | End[result.FinalResult[NodeRunEndT]]:
+        if len(outputs) != 1:
+            raise exceptions.UnexpectedModelBehavior('Received multiple structured outputs in a single response')
+        output_schema = ctx.deps.output_schema
+        if not output_schema:
+            raise exceptions.UnexpectedModelBehavior('Must specify a non-str result_type when using structured outputs')
+
+        structured_output = outputs[0]
+        try:
+            result_data = output_schema.validate(structured_output)
+            result_data = await _validate_output(result_data, ctx, None)
+        except _output.ToolRetryError as e:
+            ctx.state.increment_retries(ctx.deps.max_result_retries)
+            return ModelRequestNode[DepsT, NodeRunEndT](_messages.ModelRequest(parts=[e.tool_retry]))
+        else:
+            return self._handle_final_result(ctx, result.FinalResult(result_data, None, None), [])
+
 
 def build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, Any]]) -> RunContext[DepsT]:
     """Build a `RunContext` object from the current agent graph run context."""
@@ -773,11 +820,6 @@ async def _validate_output(
     return result_data
 
 
-def allow_text_output(output_schema: _output.OutputSchema[Any] | None) -> bool:
-    """Check if the result schema allows text results."""
-    return output_schema is None or output_schema.allow_text_output
-
-
 @dataclasses.dataclass
 class _RunMessages:
     messages: list[_messages.ModelMessage]
@@ -827,7 +869,9 @@ def get_captured_run_messages() -> _RunMessages:
 
 
 def build_agent_graph(
-    name: str | None, deps_type: type[DepsT], output_type: type[OutputT] | ToolOutput[OutputT]
+    name: str | None,
+    deps_type: type[DepsT],
+    output_type: _output.OutputType[OutputT],
 ) -> Graph[GraphAgentState, GraphAgentDeps[DepsT, result.FinalResult[OutputT]], result.FinalResult[OutputT]]:
     """Build the execution [Graph][pydantic_graph.Graph] for a given agent."""
     nodes = (