tweaks and coverage

Author: DouweM

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -140,7 +140,7 @@ class GraphAgentDeps(Generic[DepsT, OutputDataT]):
     end_strategy: EndStrategy
     get_instructions: Callable[[RunContext[DepsT]], Awaitable[str | None]]
 
-    output_schema: _output.BaseOutputSchema[OutputDataT]
+    output_schema: _output.OutputSchema[OutputDataT]
     output_validators: list[_output.OutputValidator[DepsT, OutputDataT]]
 
     history_processors: Sequence[HistoryProcessor[DepsT]]

pydantic_ai_slim/pydantic_ai/_output.py

@@ -211,28 +211,21 @@ async def validate(
 
 
 @dataclass(kw_only=True)
-class BaseOutputSchema(ABC, Generic[OutputDataT]):
+class OutputSchema(ABC, Generic[OutputDataT]):
     text_processor: BaseOutputProcessor[OutputDataT] | None = None
     toolset: OutputToolset[Any] | None = None
     object_def: OutputObjectDefinition | None = None
     allows_deferred_tools: bool = False
     allows_image: bool = False
 
     @property
-    def mode(self) -> OutputMode | None:
+    def mode(self) -> OutputMode:
         raise NotImplementedError()
 
     @property
     def allows_text(self) -> bool:
         return self.text_processor is not None
 
-
-@dataclass(init=False)
-class OutputSchema(BaseOutputSchema[OutputDataT], ABC):
-    """Model the final output from an agent run."""
-
-    # TODO (DouweM): Rename/merge this, BaseOutputSchema, and OutputSchemaWithoutMode
-
     @classmethod
     def build(  # noqa: C901
         cls,
@@ -241,7 +234,7 @@ def build(  # noqa: C901
         name: str | None = None,
         description: str | None = None,
         strict: bool | None = None,
-    ) -> BaseOutputSchema[OutputDataT]:
+    ) -> OutputSchema[OutputDataT]:
         """Build an OutputSchema dataclass from an output type."""
         outputs = _flatten_output_spec(output_spec)
 
@@ -359,7 +352,7 @@ def build(  # noqa: C901
             )
 
         if len(other_outputs) > 0:
-            return OutputSchemaWithoutMode(
+            return AutoOutputSchema(
                 processor=cls._build_processor(other_outputs, name=name, description=description, strict=strict),
                 toolset=toolset,
                 allows_deferred_tools=allows_deferred_tools,
@@ -386,7 +379,7 @@ def _build_processor(
 
 
 @dataclass(init=False)
-class OutputSchemaWithoutMode(BaseOutputSchema[OutputDataT]):
+class AutoOutputSchema(OutputSchema[OutputDataT]):
     processor: BaseObjectOutputProcessor[OutputDataT]
 
     def __init__(
@@ -400,18 +393,17 @@ def __init__(
         # 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,
             object_def=processor.object_def,
             text_processor=processor,
+            allows_deferred_tools=allows_deferred_tools,
             allows_image=allows_image,
         )
         self.processor = processor
 
     @property
-    def mode(self) -> OutputMode | None:
-        # TODO (DouweM): Could this be a field?
-        return None
+    def mode(self) -> OutputMode:
+        return 'auto'
 
 
 @dataclass(init=False)
@@ -430,7 +422,7 @@ def __init__(
         )
 
     @property
-    def mode(self) -> OutputMode | None:
+    def mode(self) -> OutputMode:
         return 'text'
 
 
@@ -439,7 +431,7 @@ def __init__(self, *, allows_deferred_tools: bool):
         super().__init__(allows_deferred_tools=allows_deferred_tools, allows_image=True)
 
     @property
-    def mode(self) -> OutputMode | None:
+    def mode(self) -> OutputMode:
         return 'image'
 
 
@@ -461,7 +453,7 @@ def __init__(
 
 class NativeOutputSchema(StructuredTextOutputSchema[OutputDataT]):
     @property
-    def mode(self) -> OutputMode | None:
+    def mode(self) -> OutputMode:
         return 'native'
 
 
@@ -485,7 +477,7 @@ def __init__(
         self.template = template
 
     @property
-    def mode(self) -> OutputMode | None:
+    def mode(self) -> OutputMode:
         return 'prompted'
 
     @classmethod
@@ -528,7 +520,7 @@ def __init__(
         )
 
     @property
-    def mode(self) -> OutputMode | None:
+    def mode(self) -> OutputMode:
         return 'tool'
 
 

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -132,7 +132,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
     _instrument_default: ClassVar[InstrumentationSettings | bool] = False
 
     _deps_type: type[AgentDepsT] = dataclasses.field(repr=False)
-    _output_schema: _output.BaseOutputSchema[OutputDataT] = dataclasses.field(repr=False)
+    _output_schema: _output.OutputSchema[OutputDataT] = dataclasses.field(repr=False)
     _output_validators: list[_output.OutputValidator[AgentDepsT, OutputDataT]] = dataclasses.field(repr=False)
     _instructions: list[str | _system_prompt.SystemPromptFunc[AgentDepsT]] = dataclasses.field(repr=False)
     _system_prompts: tuple[str, ...] = dataclasses.field(repr=False)
@@ -1409,14 +1409,14 @@ def toolsets(self) -> Sequence[AbstractToolset[AgentDepsT]]:
         return toolsets
 
     @overload
-    def _prepare_output_schema(self, output_type: None) -> _output.BaseOutputSchema[OutputDataT]: ...
+    def _prepare_output_schema(self, output_type: None) -> _output.OutputSchema[OutputDataT]: ...
 
     @overload
     def _prepare_output_schema(
         self, output_type: OutputSpec[RunOutputDataT]
-    ) -> _output.BaseOutputSchema[RunOutputDataT]: ...
+    ) -> _output.OutputSchema[RunOutputDataT]: ...
 
-    def _prepare_output_schema(self, output_type: OutputSpec[Any] | None) -> _output.BaseOutputSchema[Any]:
+    def _prepare_output_schema(self, output_type: OutputSpec[Any] | None) -> _output.OutputSchema[Any]:
         if output_type is not None:
             if self._output_validators:
                 raise exceptions.UserError('Cannot set a custom run `output_type` when the agent has output validators')
@@ -1494,15 +1494,15 @@ async def run_mcp_servers(
 
 @dataclasses.dataclass(init=False)
 class _AgentFunctionToolset(FunctionToolset[AgentDepsT]):
-    output_schema: _output.BaseOutputSchema[Any]
+    output_schema: _output.OutputSchema[Any]
 
     def __init__(
         self,
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = [],
         *,
         max_retries: int = 1,
         id: str | None = None,
-        output_schema: _output.BaseOutputSchema[Any],
+        output_schema: _output.OutputSchema[Any],
     ):
         self.output_schema = output_schema
         super().__init__(tools, max_retries=max_retries, id=id)

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -306,7 +306,7 @@ class ModelRequestParameters:
     function_tools: list[ToolDefinition] = field(default_factory=list)
     builtin_tools: list[AbstractBuiltinTool] = field(default_factory=list)
 
-    output_mode: OutputMode | None = 'text'  # TODO (DouweM): None or new `'auto'` value? AutoOutputSchema is not bad.
+    output_mode: OutputMode = 'text'
     output_object: OutputObjectDefinition | None = None
     output_tools: list[ToolDefinition] = field(default_factory=list)
     prompted_output_template: str | None = None
@@ -426,8 +426,7 @@ def prepare_request(
                 builtin_tools=list({tool.unique_id: tool for tool in builtin_tools}.values()),
             )
 
-        if not model_request_parameters.output_mode:
-            assert model_request_parameters.output_tools or model_request_parameters.output_object
+        if model_request_parameters.output_mode == 'auto':
             output_mode = self.profile.default_structured_output_mode
             model_request_parameters = replace(
                 model_request_parameters,
@@ -436,31 +435,33 @@ def prepare_request(
             )
 
         if model_request_parameters.output_mode in ('native', 'prompted'):
-            if not model_request_parameters.output_object:
-                raise UserError(  # pragma: no cover
-                    'An `output_object` is required when using `NativeOutput` or `PromptedOutput`.'
-                )
-
-            if model_request_parameters.output_mode == 'native' and not self.profile.supports_json_schema_output:
-                raise UserError('Native structured output is not supported by this model.')
+            assert model_request_parameters.output_object
 
             if model_request_parameters.output_tools:
                 model_request_parameters = replace(model_request_parameters, output_tools=[])
-
-            if not model_request_parameters.prompted_output_template:
-                model_request_parameters = replace(
-                    model_request_parameters, prompted_output_template=self.profile.prompted_output_template
-                )
         else:
-            if model_request_parameters.output_mode == 'tool':
-                if not model_request_parameters.output_tools and not model_request_parameters.function_tools:
-                    raise UserError('An `output_tools` list is required when using `ToolOutput`.')  # pragma: no cover
+            if model_request_parameters.output_object:
+                model_request_parameters = replace(model_request_parameters, output_object=None)
+
+        match model_request_parameters.output_mode:
+            case 'native':
+                if not self.profile.supports_json_schema_output:
+                    raise UserError('Native structured output is not supported by this model.')
+
+                if model_request_parameters.prompted_output_template:
+                    model_request_parameters = replace(model_request_parameters, prompted_output_template=None)
+            case 'prompted':
+                if not model_request_parameters.prompted_output_template:
+                    model_request_parameters = replace(
+                        model_request_parameters, prompted_output_template=self.profile.prompted_output_template
+                    )
+            case 'tool':
+                assert model_request_parameters.output_tools or model_request_parameters.function_tools
 
                 if not self.profile.supports_tools:
                     raise UserError('Tool output is not supported by this model.')
-
-            if model_request_parameters.output_object:
-                model_request_parameters = replace(model_request_parameters, output_object=None)
+            case _:
+                pass
 
         if model_request_parameters.allow_image_output and not self.profile.supports_image_output:
             raise UserError('Image output is not supported by this model.')

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -243,15 +243,14 @@ async def request_stream(
     def prepare_request(
         self, model_settings: ModelSettings | None, model_request_parameters: ModelRequestParameters
     ) -> tuple[ModelSettings | None, ModelRequestParameters]:
-        # TODO (DouweM): Dedupe with super
         settings = merge_model_settings(self.settings, model_settings)
         if (
             model_request_parameters.output_tools
             and settings
             and (thinking := settings.get('anthropic_thinking'))
             and thinking.get('type') == 'enabled'
         ):
-            if model_request_parameters.output_mode is None:
+            if model_request_parameters.output_mode == 'auto':
                 model_request_parameters = replace(model_request_parameters, output_mode='prompted')
             elif model_request_parameters.output_mode == 'tool' and not model_request_parameters.allow_text_output:
                 # This would result in `tool_choice=required`, which Anthropic does not support with thinking.

pydantic_ai_slim/pydantic_ai/models/google.py

@@ -228,7 +228,7 @@ def prepare_request(
         self, model_settings: ModelSettings | None, model_request_parameters: ModelRequestParameters
     ) -> tuple[ModelSettings | None, ModelRequestParameters]:
         if model_request_parameters.builtin_tools and model_request_parameters.output_tools:
-            if model_request_parameters.output_mode is None:
+            if model_request_parameters.output_mode == 'auto':
                 model_request_parameters = replace(model_request_parameters, output_mode='prompted')
             else:
                 raise UserError(

pydantic_ai_slim/pydantic_ai/models/outlines.py

@@ -528,7 +528,7 @@ async def _process_streamed_response(
 
     def customize_request_parameters(self, model_request_parameters: ModelRequestParameters) -> ModelRequestParameters:
         """Customize the model request parameters for the model."""
-        if model_request_parameters.output_mode in (None, 'native'):
+        if model_request_parameters.output_mode in ('auto', 'native'):
             # This way the JSON schema will be included in the instructions.
             return replace(model_request_parameters, output_mode='prompted')
         else:

pydantic_ai_slim/pydantic_ai/output.py

@@ -37,10 +37,11 @@
 OutputDataT = TypeVar('OutputDataT', default=str, covariant=True)
 """Covariant type variable for the output data type of a run."""
 
-OutputMode = Literal['text', 'tool', 'native', 'prompted', 'tool_or_text', 'image']
+OutputMode = Literal['text', 'tool', 'native', 'prompted', 'tool_or_text', 'image', 'auto']
 """All output modes.
 
-`tool_or_text` is deprecated and no longer in use.
+- `tool_or_text` is deprecated and no longer in use.
+- `auto` means the model will automatically choose a structured output mode based on the model's `ModelProfile.default_structured_output_mode`.
 """
 StructuredOutputMode = Literal['tool', 'native', 'prompted']
 """Output modes that can be used for structured output. Used by ModelProfile.default_structured_output_mode"""

pydantic_ai_slim/pydantic_ai/result.py

@@ -11,8 +11,8 @@
 
 from . import _utils, exceptions, messages as _messages, models
 from ._output import (
-    BaseOutputSchema,
     OutputDataT_inv,
+    OutputSchema,
     OutputValidator,
     OutputValidatorFunc,
     TextOutputSchema,
@@ -46,7 +46,7 @@
 @dataclass(kw_only=True)
 class AgentStream(Generic[AgentDepsT, OutputDataT]):
     _raw_stream_response: models.StreamedResponse
-    _output_schema: BaseOutputSchema[OutputDataT]
+    _output_schema: OutputSchema[OutputDataT]
     _model_request_parameters: models.ModelRequestParameters
     _output_validators: list[OutputValidator[AgentDepsT, OutputDataT]]
     _run_ctx: RunContext[AgentDepsT]

tests/test_agent.py

@@ -3791,6 +3791,7 @@ def get_image() -> BinaryContent:
                         BinaryContent(
                             data=b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde\x00\x00\x00\x0cIDATx\x9cc```\x00\x00\x00\x04\x00\x01\xf6\x178\x00\x00\x00\x00IEND\xaeB`\x82',
                             media_type='image/png',
+                            _identifier='image_id_1',
                         ),
                     ],
                     timestamp=IsNow(tz=timezone.utc),