Support Native and Prompted output modes when using FallbackModel

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.OutputSchema[OutputDataT]
+    output_schema: _output.BaseOutputSchema[OutputDataT]
     output_validators: list[_output.OutputValidator[DepsT, OutputDataT]]
 
     history_processors: Sequence[HistoryProcessor[DepsT]]
@@ -374,9 +374,10 @@ async def _prepare_request_parameters(
 ) -> models.ModelRequestParameters:
     """Build tools and create an agent model."""
     output_schema = ctx.deps.output_schema
-    output_object = None
-    if isinstance(output_schema, _output.NativeOutputSchema):
-        output_object = output_schema.object_def
+
+    prompted_output_template = (
+        output_schema.template if isinstance(output_schema, _output.PromptedOutputSchema) else None
+    )
 
     function_tools: list[ToolDefinition] = []
     output_tools: list[ToolDefinition] = []
@@ -391,7 +392,8 @@ async def _prepare_request_parameters(
         builtin_tools=ctx.deps.builtin_tools,
         output_mode=output_schema.mode,
         output_tools=output_tools,
-        output_object=output_object,
+        output_object=output_schema.object_def,
+        prompted_output_template=prompted_output_template,
         allow_text_output=output_schema.allows_text,
         allow_image_output=output_schema.allows_image,
     )
