Support NativeOutput and PromptedOutput modes in addition to ToolOutput (#1628)

Author: DouweM

docs/api/output.md

@@ -0,0 +1,11 @@
+# `pydantic_ai.output`
+
+::: pydantic_ai.output
+    options:
+        inherited_members: true
+        members:
+            - OutputDataT
+            - ToolOutput
+            - NativeOutput
+            - PromptedOutput
+            - TextOutput

docs/api/result.md

@@ -4,5 +4,4 @@
     options:
         inherited_members: true
         members:
-            - OutputDataT
             - StreamedRunResult

docs/output.md

@@ -64,6 +64,7 @@ nav:
       - api/agent.md
       - api/tools.md
       - api/common_tools.md
+      - api/output.md
       - api/result.md
       - api/messages.md
       - api/exceptions.md

mkdocs.yml

@@ -12,7 +12,7 @@
 )
 from .format_prompt import format_as_xml
 from .messages import AudioUrl, BinaryContent, DocumentUrl, ImageUrl, VideoUrl
-from .result import ToolOutput
+from .output import NativeOutput, PromptedOutput, TextOutput, ToolOutput
 from .tools import RunContext, Tool
 
 __all__ = (
@@ -41,8 +41,11 @@
     # tools
     'Tool',
     'RunContext',
-    # result
+    # output
     'ToolOutput',
+    'NativeOutput',
+    'PromptedOutput',
+    'TextOutput',
     # format_prompt
     'format_as_xml',
 )

pydantic_ai_slim/pydantic_ai/__init__.py

@@ -18,7 +18,7 @@
 from pydantic_graph.nodes import End, NodeRunEndT
 
 from . import _output, _system_prompt, exceptions, messages as _messages, models, result, usage as _usage
-from .result import OutputDataT
+from .output import OutputDataT, OutputSpec
 from .settings import ModelSettings, merge_model_settings
 from .tools import RunContext, Tool, ToolDefinition, ToolsPrepareFunc
 
@@ -102,7 +102,7 @@ class GraphAgentDeps(Generic[DepsT, OutputDataT]):
     end_strategy: EndStrategy
     get_instructions: Callable[[RunContext[DepsT]], Awaitable[str | None]]
 
-    output_schema: _output.OutputSchema[OutputDataT] | None
+    output_schema: _output.OutputSchema[OutputDataT]
     output_validators: list[_output.OutputValidator[DepsT, OutputDataT]]
 
     history_processors: Sequence[HistoryProcessor[DepsT]]
@@ -286,10 +286,23 @@ async def add_mcp_server_tools(server: MCPServer) -> None:
         function_tool_defs = await ctx.deps.prepare_tools(run_context, function_tool_defs) or []
 
     output_schema = ctx.deps.output_schema
+
+    output_tools = []
+    output_object = None
+    if isinstance(output_schema, _output.ToolOutputSchema):
+        output_tools = output_schema.tool_defs()
+    elif isinstance(output_schema, _output.NativeOutputSchema):
+        output_object = output_schema.object_def
+
+    # ToolOrTextOutputSchema, NativeOutputSchema, and PromptedOutputSchema all inherit from TextOutputSchema
+    allow_text_output = isinstance(output_schema, _output.TextOutputSchema)
+
     return models.ModelRequestParameters(
         function_tools=function_tool_defs,
-        allow_text_output=_output.allow_text_output(output_schema),
-        output_tools=output_schema.tool_defs() if output_schema is not None else [],
+        output_mode=output_schema.mode,
+        output_tools=output_tools,
+        output_object=output_object,
+        allow_text_output=allow_text_output,
     )
 
 
@@ -484,7 +497,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 _output.allow_text_output(ctx.deps.output_schema):
+                    if isinstance(ctx.deps.output_schema, _output.TextOutputSchema):
                         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)]
@@ -507,10 +520,11 @@ async def _handle_tool_calls(
         output_schema = ctx.deps.output_schema
         run_context = build_run_context(ctx)
 
-        # first, look for the output tool call
         final_result: result.FinalResult[NodeRunEndT] | None = None
         parts: list[_messages.ModelRequestPart] = []
-        if output_schema is not None:
+
+        # first, look for the output tool call
+        if isinstance(output_schema, _output.ToolOutputSchema):
             for call, output_tool in output_schema.find_tool(tool_calls):
                 try:
                     result_data = await output_tool.process(call, run_context)
@@ -568,9 +582,9 @@ async def _handle_text_response(
 
         text = '\n\n'.join(texts)
         try:
-            if _output.allow_text_output(output_schema):
-                # The following cast is safe because we know `str` is an allowed result type
-                result_data = cast(NodeRunEndT, text)
+            if isinstance(output_schema, _output.TextOutputSchema):
+                run_context = build_run_context(ctx)
+                result_data = await output_schema.process(text, run_context)
             else:
                 m = _messages.RetryPromptPart(
                     content='Plain text responses are not permitted, please include your response in a tool call',
@@ -669,7 +683,7 @@ async def process_function_tools(  # noqa C901
                 yield event
                 call_index_to_event_id[len(calls_to_run)] = event.call_id
                 calls_to_run.append((mcp_tool, call))
-        elif output_schema is not None and call.tool_name in output_schema.tools:
+        elif call.tool_name in output_schema.tools:
             # if tool_name is in output_schema, it means we found a output tool but an error occurred in
             # validation, we don't add another part here
             if output_tool_name is not None:
@@ -809,7 +823,9 @@ def _unknown_tool(
 ) -> _messages.RetryPromptPart:
     ctx.state.increment_retries(ctx.deps.max_result_retries)
     tool_names = list(ctx.deps.function_tools.keys())
-    if output_schema := ctx.deps.output_schema:
+
+    output_schema = ctx.deps.output_schema
+    if isinstance(output_schema, _output.ToolOutputSchema):
         tool_names.extend(output_schema.tool_names())
 
     if tool_names:
@@ -886,7 +902,7 @@ def get_captured_run_messages() -> _RunMessages:
 def build_agent_graph(
     name: str | None,
     deps_type: type[DepsT],
-    output_type: _output.OutputType[OutputT],
+    output_type: OutputSpec[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/_agent_graph.py

@@ -14,14 +14,13 @@
 
 from typing_inspection.introspection import get_literal_values
 
-from pydantic_ai.result import OutputDataT
-from pydantic_ai.tools import AgentDepsT
-
 from . import __version__
+from ._run_context import AgentDepsT
 from .agent import Agent
 from .exceptions import UserError
 from .messages import ModelMessage
 from .models import KnownModelName, infer_model
+from .output import OutputDataT
 
 try:
     import argcomplete

pydantic_ai_slim/pydantic_ai/_cli.py

@@ -19,9 +19,8 @@
 from pydantic_core import SchemaValidator, core_schema
 from typing_extensions import Concatenate, ParamSpec, TypeIs, TypeVar, get_origin
 
-from pydantic_ai.tools import RunContext
-
 from ._griffe import doc_descriptions
+from ._run_context import RunContext
 from ._utils import check_object_json_schema, is_async_callable, is_model_like, run_in_executor
 
 if TYPE_CHECKING:
@@ -281,6 +280,4 @@ def _build_schema(
 
 def _is_call_ctx(annotation: Any) -> bool:
     """Return whether the annotation is the `RunContext` class, parameterized or not."""
-    from .tools import RunContext
-
     return annotation is RunContext or get_origin(annotation) is RunContext