WIP: Support structured and manual JSON output_type modes in addition to tool calls

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 StructuredOutput, ToolOutput
 from .tools import RunContext, Tool
 
 __all__ = (
@@ -43,6 +43,7 @@
     'RunContext',
     # result
     'ToolOutput',
+    'StructuredOutput',
     # 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, StructuredOutput, ToolOutput
 from .settings import ModelSettings, merge_model_settings
 from .tools import RunContext, Tool, ToolDefinition
 
@@ -125,9 +125,6 @@ def is_agent_node(
 class UserPromptNode(AgentNode[DepsT, NodeRunEndT]):
     user_prompt: str | Sequence[_messages.UserContent] | None
 
-    instructions: str | None
-    instructions_functions: list[_system_prompt.SystemPromptRunner[DepsT]]
-
     system_prompts: tuple[str, ...]
     system_prompt_functions: list[_system_prompt.SystemPromptRunner[DepsT]]
     system_prompt_dynamic_functions: dict[str, _system_prompt.SystemPromptRunner[DepsT]]
@@ -244,6 +241,8 @@ async def add_mcp_server_tools(server: MCPServer) -> None:
         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_schema=output_schema.json_schema if output_schema is not None else None,
+        preferred_output_mode=output_schema.preferred_mode if output_schema is not None else None,
     )
 
 
@@ -396,20 +395,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] = []
+                structured_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.StructuredOutputPart):
+                        if part.content:
+                            structured_outputs.append(part.content)
                     elif isinstance(part, _messages.ToolCallPart):
                         tool_calls.append(part)
                     else:
@@ -422,6 +425,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 structured_outputs:
+                    # No events are emitted during the handling of structured outputs, so we don't need to yield anything
+                    self._next_node = await self._handle_structured_outputs(ctx, structured_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)
@@ -535,6 +541,27 @@ async def _handle_text_response(
                 )
             )
 
