WIP: More output modes

Author: DouweM

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -7,7 +7,7 @@
 from contextlib import asynccontextmanager, contextmanager
 from contextvars import ContextVar
 from dataclasses import field
-from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Union, cast
+from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Union
 
 from opentelemetry.trace import Tracer
 from typing_extensions import TypeGuard, TypeVar, assert_never
@@ -90,7 +90,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]]
 
     function_tools: dict[str, Tool[DepsT]] = dataclasses.field(repr=False)
@@ -264,29 +264,14 @@ 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
-    model = ctx.deps.model
-
-    # TODO: This is horrible
-    output_mode = None
-    output_object = None
-    output_tools = []
-    require_tool_use = False
-    if output_schema:
-        output_mode = output_schema.forced_mode or model.default_output_mode
-        output_object = output_schema.object_schema.definition
-        output_tools = output_schema.tool_defs()
-        require_tool_use = output_mode == 'tool' and output_schema.allow_text_output != 'plain'
-
-        supported_modes = model.supported_output_modes
-        if output_mode not in supported_modes:
-            raise exceptions.UserError(f"Output mode '{output_mode}' is not among supported modes: {supported_modes}")
+    assert output_schema.mode is not None  # Should have been set in agent._prepare_output_schema
 
     return models.ModelRequestParameters(
         function_tools=function_tool_defs,
-        output_mode=output_mode,
-        output_object=output_object,
-        output_tools=output_tools,
-        require_tool_use=require_tool_use,
+        output_mode=output_schema.mode,
+        output_object=output_schema.object_schema.definition if output_schema.object_schema else None,
+        output_tools=output_schema.tool_defs(),
+        allow_text_output=output_schema.allow_text_output == 'plain',
     )
 
 
@@ -471,7 +456,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 ctx.deps.output_schema.allow_text_output:
                         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)]
@@ -497,19 +482,18 @@ async def _handle_tool_calls(
         # first, look for the output tool call
         final_result: result.FinalResult[NodeRunEndT] | None = None
         parts: list[_messages.ModelRequestPart] = []
-        if output_schema is not None:
-            for call, output_tool in output_schema.find_tool(tool_calls):
-                try:
-                    result_data = await output_tool.process(call, run_context)
-                    result_data = await _validate_output(result_data, ctx, call)
-                except _output.ToolRetryError as e:
-                    # TODO: Should only increment retry stuff once per node execution, not for each tool call
-                    #   Also, should increment the tool-specific retry count rather than the run retry count
-                    ctx.state.increment_retries(ctx.deps.max_result_retries, e)
-                    parts.append(e.tool_retry)
-                else:
-                    final_result = result.FinalResult(result_data, call.tool_name, call.tool_call_id)
-                    break
+        for call, output_tool in output_schema.find_tool(tool_calls):
+            try:
+                result_data = await output_tool.process(call, run_context)
+                result_data = await _validate_output(result_data, ctx, call)
+            except _output.ToolRetryError as e:
+                # TODO: Should only increment retry stuff once per node execution, not for each tool call
+                #   Also, should increment the tool-specific retry count rather than the run retry count
+                ctx.state.increment_retries(ctx.deps.max_result_retries, e)
+                parts.append(e.tool_retry)
+            else:
+                final_result = result.FinalResult(result_data, call.tool_name, call.tool_call_id)
+                break
 
         # Then build the other request parts based on end strategy
         tool_responses: list[_messages.ModelRequestPart] = self._tool_responses
@@ -555,10 +539,7 @@ async def _handle_text_response(
 
         text = '\n\n'.join(texts)
         try:
-            if output_schema is None or output_schema.allow_text_output == 'plain':
-                # The following cast is safe because we know `str` is an allowed result type
-                result_data = cast(NodeRunEndT, text)
-            elif output_schema.allow_text_output == 'json':
+            if output_schema.allow_text_output:
                 run_context = build_run_context(ctx)
                 result_data = await output_schema.process(text, run_context)
             else:
@@ -659,7 +640,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:
@@ -788,8 +769,7 @@ 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:
-        tool_names.extend(output_schema.tool_names())
+    tool_names.extend(ctx.deps.output_schema.tool_names())
 
     if tool_names:
         msg = f'Available tools: {", ".join(tool_names)}'

pydantic_ai_slim/pydantic_ai/_output.py

@@ -13,6 +13,8 @@
 from typing_inspection import typing_objects
 from typing_inspection.introspection import is_union_origin
 
+from pydantic_ai.profiles import ModelProfile
+
 from . import _function_schema, _utils, messages as _messages
 from .exceptions import ModelRetry
 from .tools import AgentDepsT, GenerateToolJsonSchema, ObjectJsonSchema, RunContext, ToolDefinition
@@ -208,7 +210,7 @@ def __init__(
 )
 
 # TODO: Add `json_object` for old OpenAI models, or rename `json_schema` to `json` and choose automatically, relying on Pydantic validation
-type OutputMode = Literal['tool', 'json_schema', 'manual_json']
+type OutputMode = Literal['text', 'tool', 'tool_or_text', 'json_schema', 'manual_json']
 
 
 @dataclass
@@ -218,50 +220,46 @@ class OutputSchema(Generic[OutputDataT]):
     Similar to `Tool` but for the final output of running an agent.
     """
 
-    forced_mode: OutputMode | None
-    object_schema: OutputObjectSchema[OutputDataT] | OutputUnionSchema[OutputDataT]
-    tools: dict[str, OutputTool[OutputDataT]]
-    allow_text_output: Literal['plain', 'json'] | None = None
+    mode: OutputMode | None
+    object_schema: OutputObjectSchema[OutputDataT] | OutputUnionSchema[OutputDataT] | None = None
+    tools: dict[str, OutputTool[OutputDataT]] = field(default_factory=dict)
 
     @classmethod
     def build(
         cls: type[OutputSchema[OutputDataT]],
         output_type: OutputType[OutputDataT],
-        name: str | None = None,
-        description: str | None = None,
-        strict: bool | None = None,
-    ) -> OutputSchema[OutputDataT] | None:
+        name: str | None,
+        description: str | None,
+    ) -> OutputSchema[OutputDataT]:
         """Build an OutputSchema dataclass from an output type."""
         if output_type is str:
-            return None
+            return cls(mode='text')
 
-        forced_mode: OutputMode | None = None
-        allow_text_output: Literal['plain', 'json'] | None = 'plain'
+        mode: OutputMode | None = None
         tools: dict[str, OutputTool[OutputDataT]] = {}
+        strict: bool | None = None
 
         output_types: Sequence[OutputTypeOrFunction[OutputDataT]]
         if isinstance(output_type, JSONSchemaOutput):
-            forced_mode = 'json_schema'
+            mode = 'json_schema'
             output_types = output_type.output_types
             name = output_type.name  # TODO: If not set, use method arg?
             description = output_type.description
             strict = output_type.strict
-            allow_text_output = 'json'
         elif isinstance(output_type, ManualJSONOutput):
-            forced_mode = 'manual_json'
+            mode = 'manual_json'
             output_types = output_type.output_types
             name = output_type.name
             description = output_type.description
-            allow_text_output = 'json'
         else:
-            # TODO: We can't always force tool mode here, because some models may not support tools but will work with manual_json
             output_types_or_tool_outputs = flatten_output_types(output_type)
 
             if str in output_types_or_tool_outputs:
-                forced_mode = 'tool'
-                allow_text_output = 'plain'
-                # TODO: What if str is the only item, e.g. `output_type=[str]`
-                output_types_or_tool_outputs = [t for t in output_types_or_tool_outputs if t is not str]
+                if len(output_types_or_tool_outputs) == 1:
+                    return cls(mode='text')
+                else:
+                    mode = 'tool_or_text'
+                    output_types_or_tool_outputs = [t for t in output_types_or_tool_outputs if t is not str]
 
             multiple = len(output_types_or_tool_outputs) > 1
 
@@ -275,7 +273,9 @@ def build(
                 tool_description = None
                 tool_strict = None
                 if isinstance(output_type_or_tool_output, ToolOutput):
-                    forced_mode = 'tool'
+                    if mode is None:
+                        mode = 'tool'
+
                     tool_output = output_type_or_tool_output
                     output_type = tool_output.output_type
                     # do we need to error on conflicts here? (DavidM): If this is internal maybe doesn't matter, if public, use overloads
@@ -307,7 +307,6 @@ def build(
                 output_types.append(output_type)
 
         output_types = flatten_output_types(output_types)
-
         if len(output_types) > 1:
             output_object_schema = OutputUnionSchema(
                 output_types=output_types, name=name, description=description, strict=strict
@@ -318,12 +317,30 @@ def build(
             )
 
         return cls(
-            forced_mode=forced_mode,
+            mode=mode,
             object_schema=output_object_schema,
             tools=tools,
-            allow_text_output=allow_text_output,
         )
 
+    @property
+    def allow_text_output(self) -> Literal['plain', 'json', False]:
+        """Whether the model allows text output."""
+        if self.mode in ('text', 'tool_or_text'):
+            return 'plain'
+        elif self.mode in ('json_schema', 'manual_json'):
+            return 'json'
+        else:  # tool-only mode
+            return False
+
+    def is_mode_supported(self, profile: ModelProfile) -> bool:
+        """Whether the model supports the output mode."""
+        mode = self.mode
+        if mode in ('text', 'manual_json'):
+            return True
+        if self.mode == 'tool_or_text':
+            mode = 'tool'
+        return mode in profile.output_modes
+
     def find_named_tool(
         self, parts: Iterable[_messages.ModelResponsePart], tool_name: str
     ) -> tuple[_messages.ToolCallPart, OutputTool[OutputDataT]] | None:
@@ -369,16 +386,18 @@ async def process(
         Returns:
             Either the validated output data (left) or a retry message (right).
         """
+        assert self.allow_text_output is not False
+
+        if self.allow_text_output == 'plain':
+            return cast(OutputDataT, data)
+
+        assert self.object_schema is not None
+
         return await self.object_schema.process(
             data, run_context, allow_partial=allow_partial, wrap_validation_errors=wrap_validation_errors
         )
 
 
-def allow_text_output(output_schema: OutputSchema[Any] | None) -> bool:
-    # TODO: Add plain/json argument?
-    return output_schema is None or output_schema.allow_text_output is not None
-
-
 @dataclass
 class OutputObjectDefinition:
     name: str
@@ -389,6 +408,7 @@ class OutputObjectDefinition:
     @property
     def manual_json_instructions(self) -> str:
         """Get instructions for model to output manual JSON matching the schema."""
+        # TODO: Move to ModelProfile so it can be tweaked
         description = ': '.join([v for v in [self.name, self.description] if v])
         return DEFAULT_MANUAL_JSON_PROMPT.format(schema=json.dumps(self.json_schema), description=description)