pydantic-ai-slim 1.0.13__py3-none-any.whl → 1.0.15__py3-none-any.whl

pydantic_ai/_output.py

@@ -75,7 +75,7 @@ DEFAULT_OUTPUT_TOOL_DESCRIPTION = 'The final response which ends this conversati
 async def execute_traced_output_function(
     function_schema: _function_schema.FunctionSchema,
     run_context: RunContext[AgentDepsT],
-    args: dict[str, Any] | Any,
+    args: dict[str, Any],
     wrap_validation_errors: bool = True,
 ) -> Any:
     """Execute an output function within a traced span with error handling.
@@ -209,18 +209,20 @@ class OutputValidator(Generic[AgentDepsT, OutputDataT_inv]):
             return result_data
 
 
-@dataclass
+@dataclass(kw_only=True)
 class BaseOutputSchema(ABC, Generic[OutputDataT]):
-    allows_deferred_tools: bool
+    text_processor: BaseOutputProcessor[OutputDataT] | None = None
+    toolset: OutputToolset[Any] | None = None
+    allows_deferred_tools: bool = False
+    allows_image: bool = False
 
     @abstractmethod
     def with_default_mode(self, mode: StructuredOutputMode) -> OutputSchema[OutputDataT]:
         raise NotImplementedError()
 
     @property
-    def toolset(self) -> OutputToolset[Any] | None:
-        """Get the toolset for this output schema."""
-        return None
+    def allows_text(self) -> bool:
+        return self.text_processor is not None
 
 
 @dataclass(init=False)
@@ -262,38 +264,67 @@ class OutputSchema(BaseOutputSchema[OutputDataT], ABC):
         strict: bool | None = None,
     ) -> BaseOutputSchema[OutputDataT]:
         """Build an OutputSchema dataclass from an output type."""
-        raw_outputs = _flatten_output_spec(output_spec)
+        outputs = _flatten_output_spec(output_spec)
+
+        allows_deferred_tools = DeferredToolRequests in outputs
+        if allows_deferred_tools:
+            outputs = [output for output in outputs if output is not DeferredToolRequests]
+            if len(outputs) == 0:
+                raise UserError('At least one output type must be provided other than `DeferredToolRequests`.')
 
-        outputs = [output for output in raw_outputs if output is not DeferredToolRequests]
-        allows_deferred_tools = len(outputs) < len(raw_outputs)
-        if len(outputs) == 0 and allows_deferred_tools:
-            raise UserError('At least one output type must be provided other than `DeferredToolRequests`.')
+        allows_image = _messages.BinaryImage in outputs
+        if allows_image:
+            outputs = [output for output in outputs if output is not _messages.BinaryImage]
 
         if output := next((output for output in outputs if isinstance(output, NativeOutput)), None):
             if len(outputs) > 1:
                 raise UserError('`NativeOutput` must be the only output type.')  # pragma: no cover
 
+            flattened_outputs = _flatten_output_spec(output.outputs)
+
+            if DeferredToolRequests in flattened_outputs:
+                raise UserError(  # pragma: no cover
+                    '`NativeOutput` cannot contain `DeferredToolRequests`. Include it alongside the native output marker instead: `output_type=[NativeOutput(...), DeferredToolRequests]`'
+                )
+            if _messages.BinaryImage in flattened_outputs:
+                raise UserError(  # pragma: no cover
+                    '`NativeOutput` cannot contain `BinaryImage`. Include it alongside the native output marker instead: `output_type=[NativeOutput(...), BinaryImage]`'
+                )
+
             return NativeOutputSchema(
                 processor=cls._build_processor(
-                    _flatten_output_spec(output.outputs),
+                    flattened_outputs,
                     name=output.name,
                     description=output.description,
                     strict=output.strict,
                 ),
                 allows_deferred_tools=allows_deferred_tools,
+                allows_image=allows_image,
             )
         elif output := next((output for output in outputs if isinstance(output, PromptedOutput)), None):
             if len(outputs) > 1:
                 raise UserError('`PromptedOutput` must be the only output type.')  # pragma: no cover
 
+            flattened_outputs = _flatten_output_spec(output.outputs)
+
+            if DeferredToolRequests in flattened_outputs:
+                raise UserError(  # pragma: no cover
+                    '`PromptedOutput` cannot contain `DeferredToolRequests`. Include it alongside the prompted output marker instead: `output_type=[PromptedOutput(...), DeferredToolRequests]`'
+                )
+            if _messages.BinaryImage in flattened_outputs:
+                raise UserError(  # pragma: no cover
+                    '`PromptedOutput` cannot contain `BinaryImage`. Include it alongside the prompted output marker instead: `output_type=[PromptedOutput(...), BinaryImage]`'
+                )
+
             return PromptedOutputSchema(
+                template=output.template,
                 processor=cls._build_processor(
-                    _flatten_output_spec(output.outputs),
+                    flattened_outputs,
                     name=output.name,
                     description=output.description,
                 ),
-                template=output.template,
                 allows_deferred_tools=allows_deferred_tools,
+                allows_image=allows_image,
             )
 
         text_outputs: Sequence[type[str] | TextOutput[OutputDataT]] = []
@@ -317,37 +348,51 @@ class OutputSchema(BaseOutputSchema[OutputDataT], ABC):
 
         toolset = OutputToolset.build(tool_outputs + other_outputs, name=name, description=description, strict=strict)
 
+        text_processor: BaseOutputProcessor[OutputDataT] | None = None
+
         if len(text_outputs) > 0:
             if len(text_outputs) > 1:
                 raise UserError('Only one `str` or `TextOutput` is allowed.')
             text_output = text_outputs[0]
 
-            text_output_schema = None
             if isinstance(text_output, TextOutput):
-                text_output_schema = PlainTextOutputProcessor(text_output.output_function)
+                text_processor = TextFunctionOutputProcessor(text_output.output_function)
+            else:
+                text_processor = TextOutputProcessor()
 
             if toolset:
-                return ToolOrTextOutputSchema(
-                    processor=text_output_schema,
+                return ToolOutputSchema(
                     toolset=toolset,
+                    text_processor=text_processor,
                     allows_deferred_tools=allows_deferred_tools,
+                    allows_image=allows_image,
                 )
             else:
-                return PlainTextOutputSchema(processor=text_output_schema, allows_deferred_tools=allows_deferred_tools)
+                return TextOutputSchema(
+                    text_processor=text_processor,
+                    allows_deferred_tools=allows_deferred_tools,
+                    allows_image=allows_image,
+                )
 
         if len(tool_outputs) > 0:
-            return ToolOutputSchema(toolset=toolset, allows_deferred_tools=allows_deferred_tools)
+            return ToolOutputSchema(
+                toolset=toolset, allows_deferred_tools=allows_deferred_tools, allows_image=allows_image
+            )
 
         if len(other_outputs) > 0:
             schema = OutputSchemaWithoutMode(
                 processor=cls._build_processor(other_outputs, name=name, description=description, strict=strict),
                 toolset=toolset,
                 allows_deferred_tools=allows_deferred_tools,
+                allows_image=allows_image,
             )
             if default_mode:
                 schema = schema.with_default_mode(default_mode)
             return schema
 
+        if allows_image:
+            return ImageOutputSchema(allows_deferred_tools=allows_deferred_tools)
+
         raise UserError('At least one output type must be provided.')
 
     @staticmethod
@@ -356,7 +401,7 @@ class OutputSchema(BaseOutputSchema[OutputDataT], ABC):
         name: str | None = None,
         description: str | None = None,
         strict: bool | None = None,
-    ) -> ObjectOutputProcessor[OutputDataT] | UnionOutputProcessor[OutputDataT]:
+    ) -> BaseObjectOutputProcessor[OutputDataT]:
         outputs = _flatten_output_spec(outputs)
         if len(outputs) == 1:
             return ObjectOutputProcessor(output=outputs[0], name=name, description=description, strict=strict)
@@ -368,10 +413,10 @@ class OutputSchema(BaseOutputSchema[OutputDataT], ABC):
     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()
+        """Raise an error if the mode is not supported by this model."""
+        if self.allows_image and not profile.supports_image_output:
+            raise UserError('Image output is not supported by this model.')
 
     def with_default_mode(self, mode: StructuredOutputMode) -> OutputSchema[OutputDataT]:
         return self
@@ -379,142 +424,139 @@ class OutputSchema(BaseOutputSchema[OutputDataT], ABC):
 
 @dataclass(init=False)
 class OutputSchemaWithoutMode(BaseOutputSchema[OutputDataT]):
-    processor: ObjectOutputProcessor[OutputDataT] | UnionOutputProcessor[OutputDataT]
-    _toolset: OutputToolset[Any] | None
+    processor: BaseObjectOutputProcessor[OutputDataT]
 
     def __init__(
         self,
-        processor: ObjectOutputProcessor[OutputDataT] | UnionOutputProcessor[OutputDataT],
+        processor: BaseObjectOutputProcessor[OutputDataT],
         toolset: OutputToolset[Any] | None,
         allows_deferred_tools: bool,
+        allows_image: bool,
     ):
-        super().__init__(allows_deferred_tools)
+        # We set a toolset here as they're checked for name conflicts with other toolsets in the Agent constructor.
+        # At that point we may not know yet what output mode we're going to use if no model was provided or it was deferred until agent.run time,
+        # but we cover ourselves just in case we end up using the tool output mode.
+        super().__init__(
+            allows_deferred_tools=allows_deferred_tools,
+            toolset=toolset,
+            text_processor=processor,
+            allows_image=allows_image,
+        )
         self.processor = processor
-        self._toolset = toolset
 
     def with_default_mode(self, mode: StructuredOutputMode) -> OutputSchema[OutputDataT]:
         if mode == 'native':
-            return NativeOutputSchema(processor=self.processor, allows_deferred_tools=self.allows_deferred_tools)
+            return NativeOutputSchema(
+                processor=self.processor,
+                allows_deferred_tools=self.allows_deferred_tools,
+                allows_image=self.allows_image,
+            )
         elif mode == 'prompted':
-            return PromptedOutputSchema(processor=self.processor, allows_deferred_tools=self.allows_deferred_tools)
+            return PromptedOutputSchema(
+                processor=self.processor,
+                allows_deferred_tools=self.allows_deferred_tools,
+                allows_image=self.allows_image,
+            )
         elif mode == 'tool':
-            return ToolOutputSchema(toolset=self.toolset, allows_deferred_tools=self.allows_deferred_tools)
+            return ToolOutputSchema(
+                toolset=self.toolset, allows_deferred_tools=self.allows_deferred_tools, allows_image=self.allows_image
+            )
         else:
             assert_never(mode)
 
-    @property
-    def toolset(self) -> OutputToolset[Any] | None:
-        """Get the toolset for this output schema."""
-        # We return a toolset here as they're checked for name conflicts with other toolsets in the Agent constructor.
-        # At that point we may not know yet what output mode we're going to use if no model was provided or it was deferred until agent.run time,
-        # but we cover ourselves just in case we end up using the tool output mode.
-        return self._toolset
-
 
-class TextOutputSchema(OutputSchema[OutputDataT], ABC):
-    @abstractmethod
-    async def process(
+@dataclass(init=False)
+class TextOutputSchema(OutputSchema[OutputDataT]):
+    def __init__(
         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
+        *,
+        text_processor: TextOutputProcessor[OutputDataT],
+        allows_deferred_tools: bool,
+        allows_image: bool,
+    ):
+        super().__init__(
+            text_processor=text_processor,
+            allows_deferred_tools=allows_deferred_tools,
+            allows_image=allows_image,
+        )
 
     @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
+        """Raise an error if the mode is not supported by this model."""
+        super().raise_if_unsupported(profile)
 
-    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.
+class ImageOutputSchema(OutputSchema[OutputDataT]):
+    def __init__(self, *, allows_deferred_tools: bool):
+        super().__init__(allows_deferred_tools=allows_deferred_tools, allows_image=True)
 
-        Returns:
-            Either the validated output data (left) or a retry message (right).
-        """
-        if self.processor is None:
-            return cast(OutputDataT, text)
+    @property
+    def mode(self) -> OutputMode:
+        return 'image'
 