@@ -489,7 +491,6 @@ async def _prepare_request(
         message_history = _clean_message_history(message_history)
 
         model_request_parameters = await _prepare_request_parameters(ctx)
-        model_request_parameters = ctx.deps.model.customize_request_parameters(model_request_parameters)
 
         model_settings = ctx.deps.model_settings
         usage = ctx.state.usage

pydantic_ai_slim/pydantic_ai/_output.py

@@ -10,7 +10,7 @@
 
 from pydantic import Json, TypeAdapter, ValidationError
 from pydantic_core import SchemaValidator, to_json
-from typing_extensions import Self, TypedDict, TypeVar, assert_never
+from typing_extensions import Self, TypedDict, TypeVar
 
 from pydantic_ai._instrumentation import InstrumentationNames
 
@@ -215,11 +215,12 @@ async def validate(
 class BaseOutputSchema(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
 
-    @abstractmethod
-    def with_default_mode(self, mode: StructuredOutputMode) -> OutputSchema[OutputDataT]:
+    @property
+    def mode(self) -> OutputMode | None:
         raise NotImplementedError()
 
     @property
@@ -231,6 +232,8 @@ def allows_text(self) -> bool:
 class OutputSchema(BaseOutputSchema[OutputDataT], ABC):
     """Model the final output from an agent run."""
 
+    # TODO (DouweM): Rename/merge this, BaseOutputSchema, and OutputSchemaWithoutMode
+
     @classmethod
     @overload
     def build(
@@ -260,7 +263,6 @@ def build(  # noqa: C901
         cls,
         output_spec: OutputSpec[OutputDataT],
         *,
-        default_mode: StructuredOutputMode | None = None,
         name: str | None = None,
         description: str | None = None,
         strict: bool | None = None,
@@ -382,15 +384,12 @@ def build(  # noqa: C901
             )
 
         if len(other_outputs) > 0:
-            schema = OutputSchemaWithoutMode(
+            return 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)
@@ -410,19 +409,12 @@ def _build_processor(
 
         return UnionOutputProcessor(outputs=outputs, strict=strict, name=name, description=description)
 
-    @property
-    @abstractmethod
-    def mode(self) -> OutputMode:
-        raise NotImplementedError()
-
     def raise_if_unsupported(self, profile: ModelProfile) -> None:
         """Raise an error if the mode is not supported by this model."""
+        # TODO (DouweM): Remove method?
         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
-
 
 @dataclass(init=False)
 class OutputSchemaWithoutMode(BaseOutputSchema[OutputDataT]):
@@ -441,30 +433,16 @@ def __init__(
         super().__init__(
             allows_deferred_tools=allows_deferred_tools,
             toolset=toolset,
+            object_def=processor.object_def,
             text_processor=processor,
             allows_image=allows_image,
         )
         self.processor = processor
 
-    def with_default_mode(self, mode: StructuredOutputMode) -> OutputSchema[OutputDataT]:
-        if mode == 'native':
-            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,
-                allows_image=self.allows_image,
-            )
-        elif mode == 'tool':
-            return ToolOutputSchema(
-                toolset=self.toolset, allows_deferred_tools=self.allows_deferred_tools, allows_image=self.allows_image
-            )
-        else:
-            assert_never(mode)
+    @property
+    def mode(self) -> OutputMode | None:
+        # TODO (DouweM): Could this be a field?
+        return None
 
 
 @dataclass(init=False)
@@ -483,7 +461,7 @@ def __init__(
         )
 
     @property
-    def mode(self) -> OutputMode:
+    def mode(self) -> OutputMode | None:
         return 'text'
 
     def raise_if_unsupported(self, profile: ModelProfile) -> None:
@@ -496,7 +474,7 @@ def __init__(self, *, allows_deferred_tools: bool):
         super().__init__(allows_deferred_tools=allows_deferred_tools, allows_image=True)
 
     @property
-    def mode(self) -> OutputMode:
+    def mode(self) -> OutputMode | None:
         return 'image'
 
     def raise_if_unsupported(self, profile: ModelProfile) -> None:
@@ -513,18 +491,17 @@ 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
+            text_processor=processor,
+            object_def=processor.object_def,
+            allows_deferred_tools=allows_deferred_tools,
+            allows_image=allows_image,
         )
         self.processor = processor
 
-    @property
-    def object_def(self) -> OutputObjectDefinition:
-        return self.processor.object_def
-
 
 class NativeOutputSchema(StructuredTextOutputSchema[OutputDataT]):
     @property
-    def mode(self) -> OutputMode:
+    def mode(self) -> OutputMode | None:
         return 'native'
 
     def raise_if_unsupported(self, profile: ModelProfile) -> None:
@@ -553,7 +530,7 @@ def __init__(
         self.template = template
 
     @property
-    def mode(self) -> OutputMode:
+    def mode(self) -> OutputMode | None:
         return 'prompted'
 
     @classmethod
@@ -599,7 +576,7 @@ def __init__(
         )
 
     @property
-    def mode(self) -> OutputMode:
+    def mode(self) -> OutputMode | None:
         return 'tool'
 
     def raise_if_unsupported(self, profile: ModelProfile) -> None:

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -303,11 +303,7 @@ def __init__(
 
         _utils.validate_empty_kwargs(_deprecated_kwargs)
 
-        default_output_mode = (
-            self.model.profile.default_structured_output_mode if isinstance(self.model, models.Model) else None
-        )
-
-        self._output_schema = _output.OutputSchema[OutputDataT].build(output_type, default_mode=default_output_mode)
+        self._output_schema = _output.OutputSchema[OutputDataT].build(output_type)
         self._output_validators = []
 
         self._instructions = self._normalize_instructions(instructions)
@@ -545,7 +541,7 @@ async def main():
         del model
 
         deps = self._get_deps(deps)
-        output_schema = self._prepare_output_schema(output_type, model_used.profile)
+        output_schema = self._prepare_output_schema(output_type)
 
         output_type_ = output_type or self.output_type
 
@@ -556,7 +552,7 @@ async def main():
 
         output_toolset = self._output_toolset
         if output_schema != self._output_schema or output_validators:
-            output_toolset = cast(OutputToolset[AgentDepsT], output_schema.toolset)
+            output_toolset = output_schema.toolset
             if output_toolset:
                 output_toolset.max_retries = self._max_result_retries
                 output_toolset.output_validators = output_validators
@@ -588,11 +584,6 @@ async def get_instructions(run_context: RunContext[AgentDepsT]) -> str | None:
                 *[await func.run(run_context) for func in instructions_functions],
             ]
 
-            model_profile = model_used.profile
-            if isinstance(output_schema, _output.PromptedOutputSchema):
-                instructions = output_schema.instructions(model_profile.prompted_output_template)
-                parts.append(instructions)
-
             parts = [p for p in parts if p]
             if not parts:
                 return None
@@ -1409,20 +1400,16 @@ def toolsets(self) -> Sequence[AbstractToolset[AgentDepsT]]:
         return toolsets
 
     def _prepare_output_schema(
-        self, output_type: OutputSpec[RunOutputDataT] | None, model_profile: ModelProfile
-    ) -> _output.OutputSchema[RunOutputDataT]:
+        self, output_type: OutputSpec[RunOutputDataT] | None
+    ) -> _output.BaseOutputSchema[RunOutputDataT]:
         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')
-            schema = _output.OutputSchema[RunOutputDataT].build(
-                output_type, default_mode=model_profile.default_structured_output_mode
-            )
+            schema = _output.OutputSchema[RunOutputDataT].build(output_type)
         else:
-            schema = self._output_schema.with_default_mode(model_profile.default_structured_output_mode)
-
-        schema.raise_if_unsupported(model_profile)
+            schema = self._output_schema
 
-        return schema  # pyright: ignore[reportReturnType]
+        return schema
 
     async def __aenter__(self) -> Self:
         """Enter the agent context.