+    async def _handle_structured_outputs(
+        self,
+        ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
+        structured_outputs: list[str],
+    ) -> ModelRequestNode[DepsT, NodeRunEndT] | End[result.FinalResult[NodeRunEndT]]:
+        if len(structured_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 = structured_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."""
@@ -829,7 +856,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: type[OutputT] | ToolOutput[OutputT] | StructuredOutput[OutputT],
 ) -> Graph[GraphAgentState, GraphAgentDeps[DepsT, result.FinalResult[OutputT]], result.FinalResult[OutputT]]:
     """Build the execution [Graph][pydantic_graph.Graph] for a given agent."""
     nodes = (

pydantic_ai_slim/pydantic_ai/_output.py

@@ -12,8 +12,15 @@
 
 from . import _utils, messages as _messages
 from .exceptions import ModelRetry
-from .result import DEFAULT_OUTPUT_TOOL_NAME, OutputDataT, OutputDataT_inv, OutputValidatorFunc, ToolOutput
-from .tools import AgentDepsT, GenerateToolJsonSchema, RunContext, ToolDefinition
+from .result import (
+    DEFAULT_OUTPUT_TOOL_NAME,
+    OutputDataT,
+    OutputDataT_inv,
+    OutputValidatorFunc,
+    StructuredOutput,
+    ToolOutput,
+)
+from .tools import AgentDepsT, GenerateToolJsonSchema, ObjectJsonSchema, RunContext, ToolDefinition
 
 T = TypeVar('T')
 """An invariant TypeVar."""
@@ -83,13 +90,18 @@ class OutputSchema(Generic[OutputDataT]):
     Similar to `Tool` but for the final output of running an agent.
     """
 
+    # TODO: Since this is currently called "preferred", models that don't have structured output implemented yet ignore it and use tools (except for Mistral).
+    # We should likely raise an error if an unsupported mode is used, _and_ allow the model to pick its own preferred mode if none is forced.
+    preferred_mode: Literal['tool', 'structured'] | None  # TODO: Add mode for manual JSON
+    type_adapter: TypeAdapter[OutputDataT]
     tools: dict[str, OutputSchemaTool[OutputDataT]]
-    allow_text_output: bool
+    allow_text_output: bool  # TODO: Verify structured output works correctly with string as a union member
+    json_schema: ObjectJsonSchema  # TODO: Verify structured output works correctly with a union
 
     @classmethod
     def build(
         cls: type[OutputSchema[T]],
-        output_type: type[T] | ToolOutput[T],
+        output_type: type[T] | ToolOutput[T] | StructuredOutput[T],  # TODO: Support a list of output types/markers
         name: str | None = None,
         description: str | None = None,
         strict: bool | None = None,
@@ -98,15 +110,34 @@ def build(
         if output_type is str:
             return None
 
+        preferred_mode = None
         if isinstance(output_type, ToolOutput):
             # do we need to error on conflicts here? (DavidM): If this is internal maybe doesn't matter, if public, use overloads
             name = output_type.name
             description = output_type.description
             output_type_ = output_type.output_type
             strict = output_type.strict
+            preferred_mode = 'tool'
+        elif isinstance(output_type, StructuredOutput):
+            name = output_type.name  # TODO: Get this to the response_format model request arg
+            description = output_type.description  # TODO: Get this to the response_format model request arg
+            output_type_ = output_type.output_type
+            strict = output_type.strict  # TODO: Get this to the response_format model request arg
+            preferred_mode = 'structured'
         else:
             output_type_ = output_type
 
+        type_adapter = cast(TypeAdapter[T], TypeAdapter(output_type_))
+        json_schema = _utils.check_object_json_schema(type_adapter.json_schema(schema_generator=GenerateToolJsonSchema))
+
+        # TODO: Make this description available to the model params
+        if json_schema_description := json_schema.pop('description', None):
+            if description is None:
+                description = json_schema_description
+            else:
+                description = f'{description}. {json_schema_description}'
+
+        # No need to include an output tool for string output
         if output_type_option := extract_str_from_union(output_type):
             output_type_ = output_type_option.value
             allow_text_output = True
@@ -134,7 +165,13 @@ def build(
                 ),
             )
 
-        return cls(tools=tools, allow_text_output=allow_text_output)
+        return cls(
+            preferred_mode=preferred_mode,
+            tools=tools,
+            allow_text_output=allow_text_output,
+            type_adapter=type_adapter,
+            json_schema=json_schema,
+        )
 
     def find_named_tool(
         self, parts: Iterable[_messages.ModelResponsePart], tool_name: str
@@ -163,6 +200,35 @@ def tool_defs(self) -> list[ToolDefinition]:
         """Get tool definitions to register with the model."""
         return [t.tool_def for t in self.tools.values()]
 
+    def validate(
+        self, output_text: str, allow_partial: bool = False, wrap_validation_errors: bool = True
+    ) -> OutputDataT:
+        """Validate a structured output message.
+
+        Args:
+            output_text: The structured output from the LLM to validate.
+            allow_partial: If true, allow partial validation.
+            wrap_validation_errors: If true, wrap the validation errors in a retry message.
+
+        Returns:
+            Either the validated output data (left) or a retry message (right).
+        """
+        try:
+            pyd_allow_partial: Literal['off', 'trailing-strings'] = 'trailing-strings' if allow_partial else 'off'
+            output = self.type_adapter.validate_json(output_text, experimental_allow_partial=pyd_allow_partial)
+        except ValidationError as e:
+            if wrap_validation_errors:
+                m = _messages.RetryPromptPart(
+                    content=e.errors(include_url=False),
+                )
+                raise ToolRetryError(m) from e
+            else:
+                raise
+        else:
+            return output
+
+    # TODO: Build instructions for manual JSON
+
 
 DEFAULT_DESCRIPTION = 'The final response which ends this conversation'
 

pydantic_ai_slim/pydantic_ai/_parts_manager.py

@@ -23,6 +23,7 @@
     ModelResponseStreamEvent,
     PartDeltaEvent,
     PartStartEvent,
+    StructuredOutputPartDelta,
     TextPart,
     TextPartDelta,
     ToolCallPart,
@@ -57,12 +58,12 @@ class ModelResponsePartsManager:
     """Maps a vendor's "part" ID (if provided) to the index in `_parts` where that part resides."""
 
     def get_parts(self) -> list[ModelResponsePart]:
-        """Return only model response parts that are complete (i.e., not ToolCallPartDelta's).
+        """Return only model response parts that are complete (i.e., not ToolCallPartDelta's or StructuredOutputPartDelta's).
 
         Returns:
-            A list of ModelResponsePart objects. ToolCallPartDelta objects are excluded.
+            A list of ModelResponsePart objects. ToolCallPartDelta and StructuredOutputPartDelta objects are excluded.
         """
-        return [p for p in self._parts if not isinstance(p, ToolCallPartDelta)]
+        return [p for p in self._parts if not isinstance(p, (ToolCallPartDelta, StructuredOutputPartDelta))]
 
     def handle_text_delta(
         self,
@@ -91,6 +92,8 @@ def handle_text_delta(
         """
         existing_text_part_and_index: tuple[TextPart, int] | None = None
 
+        # TODO: Parse out structured output or manual JSON, with a separate message?
+
         if vendor_part_id is None:
             # If the vendor_part_id is None, check if the latest part is a TextPart to update
             if self._parts: