pydantic-ai-slim 0.3.2__py3-none-any.whl → 0.3.3__py3-none-any.whl

pydantic_ai/_output.py

@@ -1,19 +1,35 @@
 from __future__ import annotations as _annotations
 
 import inspect
+import json
+from abc import ABC, abstractmethod
 from collections.abc import Awaitable, Iterable, Iterator, Sequence
 from dataclasses import dataclass, field
-from typing import Any, Callable, Generic, Literal, Union, cast
+from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Union, cast, overload
 
 from pydantic import TypeAdapter, ValidationError
 from pydantic_core import SchemaValidator
-from typing_extensions import TypeAliasType, TypedDict, TypeVar, get_args, get_origin
-from typing_inspection import typing_objects
-from typing_inspection.introspection import is_union_origin
+from typing_extensions import TypedDict, TypeVar, assert_never
 
 from . import _function_schema, _utils, messages as _messages
-from .exceptions import ModelRetry
-from .tools import AgentDepsT, GenerateToolJsonSchema, ObjectJsonSchema, RunContext, ToolDefinition
+from ._run_context import AgentDepsT, RunContext
+from .exceptions import ModelRetry, UserError
+from .output import (
+    NativeOutput,
+    OutputDataT,
+    OutputMode,
+    OutputSpec,
+    OutputTypeOrFunction,
+    PromptedOutput,
+    StructuredOutputMode,
+    TextOutput,
+    TextOutputFunc,
+    ToolOutput,
+)
+from .tools import GenerateToolJsonSchema, ObjectJsonSchema, ToolDefinition
+
+if TYPE_CHECKING:
+    from .profiles import ModelProfile
 
 T = TypeVar('T')
 """An invariant TypeVar."""
@@ -29,8 +45,6 @@ changing it would have negative consequences for the ergonomics of the library.
 At some point, it may make sense to change the input to OutputValidatorFunc to be `Any` or `object` as doing that would
 resolve these potential variance issues.
 """
-OutputDataT = TypeVar('OutputDataT', default=str, covariant=True)
-"""Covariant type variable for the result data type of a run."""
 
 OutputValidatorFunc = Union[
     Callable[[RunContext[AgentDepsT], OutputDataT_inv], OutputDataT_inv],
@@ -52,6 +66,14 @@ DEFAULT_OUTPUT_TOOL_NAME = 'final_result'
 DEFAULT_OUTPUT_TOOL_DESCRIPTION = 'The final response which ends this conversation'
 
 
+class ToolRetryError(Exception):
+    """Exception used to signal a `ToolRetry` message should be returned to the LLM."""
+
+    def __init__(self, tool_retry: _messages.RetryPromptPart):
+        self.tool_retry = tool_retry
+        super().__init__()
+
+
 @dataclass
 class OutputValidator(Generic[AgentDepsT, OutputDataT_inv]):
     function: OutputValidatorFunc[AgentDepsT, OutputDataT_inv]
@@ -101,140 +123,399 @@ class OutputValidator(Generic[AgentDepsT, OutputDataT_inv]):
             return result_data
 
 
-class ToolRetryError(Exception):
-    """Internal exception used to signal a `ToolRetry` message should be returned to the LLM."""
+class BaseOutputSchema(ABC, Generic[OutputDataT]):
+    @abstractmethod
+    def with_default_mode(self, mode: StructuredOutputMode) -> OutputSchema[OutputDataT]:
+        raise NotImplementedError()
 
-    def __init__(self, tool_retry: _messages.RetryPromptPart):
-        self.tool_retry = tool_retry
-        super().__init__()
+    @property
+    def tools(self) -> dict[str, OutputTool[OutputDataT]]:
+        """Get the tools for this output schema."""
+        return {}
 
 
 @dataclass(init=False)
-class ToolOutput(Generic[OutputDataT]):
-    """Marker class to use tools for outputs, and customize the tool."""
-
-    output_type: SimpleOutputType[OutputDataT]
-    name: str | None
-    description: str | None
-    max_retries: int | None
-    strict: bool | None
+class OutputSchema(BaseOutputSchema[OutputDataT], ABC):
+    """Model the final output from an agent run."""
 
-    def __init__(
-        self,
-        type_: SimpleOutputType[OutputDataT],
+    @classmethod
+    @overload
+    def build(
+        cls,
+        output_spec: OutputSpec[OutputDataT],
         *,
+        default_mode: StructuredOutputMode,
         name: str | None = None,
         description: str | None = None,
-        max_retries: int | None = None,
         strict: bool | None = None,
-    ):
-        self.output_type = type_
-        self.name = name
-        self.description = description
-        self.max_retries = max_retries
-        self.strict = strict
-
-
-T_co = TypeVar('T_co', covariant=True)
-# output_type=Type or output_type=function or output_type=object.method
-SimpleOutputType = TypeAliasType(
-    'SimpleOutputType', Union[type[T_co], Callable[..., Union[Awaitable[T_co], T_co]]], type_params=(T_co,)
-)
-# output_type=ToolOutput(<see above>) or <see above>
-SimpleOutputTypeOrMarker = TypeAliasType(
-    'SimpleOutputTypeOrMarker', Union[SimpleOutputType[T_co], ToolOutput[T_co]], type_params=(T_co,)
-)
-# output_type=<see above> or [<see above>, ...]
-OutputType = TypeAliasType(
-    'OutputType', Union[SimpleOutputTypeOrMarker[T_co], Sequence[SimpleOutputTypeOrMarker[T_co]]], type_params=(T_co,)
-)
-
-
-@dataclass
-class OutputSchema(Generic[OutputDataT]):
-    """Model the final output from an agent run.
+    ) -> OutputSchema[OutputDataT]: ...
 