-        return await self.processor.process(
-            text, run_context, allow_partial=allow_partial, wrap_validation_errors=wrap_validation_errors
-        )
+    def raise_if_unsupported(self, profile: ModelProfile) -> None:
+        """Raise an error if the mode is not supported by this model."""
+        # This already raises if image output is not supported by this model.
+        super().raise_if_unsupported(profile)
 
 
-@dataclass
-class StructuredTextOutputSchema(TextOutputSchema[OutputDataT], ABC):
-    processor: ObjectOutputProcessor[OutputDataT] | UnionOutputProcessor[OutputDataT]
+@dataclass(init=False)
+class StructuredTextOutputSchema(OutputSchema[OutputDataT], ABC):
+    processor: BaseObjectOutputProcessor[OutputDataT]
+
+    def __init__(
+        self, *, processor: BaseObjectOutputProcessor[OutputDataT], allows_deferred_tools: bool, allows_image: bool
+    ):
+        super().__init__(
+            text_processor=processor, allows_deferred_tools=allows_deferred_tools, allows_image=allows_image
+        )
+        self.processor = processor
 
     @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."""
+        """Raise an error if the mode is not supported by this model."""
         if not profile.supports_json_schema_output:
-            raise UserError('Native structured output is not supported by the model.')
+            raise UserError('Native structured output is not supported by this 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.
+@dataclass(init=False)
+class PromptedOutputSchema(StructuredTextOutputSchema[OutputDataT]):
+    template: str | None
 
-        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
+    def __init__(
+        self,
+        *,
+        template: str | None = None,
+        processor: BaseObjectOutputProcessor[OutputDataT],
+        allows_deferred_tools: bool,
+        allows_image: bool,
+    ):
+        super().__init__(
+            processor=PromptedOutputProcessor(processor),
+            allows_deferred_tools=allows_deferred_tools,
+            allows_image=allows_image,
         )
-
-
-@dataclass
-class PromptedOutputSchema(StructuredTextOutputSchema[OutputDataT]):
-    template: str | None = None
+        self.template = template
 
     @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
+        """Raise an error if the mode is not supported by this model."""
+        super().raise_if_unsupported(profile)
 
     def instructions(self, default_template: str) -> str:
         """Get instructions to tell model to output JSON matching the schema."""
@@ -532,71 +574,35 @@ class PromptedOutputSchema(StructuredTextOutputSchema[OutputDataT]):
 
         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]):
-    _toolset: OutputToolset[Any] | None
-
-    def __init__(self, toolset: OutputToolset[Any] | None, allows_deferred_tools: bool):
-        super().__init__(allows_deferred_tools)
-        self._toolset = toolset
-
-    @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 toolset(self) -> OutputToolset[Any] | None:
-        """Get the toolset for this output schema."""
-        return self._toolset
-
-
-@dataclass(init=False)
-class ToolOrTextOutputSchema(ToolOutputSchema[OutputDataT], PlainTextOutputSchema[OutputDataT]):
     def __init__(
         self,
-        processor: PlainTextOutputProcessor[OutputDataT] | None,
+        *,
         toolset: OutputToolset[Any] | None,
+        text_processor: BaseOutputProcessor[OutputDataT] | None = None,
         allows_deferred_tools: bool,
+        allows_image: bool,
     ):
-        super().__init__(toolset=toolset, allows_deferred_tools=allows_deferred_tools)
-        self.processor = processor
+        super().__init__(
+            toolset=toolset,
+            allows_deferred_tools=allows_deferred_tools,
+            text_processor=text_processor,
+            allows_image=allows_image,
+        )
 
     @property
     def mode(self) -> OutputMode:
-        return 'tool_or_text'
+        return 'tool'
+
+    def raise_if_unsupported(self, profile: ModelProfile) -> None:
+        """Raise an error if the mode is not supported by this model."""
+        super().raise_if_unsupported(profile)
+        if not profile.supports_tools:
+            raise UserError('Tool output is not supported by this model.')
 
 
-@dataclass(init=False)
 class BaseOutputProcessor(ABC, Generic[OutputDataT]):
     @abstractmethod
     async def process(
@@ -610,9 +616,35 @@ class BaseOutputProcessor(ABC, Generic[OutputDataT]):
         raise NotImplementedError()
 
 
-@dataclass(init=False)
-class ObjectOutputProcessor(BaseOutputProcessor[OutputDataT]):
+@dataclass(kw_only=True)
+class BaseObjectOutputProcessor(BaseOutputProcessor[OutputDataT]):
     object_def: OutputObjectDefinition
+
+
+@dataclass(init=False)
+class PromptedOutputProcessor(BaseObjectOutputProcessor[OutputDataT]):
+    wrapped: BaseObjectOutputProcessor[OutputDataT]
+
+    def __init__(self, wrapped: BaseObjectOutputProcessor[OutputDataT]):
+        self.wrapped = wrapped
+        super().__init__(object_def=wrapped.object_def)
+
+    async def process(
+        self,
+        data: str,
+        run_context: RunContext[AgentDepsT],
+        allow_partial: bool = False,
+        wrap_validation_errors: bool = True,
+    ) -> OutputDataT:
+        text = _utils.strip_markdown_fences(data)
+
+        return await self.wrapped.process(
+            text, run_context, allow_partial=allow_partial, wrap_validation_errors=wrap_validation_errors
+        )
+
+
+@dataclass(init=False)
+class ObjectOutputProcessor(BaseObjectOutputProcessor[OutputDataT]):
     outer_typed_dict_key: str | None = None
     validator: SchemaValidator
     _function_schema: _function_schema.FunctionSchema | None = None
@@ -673,11 +705,13 @@ class ObjectOutputProcessor(BaseOutputProcessor[OutputDataT]):
             else:
                 description = f'{description}. {json_schema_description}'
 
-        self.object_def = OutputObjectDefinition(
-            name=name or getattr(output, '__name__', None),
-            description=description,
-            json_schema=json_schema,
-            strict=strict,
+        super().__init__(
+            object_def=OutputObjectDefinition(
+                name=name or getattr(output, '__name__', None),
+                description=description,
+                json_schema=json_schema,
+                strict=strict,
+            )
         )
 
     async def process(
@@ -726,10 +760,10 @@ class ObjectOutputProcessor(BaseOutputProcessor[OutputDataT]):
 
     async def call(
         self,
-        output: Any,
+        output: dict[str, Any],
         run_context: RunContext[AgentDepsT],
         wrap_validation_errors: bool = True,
-    ):
+    ) -> Any:
         if k := self.outer_typed_dict_key:
             output = output[k]
 
@@ -753,8 +787,7 @@ class UnionOutputModel:
 
 
 @dataclass(init=False)
-class UnionOutputProcessor(BaseOutputProcessor[OutputDataT]):
-    object_def: OutputObjectDefinition
+class UnionOutputProcessor(BaseObjectOutputProcessor[OutputDataT]):
     _union_processor: ObjectOutputProcessor[UnionOutputModel]
     _processors: dict[str, ObjectOutputProcessor[OutputDataT]]
 
@@ -830,16 +863,18 @@ class UnionOutputProcessor(BaseOutputProcessor[OutputDataT]):
         if all_defs:
             json_schema['$defs'] = all_defs
 
-        self.object_def = OutputObjectDefinition(
-            json_schema=json_schema,
-            strict=strict,
-            name=name,
-            description=description,
+        super().__init__(
+            object_def=OutputObjectDefinition(
+                json_schema=json_schema,
+                strict=strict,
+                name=name,
+                description=description,
+            )
         )
 
     async def process(
         self,
-        data: str | dict[str, Any] | None,
+        data: str,
         run_context: RunContext[AgentDepsT],
         allow_partial: bool = False,
         wrap_validation_errors: bool = True,
@@ -850,7 +885,7 @@ class UnionOutputProcessor(BaseOutputProcessor[OutputDataT]):
 
         result = union_object.result
         kind = result.kind
-        data = result.data
+        inner_data = result.data
         try:
             processor = self._processors[kind]
         except KeyError as e:  # pragma: no cover
@@ -861,12 +896,23 @@ class UnionOutputProcessor(BaseOutputProcessor[OutputDataT]):
                 raise
 
         return await processor.process(
-            data, run_context, allow_partial=allow_partial, wrap_validation_errors=wrap_validation_errors
+            inner_data, run_context, allow_partial=allow_partial, wrap_validation_errors=wrap_validation_errors
         )
 
 
+class TextOutputProcessor(BaseOutputProcessor[OutputDataT]):
+    async def process(
+        self,
+        data: str,
+        run_context: RunContext[AgentDepsT],
+        allow_partial: bool = False,
+        wrap_validation_errors: bool = True,
+    ) -> OutputDataT:
+        return cast(OutputDataT, data)
+
+
 @dataclass(init=False)
-class PlainTextOutputProcessor(BaseOutputProcessor[OutputDataT]):
+class TextFunctionOutputProcessor(TextOutputProcessor[OutputDataT]):
     _function_schema: _function_schema.FunctionSchema
     _str_argument_name: str
 
@@ -876,17 +922,15 @@ class PlainTextOutputProcessor(BaseOutputProcessor[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
+        if (
+            not (arguments_schema := self._function_schema.json_schema.get('properties', {}))
+            or len(arguments_schema) != 1
+            or not (argument_name := next(iter(arguments_schema.keys()), None))
+            or arguments_schema.get(argument_name, {}).get('type') != 'string'
+        ):
+            raise UserError('TextOutput must take a function taking a single `str` argument')
 
-        raise UserError('TextOutput must take a function taking a `str`')
-
-    @property
-    def object_def(self) -> None:
-        return None  # pragma: no cover
+        self._str_argument_name = argument_name
 
     async def process(
         self,
@@ -896,9 +940,9 @@ class PlainTextOutputProcessor(BaseOutputProcessor[OutputDataT]):
         wrap_validation_errors: bool = True,
     ) -> OutputDataT:
         args = {self._str_argument_name: data}
-        output = await execute_traced_output_function(self._function_schema, run_context, args, wrap_validation_errors)
+        data = await execute_traced_output_function(self._function_schema, run_context, args, wrap_validation_errors)
 
-        return cast(OutputDataT, output)
+        return await super().process(data, run_context, allow_partial, wrap_validation_errors)
 
 
 @dataclass(init=False)