-    Similar to `Tool` but for the final output of running an agent.
-    """
-
-    tools: dict[str, OutputTool[OutputDataT]]
-    allow_text_output: bool
+    @classmethod
+    @overload
+    def build(
+        cls,
+        output_spec: OutputSpec[OutputDataT],
+        *,
+        default_mode: None = None,
+        name: str | None = None,
+        description: str | None = None,
+        strict: bool | None = None,
+    ) -> BaseOutputSchema[OutputDataT]: ...
 
     @classmethod
     def build(
-        cls: type[OutputSchema[OutputDataT]],
-        output_type: OutputType[OutputDataT],
+        cls,
+        output_spec: OutputSpec[OutputDataT],
+        *,
+        default_mode: StructuredOutputMode | None = None,
         name: str | None = None,
         description: str | None = None,
         strict: bool | None = None,
-    ) -> OutputSchema[OutputDataT] | None:
+    ) -> BaseOutputSchema[OutputDataT]:
         """Build an OutputSchema dataclass from an output type."""
-        if output_type is str:
-            return None
+        if output_spec is str:
+            return PlainTextOutputSchema()
+
+        if isinstance(output_spec, NativeOutput):
+            return NativeOutputSchema(
+                cls._build_processor(
+                    _flatten_output_spec(output_spec.outputs),
+                    name=output_spec.name,
+                    description=output_spec.description,
+                )
+            )
+        elif isinstance(output_spec, PromptedOutput):
+            return PromptedOutputSchema(
+                cls._build_processor(
+                    _flatten_output_spec(output_spec.outputs),
+                    name=output_spec.name,
+                    description=output_spec.description,
+                ),
+                template=output_spec.template,
+            )
 
-        output_types: Sequence[SimpleOutputTypeOrMarker[OutputDataT]]
-        if isinstance(output_type, Sequence):
-            output_types = output_type
-        else:
-            output_types = (output_type,)
+        text_outputs: Sequence[type[str] | TextOutput[OutputDataT]] = []
+        tool_outputs: Sequence[ToolOutput[OutputDataT]] = []
+        other_outputs: Sequence[OutputTypeOrFunction[OutputDataT]] = []
+        for output in _flatten_output_spec(output_spec):
+            if output is str:
+                text_outputs.append(cast(type[str], output))
+            elif isinstance(output, TextOutput):
+                text_outputs.append(output)
+            elif isinstance(output, ToolOutput):
+                tool_outputs.append(output)
+            else:
+                other_outputs.append(output)
 
-        output_types_flat: list[SimpleOutputTypeOrMarker[OutputDataT]] = []
-        for output_type in output_types:
-            if union_types := get_union_args(output_type):
-                output_types_flat.extend(union_types)
+        tools = cls._build_tools(tool_outputs + other_outputs, name=name, description=description, strict=strict)
+
+        if len(text_outputs) > 0:
+            if len(text_outputs) > 1:
+                raise UserError('Only one text output is allowed.')
+            text_output = text_outputs[0]
+
+            text_output_schema = None
+            if isinstance(text_output, TextOutput):
+                text_output_schema = PlainTextOutputProcessor(text_output.output_function)
+
+            if len(tools) == 0:
+                return PlainTextOutputSchema(text_output_schema)
             else:
-                output_types_flat.append(output_type)
+                return ToolOrTextOutputSchema(processor=text_output_schema, tools=tools)
 
-        allow_text_output = False
-        if str in output_types_flat:
-            allow_text_output = True
-            output_types_flat = [t for t in output_types_flat if t is not str]
+        if len(tool_outputs) > 0:
+            return ToolOutputSchema(tools)
 
-        multiple = len(output_types_flat) > 1
+        if len(other_outputs) > 0:
+            schema = OutputSchemaWithoutMode(
+                processor=cls._build_processor(other_outputs, name=name, description=description, strict=strict),
+                tools=tools,
+            )
+            if default_mode:
+                schema = schema.with_default_mode(default_mode)
+            return schema
 
-        default_tool_name = name or DEFAULT_OUTPUT_TOOL_NAME
-        default_tool_description = description
-        default_tool_strict = strict
+        raise UserError('No output type provided.')  # pragma: no cover
 
+    @staticmethod
+    def _build_tools(
+        outputs: list[OutputTypeOrFunction[OutputDataT] | ToolOutput[OutputDataT]],
+        name: str | None = None,
+        description: str | None = None,
+        strict: bool | None = None,
+    ) -> dict[str, OutputTool[OutputDataT]]:
         tools: dict[str, OutputTool[OutputDataT]] = {}
-        for output_type in output_types_flat:
-            tool_name = None
-            tool_description = None
-            tool_strict = None
-            if isinstance(output_type, ToolOutput):
-                tool_output_type = output_type.output_type
+
+        default_name = name or DEFAULT_OUTPUT_TOOL_NAME
+        default_description = description
+        default_strict = strict
+
+        multiple = len(outputs) > 1
+        for output in outputs:
+            name = None
+            description = None
+            strict = None
+            if isinstance(output, ToolOutput):
                 # do we need to error on conflicts here? (DavidM): If this is internal maybe doesn't matter, if public, use overloads
-                tool_name = output_type.name
-                tool_description = output_type.description
-                tool_strict = output_type.strict
-            else:
-                tool_output_type = output_type
+                name = output.name
+                description = output.description
+                strict = output.strict
 
-            if tool_name is None:
-                tool_name = default_tool_name
+                output = output.output
+
+            if name is None:
+                name = default_name
                 if multiple:
-                    tool_name += f'_{tool_output_type.__name__}'
+                    name += f'_{output.__name__}'
 
             i = 1
-            original_tool_name = tool_name
-            while tool_name in tools:
+            original_name = name
+            while name in tools:
                 i += 1
-                tool_name = f'{original_tool_name}_{i}'
+                name = f'{original_name}_{i}'
 
-            tool_description = tool_description or default_tool_description
-            if tool_strict is None:
-                tool_strict = default_tool_strict
+            description = description or default_description
+            if strict is None:
+                strict = default_strict
 
-            parameters_schema = OutputObjectSchema(
-                output_type=tool_output_type, description=tool_description, strict=tool_strict
-            )
-            tools[tool_name] = OutputTool(name=tool_name, parameters_schema=parameters_schema, multiple=multiple)
+            processor = ObjectOutputProcessor(output=output, description=description, strict=strict)
+            tools[name] = OutputTool(name=name, processor=processor, multiple=multiple)
+
+        return tools
+
+    @staticmethod
+    def _build_processor(
+        outputs: Sequence[OutputTypeOrFunction[OutputDataT]],
+        name: str | None = None,
+        description: str | None = None,
+        strict: bool | None = None,
+    ) -> ObjectOutputProcessor[OutputDataT] | UnionOutputProcessor[OutputDataT]:
+        outputs = _flatten_output_spec(outputs)
+        if len(outputs) == 1:
+            return ObjectOutputProcessor(output=outputs[0], name=name, description=description, strict=strict)
+
+        return UnionOutputProcessor(outputs=outputs, strict=strict, name=name, description=description)
+
+    @property
+    @abstractmethod
+    def mode(self) -> OutputMode:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def raise_if_unsupported(self, profile: ModelProfile) -> None:
+        """Raise an error if the mode is not supported by the model."""
+        raise NotImplementedError()
+
+    def with_default_mode(self, mode: StructuredOutputMode) -> OutputSchema[OutputDataT]:
+        return self
+
+
+@dataclass(init=False)
+class OutputSchemaWithoutMode(BaseOutputSchema[OutputDataT]):
+    processor: ObjectOutputProcessor[OutputDataT] | UnionOutputProcessor[OutputDataT]
+    _tools: dict[str, OutputTool[OutputDataT]] = field(default_factory=dict)
+
+    def __init__(
+        self,
+        processor: ObjectOutputProcessor[OutputDataT] | UnionOutputProcessor[OutputDataT],
+        tools: dict[str, OutputTool[OutputDataT]],
+    ):
+        self.processor = processor
+        self._tools = tools
+
+    def with_default_mode(self, mode: StructuredOutputMode) -> OutputSchema[OutputDataT]:
+        if mode == 'native':
+            return NativeOutputSchema(self.processor)
+        elif mode == 'prompted':
+            return PromptedOutputSchema(self.processor)
+        elif mode == 'tool':
+            return ToolOutputSchema(self.tools)
+        else:
+            assert_never(mode)
 
-        return cls(
-            tools=tools,
-            allow_text_output=allow_text_output,
+    @property
+    def tools(self) -> dict[str, OutputTool[OutputDataT]]:
+        """Get the tools for this output schema."""
+        # We return tools here as they're checked in Agent._register_tool.
+        # At that point we may don't know yet what output mode we're going to use if no model was provided or it was deferred until agent.run time.
+        return self._tools
+
+
+class TextOutputSchema(OutputSchema[OutputDataT], ABC):
+    @abstractmethod
+    async def process(
+        self,
+        text: str,
+        run_context: RunContext[AgentDepsT],
+        allow_partial: bool = False,
+        wrap_validation_errors: bool = True,
+    ) -> OutputDataT:
+        raise NotImplementedError()
+
+
+@dataclass
+class PlainTextOutputSchema(TextOutputSchema[OutputDataT]):
+    processor: PlainTextOutputProcessor[OutputDataT] | None = None
+
+    @property
+    def mode(self) -> OutputMode:
+        return 'text'
+
+    def raise_if_unsupported(self, profile: ModelProfile) -> None:
+        """Raise an error if the mode is not supported by the model."""
+        pass
+
+    async def process(
+        self,
+        text: str,
+        run_context: RunContext[AgentDepsT],
+        allow_partial: bool = False,
+        wrap_validation_errors: bool = True,
+    ) -> OutputDataT:
+        """Validate an output message.
+
+        Args:
+            text: The output text to validate.
+            run_context: The current run context.
+            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).
+        """
+        if self.processor is None:
+            return cast(OutputDataT, text)
+
+        return await self.processor.process(
+            text, run_context, allow_partial=allow_partial, wrap_validation_errors=wrap_validation_errors
         )
 
+
+@dataclass
+class StructuredTextOutputSchema(TextOutputSchema[OutputDataT], ABC):
+    processor: ObjectOutputProcessor[OutputDataT] | UnionOutputProcessor[OutputDataT]
+
+    @property
+    def object_def(self) -> OutputObjectDefinition:
+        return self.processor.object_def
+
+
+@dataclass
+class NativeOutputSchema(StructuredTextOutputSchema[OutputDataT]):
+    @property
+    def mode(self) -> OutputMode:
+        return 'native'
+
+    def raise_if_unsupported(self, profile: ModelProfile) -> None:
+        """Raise an error if the mode is not supported by the model."""
+        if not profile.supports_json_schema_output:
+            raise UserError('Structured output is not supported by the model.')
+
+    async def process(
+        self,
+        text: str,
+        run_context: RunContext[AgentDepsT],
+        allow_partial: bool = False,
+        wrap_validation_errors: bool = True,
+    ) -> OutputDataT:
+        """Validate an output message.
+
+        Args:
+            text: The output text to validate.
+            run_context: The current run context.
+            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).
+        """
+        return await self.processor.process(
+            text, run_context, allow_partial=allow_partial, wrap_validation_errors=wrap_validation_errors
+        )
+
+
+@dataclass
+class PromptedOutputSchema(StructuredTextOutputSchema[OutputDataT]):
+    template: str | None = None
+
+    @property
+    def mode(self) -> OutputMode:
+        return 'prompted'
+
+    def raise_if_unsupported(self, profile: ModelProfile) -> None:
+        """Raise an error if the mode is not supported by the model."""
+        pass
+
+    def instructions(self, default_template: str) -> str:
+        """Get instructions to tell model to output JSON matching the schema."""
+        template = self.template or default_template
+
+        if '{schema}' not in template:
+            template = '\n\n'.join([template, '{schema}'])
+
+        object_def = self.object_def
+        schema = object_def.json_schema.copy()
+        if object_def.name:
+            schema['title'] = object_def.name
+        if object_def.description:
+            schema['description'] = object_def.description
+
+        return template.format(schema=json.dumps(schema))
+
+    async def process(
+        self,
+        text: str,
+        run_context: RunContext[AgentDepsT],
+        allow_partial: bool = False,
+        wrap_validation_errors: bool = True,
+    ) -> OutputDataT:
+        """Validate an output message.
+
+        Args:
+            text: The output text to validate.
+            run_context: The current run context.
+            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).
+        """
+        text = _utils.strip_markdown_fences(text)
+
+        return await self.processor.process(
+            text, run_context, allow_partial=allow_partial, wrap_validation_errors=wrap_validation_errors
+        )
+
+
+@dataclass(init=False)
+class ToolOutputSchema(OutputSchema[OutputDataT]):
+    _tools: dict[str, OutputTool[OutputDataT]] = field(default_factory=dict)
+
+    def __init__(self, tools: dict[str, OutputTool[OutputDataT]]):
+        self._tools = tools
+
+    @property
+    def mode(self) -> OutputMode:
+        return 'tool'
+
+    def raise_if_unsupported(self, profile: ModelProfile) -> None:
+        """Raise an error if the mode is not supported by the model."""
+        if not profile.supports_tools:
+            raise UserError('Output tools are not supported by the model.')
+
+    @property
+    def tools(self) -> dict[str, OutputTool[OutputDataT]]:
+        """Get the tools for this output schema."""
+        return self._tools
+
+    def tool_names(self) -> list[str]:
+        """Return the names of the tools."""
+        return list(self.tools.keys())
+
+    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 find_named_tool(
         self, parts: Iterable[_messages.ModelResponsePart], tool_name: str
     ) -> tuple[_messages.ToolCallPart, OutputTool[OutputDataT]] | None:
@@ -254,61 +535,78 @@ class OutputSchema(Generic[OutputDataT]):
                 if result := self.tools.get(part.tool_name):
                     yield part, result
 
-    def tool_names(self) -> list[str]:
-        """Return the names of the tools."""
-        return list(self.tools.keys())
-
-    def tool_defs(self) -> list[ToolDefinition]:
-        """Get tool definitions to register with the model."""
-        return [t.tool_def for t in self.tools.values()]
 
+@dataclass(init=False)
+class ToolOrTextOutputSchema(ToolOutputSchema[OutputDataT], PlainTextOutputSchema[OutputDataT]):
+    def __init__(
+        self,
+        processor: PlainTextOutputProcessor[OutputDataT] | None,
+        tools: dict[str, OutputTool[OutputDataT]],
+    ):
+        self.processor = processor
+        self._tools = tools
 
-def allow_text_output(output_schema: OutputSchema[Any] | None) -> bool:
-    return output_schema is None or output_schema.allow_text_output
+    @property
+    def mode(self) -> OutputMode:
+        return 'tool_or_text'
 
 
 @dataclass
 class OutputObjectDefinition:
-    name: str
     json_schema: ObjectJsonSchema
+    name: str | None = None
     description: str | None = None
     strict: bool | None = None
 
 
 @dataclass(init=False)
-class OutputObjectSchema(Generic[OutputDataT]):
-    definition: OutputObjectDefinition
-    validator: SchemaValidator
-    function_schema: _function_schema.FunctionSchema | None = None
+class BaseOutputProcessor(ABC, Generic[OutputDataT]):
+    @abstractmethod
+    async def process(
+        self,
+        data: str,
+        run_context: RunContext[AgentDepsT],
+        allow_partial: bool = False,
+        wrap_validation_errors: bool = True,
+    ) -> OutputDataT:
+        """Process an output message, performing validation and (if necessary) calling the output function."""
+        raise NotImplementedError()
+
+
+@dataclass(init=False)
+class ObjectOutputProcessor(BaseOutputProcessor[OutputDataT]):
+    object_def: OutputObjectDefinition
     outer_typed_dict_key: str | None = None
+    _validator: SchemaValidator
+    _function_schema: _function_schema.FunctionSchema | None = None
 
     def __init__(
         self,
+        output: OutputTypeOrFunction[OutputDataT],
         *,
-        output_type: SimpleOutputType[OutputDataT],
         name: str | None = None,
         description: str | None = None,
         strict: bool | None = None,
     ):
-        if inspect.isfunction(output_type) or inspect.ismethod(output_type):
-            self.function_schema = _function_schema.function_schema(output_type, GenerateToolJsonSchema)
-            self.validator = self.function_schema.validator
-            json_schema = self.function_schema.json_schema
-            json_schema['description'] = self.function_schema.description
+        if inspect.isfunction(output) or inspect.ismethod(output):
+            self._function_schema = _function_schema.function_schema(output, GenerateToolJsonSchema)
+            self._validator = self._function_schema.validator
+            json_schema = self._function_schema.json_schema
+            json_schema['description'] = self._function_schema.description
         else:
             type_adapter: TypeAdapter[Any]
-            if _utils.is_model_like(output_type):
-                type_adapter = TypeAdapter(output_type)
+            if _utils.is_model_like(output):
+                type_adapter = TypeAdapter(output)
             else:
                 self.outer_typed_dict_key = 'response'
                 response_data_typed_dict = TypedDict(  # noqa: UP013
                     'response_data_typed_dict',
-                    {'response': cast(type[OutputDataT], output_type)},  # pyright: ignore[reportInvalidTypeForm]
+                    {'response': cast(type[OutputDataT], output)},  # pyright: ignore[reportInvalidTypeForm]
                 )
                 type_adapter = TypeAdapter(response_data_typed_dict)
 
             # Really a PluggableSchemaValidator, but it's API-compatible
-            self.validator = cast(SchemaValidator, type_adapter.validator)
+            self._validator = cast(SchemaValidator, type_adapter.validator)
             json_schema = _utils.check_object_json_schema(
                 type_adapter.json_schema(schema_generator=GenerateToolJsonSchema)
             )
@@ -323,8 +621,8 @@ class OutputObjectSchema(Generic[OutputDataT]):
             else:
                 description = f'{description}. {json_schema_description}'
 
-        self.definition = OutputObjectDefinition(
-            name=name or getattr(output_type, '__name__', DEFAULT_OUTPUT_TOOL_NAME),
+        self.object_def = OutputObjectDefinition(
+            name=name or getattr(output, '__name__', None),
             description=description,
             json_schema=json_schema,
             strict=strict,
@@ -335,6 +633,7 @@ class OutputObjectSchema(Generic[OutputDataT]):
         data: str | dict[str, Any] | None,
         run_context: RunContext[AgentDepsT],
         allow_partial: bool = False,
+        wrap_validation_errors: bool = True,
     ) -> OutputDataT:
         """Process an output message, performing validation and (if necessary) calling the output function.
 
@@ -342,45 +641,235 @@ class OutputObjectSchema(Generic[OutputDataT]):
             data: The output data to validate.
             run_context: The current run context.
             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).
         """
-        pyd_allow_partial: Literal['off', 'trailing-strings'] = 'trailing-strings' if allow_partial else 'off'
-        if isinstance(data, str):
-            output = self.validator.validate_json(data or '{}', allow_partial=pyd_allow_partial)
-        else:
-            output = self.validator.validate_python(data or {}, allow_partial=pyd_allow_partial)
-
-        if self.function_schema:
-            output = await self.function_schema.call(output, run_context)
+        try:
+            pyd_allow_partial: Literal['off', 'trailing-strings'] = 'trailing-strings' if allow_partial else 'off'
+            if isinstance(data, str):
+                output = self._validator.validate_json(data or '{}', allow_partial=pyd_allow_partial)
+            else:
+                output = self._validator.validate_python(data or {}, 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  # pragma: lax no cover
 
         if k := self.outer_typed_dict_key:
             output = output[k]
+
+        if self._function_schema:
+            try:
+                output = await self._function_schema.call(output, run_context)
+            except ModelRetry as r:
+                if wrap_validation_errors:
+                    m = _messages.RetryPromptPart(
+                        content=r.message,
+                    )
+                    raise ToolRetryError(m) from r
+                else:
+                    raise  # pragma: lax no cover
+
         return output
 
 
+@dataclass
+class UnionOutputResult:
+    kind: str
+    data: ObjectJsonSchema
+
+
+@dataclass
+class UnionOutputModel:
+    result: UnionOutputResult
+
+
+@dataclass(init=False)
+class UnionOutputProcessor(BaseOutputProcessor[OutputDataT]):
+    object_def: OutputObjectDefinition
+    _union_processor: ObjectOutputProcessor[UnionOutputModel]
+    _processors: dict[str, ObjectOutputProcessor[OutputDataT]]
+
+    def __init__(
+        self,
+        outputs: Sequence[OutputTypeOrFunction[OutputDataT]],
+        *,
+        name: str | None = None,
+        description: str | None = None,
+        strict: bool | None = None,
+    ):
+        self._union_processor = ObjectOutputProcessor(output=UnionOutputModel)
+
+        json_schemas: list[ObjectJsonSchema] = []
+        self._processors = {}
+        for output in outputs:
+            processor = ObjectOutputProcessor(output=output, strict=strict)
+            object_def = processor.object_def
+
+            object_key = object_def.name or output.__name__
+            i = 1
+            original_key = object_key
+            while object_key in self._processors:
+                i += 1
+                object_key = f'{original_key}_{i}'
+
+            self._processors[object_key] = processor
+
+            json_schema = object_def.json_schema
+            if object_def.name:  # pragma: no branch
+                json_schema['title'] = object_def.name
+            if object_def.description:
+                json_schema['description'] = object_def.description
+
+            json_schemas.append(json_schema)
+
+        json_schemas, all_defs = _utils.merge_json_schema_defs(json_schemas)
+
+        discriminated_json_schemas: list[ObjectJsonSchema] = []
+        for object_key, json_schema in zip(self._processors.keys(), json_schemas):
+            title = json_schema.pop('title', None)
+            description = json_schema.pop('description', None)
+
+            discriminated_json_schema = {
+                'type': 'object',
+                'properties': {
+                    'kind': {
+                        'type': 'string',
+                        'const': object_key,
+                    },
+                    'data': json_schema,
+                },
+                'required': ['kind', 'data'],
+                'additionalProperties': False,
+            }
+            if title:  # pragma: no branch
+                discriminated_json_schema['title'] = title
+            if description:
+                discriminated_json_schema['description'] = description
+
+            discriminated_json_schemas.append(discriminated_json_schema)
+
+        json_schema = {
+            'type': 'object',
+            'properties': {
+                'result': {
+                    'anyOf': discriminated_json_schemas,
+                }
+            },
+            'required': ['result'],
+            'additionalProperties': False,
+        }
+        if all_defs:
+            json_schema['$defs'] = all_defs
+
+        self.object_def = OutputObjectDefinition(
+            json_schema=json_schema,
+            strict=strict,
+            name=name,
+            description=description,
+        )
+
+    async def process(
+        self,
+        data: str | dict[str, Any] | None,
+        run_context: RunContext[AgentDepsT],
+        allow_partial: bool = False,
+        wrap_validation_errors: bool = True,
+    ) -> OutputDataT:
+        union_object = await self._union_processor.process(
+            data, run_context, allow_partial=allow_partial, wrap_validation_errors=wrap_validation_errors
+        )
+
+        result = union_object.result
+        kind = result.kind
+        data = result.data
+        try:
+            processor = self._processors[kind]
+        except KeyError as e:  # pragma: no cover
+            if wrap_validation_errors:
+                m = _messages.RetryPromptPart(content=f'Invalid kind: {kind}')
+                raise ToolRetryError(m) from e
+            else:
+                raise
+
+        return await processor.process(
+            data, run_context, allow_partial=allow_partial, wrap_validation_errors=wrap_validation_errors
+        )
+
+
+@dataclass(init=False)
+class PlainTextOutputProcessor(BaseOutputProcessor[OutputDataT]):
+    _function_schema: _function_schema.FunctionSchema
+    _str_argument_name: str
+
+    def __init__(
+        self,
+        output_function: TextOutputFunc[OutputDataT],
+    ):
+        self._function_schema = _function_schema.function_schema(output_function, GenerateToolJsonSchema)
+
+        arguments_schema = self._function_schema.json_schema.get('properties', {})
+        argument_name = next(iter(arguments_schema.keys()), None)
+        if argument_name and arguments_schema.get(argument_name, {}).get('type') == 'string':
+            self._str_argument_name = argument_name
+            return
+
+        raise UserError('TextOutput must take a function taking a `str`')
+
+    @property
+    def object_def(self) -> None:
+        return None  # pragma: no cover
+
+    async def process(
+        self,
+        data: str,
+        run_context: RunContext[AgentDepsT],
+        allow_partial: bool = False,
+        wrap_validation_errors: bool = True,
+    ) -> OutputDataT:
+        args = {self._str_argument_name: data}
+
+        try:
+            output = await self._function_schema.call(args, run_context)
+        except ModelRetry as r:
+            if wrap_validation_errors:
+                m = _messages.RetryPromptPart(
+                    content=r.message,
+                )
+                raise ToolRetryError(m) from r
+            else:
+                raise  # pragma: lax no cover
+
+        return cast(OutputDataT, output)
+
+
 @dataclass(init=False)
 class OutputTool(Generic[OutputDataT]):
-    parameters_schema: OutputObjectSchema[OutputDataT]
+    processor: ObjectOutputProcessor[OutputDataT]
     tool_def: ToolDefinition
 
-    def __init__(self, *, name: str, parameters_schema: OutputObjectSchema[OutputDataT], multiple: bool):
-        self.parameters_schema = parameters_schema
-        definition = parameters_schema.definition
+    def __init__(self, *, name: str, processor: ObjectOutputProcessor[OutputDataT], multiple: bool):
+        self.processor = processor
+        object_def = processor.object_def
 
-        description = definition.description
+        description = object_def.description
         if not description:
             description = DEFAULT_OUTPUT_TOOL_DESCRIPTION
             if multiple:
-                description = f'{definition.name}: {description}'
+                description = f'{object_def.name}: {description}'
 
         self.tool_def = ToolDefinition(
             name=name,
             description=description,
-            parameters_json_schema=definition.json_schema,
-            strict=definition.strict,
-            outer_typed_dict_key=parameters_schema.outer_typed_dict_key,
+            parameters_json_schema=object_def.json_schema,
+            strict=object_def.strict,
+            outer_typed_dict_key=processor.outer_typed_dict_key,
         )
 
     async def process(
@@ -402,7 +891,9 @@ class OutputTool(Generic[OutputDataT]):
             Either the validated output data (left) or a retry message (right).
         """
         try:
-            output = await self.parameters_schema.process(tool_call.args, run_context, allow_partial=allow_partial)
+            output = await self.processor.process(
+                tool_call.args, run_context, allow_partial=allow_partial, wrap_validation_errors=False
+            )
         except ValidationError as e:
             if wrap_validation_errors:
                 m = _messages.RetryPromptPart(
@@ -427,13 +918,17 @@ class OutputTool(Generic[OutputDataT]):
             return output
 
 
-def get_union_args(tp: Any) -> tuple[Any, ...]:
-    """Extract the arguments of a Union type if `output_type` is a union, otherwise return an empty tuple."""
-    if typing_objects.is_typealiastype(tp):
-        tp = tp.__value__
-
-    origin = get_origin(tp)
-    if is_union_origin(origin):
-        return get_args(tp)
+def _flatten_output_spec(output_spec: T | Sequence[T]) -> list[T]:
+    outputs: Sequence[T]
+    if isinstance(output_spec, Sequence):
+        outputs = output_spec
     else:
-        return ()
+        outputs = (output_spec,)
+
+    outputs_flat: list[T] = []
+    for output in outputs:
+        if union_types := _utils.get_union_args(output):
+            outputs_flat.extend(union_types)
+        else:
+            outputs_flat.append(output)
+    return outputs_flat