WIP: Remove OutputPart, work around allow_text_output instead

Author: DouweM

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -245,13 +245,12 @@ async def add_mcp_server_tools(server: MCPServer) -> None:
     output_mode = None
     output_object = None
     output_tools = []
-    allow_text_output = _output.allow_text_output(output_schema)
+    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()
-        if output_mode != 'tool':
-            allow_text_output = False
+        require_tool_use = output_mode == 'tool' and not output_schema.allow_plain_text_output
 
         supported_modes = model.supported_output_modes
         if output_mode not in supported_modes:
@@ -262,7 +261,7 @@ async def add_mcp_server_tools(server: MCPServer) -> None:
         output_mode=output_mode,
         output_object=output_object,
         output_tools=output_tools,
-        allow_text_output=allow_text_output,
+        require_tool_use=require_tool_use,
     )
 
 
@@ -413,24 +412,20 @@ async def stream(
         async for _event in stream:
             pass
 
-    async def _run_stream(  # noqa: C901
+    async def _run_stream(
         self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]
     ) -> AsyncIterator[_messages.HandleResponseEvent]:
         if self._events_iterator is None:
             # Ensure that the stream is only run once
 
             async def _run_stream() -> AsyncIterator[_messages.HandleResponseEvent]:
                 texts: list[str] = []
-                outputs: list[str] = []
                 tool_calls: list[_messages.ToolCallPart] = []
                 for part in self.model_response.parts:
                     if isinstance(part, _messages.TextPart):
                         # ignore empty content for text parts, see #437
                         if part.content:
                             texts.append(part.content)
-                    elif isinstance(part, _messages.OutputPart):
-                        if part.content:
-                            outputs.append(part.content)
                     elif isinstance(part, _messages.ToolCallPart):
                         tool_calls.append(part)
                     else:
@@ -443,9 +438,6 @@ async def _run_stream() -> AsyncIterator[_messages.HandleResponseEvent]:
                 if tool_calls:
                     async for event in self._handle_tool_calls(ctx, tool_calls):
                         yield event
-                elif outputs:  # TODO: Can we have tool calls and structured output? Should we handle both?
-                    # No events are emitted during the handling of structured outputs, so we don't need to yield anything
-                    self._next_node = await self._handle_outputs(ctx, outputs)
                 elif texts:
                     # No events are emitted during the handling of text responses, so we don't need to yield anything
                     self._next_node = await self._handle_text_response(ctx, texts)
@@ -537,42 +529,18 @@ async def _handle_text_response(
         output_schema = ctx.deps.output_schema
 
         text = '\n\n'.join(texts)
-        if _output.allow_text_output(output_schema):
-            # The following cast is safe because we know `str` is an allowed result type
-            result_data_input = cast(NodeRunEndT, text)
-            try:
-                result_data = await _validate_output(result_data_input, ctx, None)
-            except _output.ToolRetryError as e:
-                ctx.state.increment_retries(ctx.deps.max_result_retries)
-                return ModelRequestNode[DepsT, NodeRunEndT](_messages.ModelRequest(parts=[e.tool_retry]))
+        try:
+            if output_schema is None or output_schema.allow_plain_text_output:
+                # The following cast is safe because we know `str` is an allowed result type
+                result_data = cast(NodeRunEndT, text)
+            elif output_schema.allow_json_text_output:
+                result_data = output_schema.validate(text)
             else:
-                return self._handle_final_result(ctx, result.FinalResult(result_data, None, None), [])
-        else:
-            ctx.state.increment_retries(ctx.deps.max_result_retries)
-            return ModelRequestNode[DepsT, NodeRunEndT](
-                _messages.ModelRequest(
-                    parts=[
-                        _messages.RetryPromptPart(
-                            content='Plain text responses are not permitted, please include your response in a tool call',
-                        )
-                    ]
+                m = _messages.RetryPromptPart(
+                    content='Plain text responses are not permitted, please include your response in a tool call',
                 )
-            )
+                raise _output.ToolRetryError(m)
 
-    async def _handle_outputs(
-        self,
-        ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
-        outputs: list[str],
-    ) -> ModelRequestNode[DepsT, NodeRunEndT] | End[result.FinalResult[NodeRunEndT]]:
-        if len(outputs) != 1:
-            raise exceptions.UnexpectedModelBehavior('Received multiple structured outputs in a single response')
-        output_schema = ctx.deps.output_schema
-        if not output_schema:
-            raise exceptions.UnexpectedModelBehavior('Must specify a non-str result_type when using structured outputs')
-
-        structured_output = outputs[0]
-        try:
-            result_data = output_schema.validate(structured_output)
             result_data = await _validate_output(result_data, ctx, None)
         except _output.ToolRetryError as e:
             ctx.state.increment_retries(ctx.deps.max_result_retries)

pydantic_ai_slim/pydantic_ai/_output.py

@@ -220,7 +220,8 @@ class OutputSchema(Generic[OutputDataT]):
     forced_mode: OutputMode | None
     object_schema: OutputObjectSchema[OutputDataT]
     tools: dict[str, OutputTool[OutputDataT]]
-    allow_text_output: bool  # TODO: Verify structured output works correctly with string as a union member
+    allow_plain_text_output: bool
+    allow_json_text_output: bool  # TODO: Turn into allowed_text_output: Literal['plain', 'json'] | None
 
     @classmethod
     def build(
@@ -235,11 +236,14 @@ def build(
             return None
 
         forced_mode = None
+        allow_json_text_output = True
+        allow_plain_text_output = False
         tool_output_type = None
-        allow_text_output = False
         if isinstance(output_type, ToolOutput):
-            # do we need to error on conflicts here? (DavidM): If this is internal maybe doesn't matter, if public, use overloads
             forced_mode = 'tool'
+            allow_json_text_output = False
+
+            # do we need to error on conflicts here? (DavidM): If this is internal maybe doesn't matter, if public, use overloads
             name = output_type.name
             description = output_type.description
             output_type_ = output_type.output_type
@@ -255,12 +259,15 @@ def build(
             name = output_type.name
             description = output_type.description
             output_type_ = output_type.output_type
-        else:
+        elif output_type_other_than_str := extract_str_from_union(output_type):
+            forced_mode = 'tool'
             output_type_ = output_type
 
-            if output_type_other_than_str := extract_str_from_union(output_type):
-                allow_text_output = True
-                tool_output_type = output_type_other_than_str.value
+            allow_json_text_output = False
+            allow_plain_text_output = True
+            tool_output_type = output_type_other_than_str.value
+        else:
+            output_type_ = output_type
 
         output_object_schema = OutputObjectSchema(
             output_type=output_type_, name=name, description=description, strict=strict
@@ -292,7 +299,8 @@ def build(
             forced_mode=forced_mode,
             object_schema=output_object_schema,
             tools=tools,
-            allow_text_output=allow_text_output,
+            allow_plain_text_output=allow_plain_text_output,
+            allow_json_text_output=allow_json_text_output,
         )
 
     def find_named_tool(
@@ -341,8 +349,7 @@ def validate(
 
 
 def allow_text_output(output_schema: OutputSchema[Any] | None) -> bool:
-    """Check if the result schema allows text results."""
-    return output_schema is None or output_schema.allow_text_output
+    return output_schema is None or output_schema.allow_plain_text_output or output_schema.allow_json_text_output
 
 
 @dataclass

pydantic_ai_slim/pydantic_ai/agent.py

@@ -997,8 +997,6 @@ async def stream_to_final(
                                     elif isinstance(new_part, _messages.ToolCallPart) and output_schema:
                                         for call, _ in output_schema.find_tool([new_part]):
                                             return FinalResult(s, call.tool_name, call.tool_call_id)
-                                    elif isinstance(new_part, _messages.OutputPart) and output_schema:
-                                        return FinalResult(s, None, None)
                             return None
 
                         final_result_details = await stream_to_final(streamed_response)

pydantic_ai_slim/pydantic_ai/messages.py

@@ -494,21 +494,6 @@ def has_content(self) -> bool:
         return bool(self.content)
 
 
-@dataclass
-class OutputPart:
-    """An output response from a model."""
-
-    content: str
-    """The output content of the response as a JSON-serialized string."""
-
-    part_kind: Literal['output'] = 'output'
-    """Part type identifier, this is available on all parts as a discriminator."""
-
-    def has_content(self) -> bool:
-        """Return `True` if the output content is non-empty."""
-        return bool(self.content)
-
-
 @dataclass
 class ToolCallPart:
     """A tool call from a model."""
@@ -563,7 +548,7 @@ def has_content(self) -> bool:
             return bool(self.args)
 
 
-ModelResponsePart = Annotated[Union[TextPart, OutputPart, ToolCallPart], pydantic.Discriminator('part_kind')]
+ModelResponsePart = Annotated[Union[TextPart, ToolCallPart], pydantic.Discriminator('part_kind')]
 """A message part returned by a model."""
 
 
@@ -659,33 +644,6 @@ def apply(self, part: ModelResponsePart) -> TextPart:
         return replace(part, content=part.content + self.content_delta)
 
 
-@dataclass
-class OutputPartDelta:
-    """A partial update (delta) for a `OutputPart` to append new structured output content."""
-
-    content_delta: str
-    """The incremental structured output content to add to the existing `OutputPart` content."""
-
-    part_delta_kind: Literal['output'] = 'output'
-    """Part delta type identifier, used as a discriminator."""
-
-    def apply(self, part: ModelResponsePart) -> OutputPart:
-        """Apply this structured output delta to an existing `OutputPart`.
-
-        Args:
-            part: The existing model response part, which must be a `OutputPart`.
-
-        Returns:
-            A new `OutputPart` with updated structured output content.
-
-        Raises:
-            ValueError: If `part` is not a `OutputPart`.
-        """
-        if not isinstance(part, OutputPart):
-            raise ValueError('Cannot apply OutputPartDeltas to non-OutputParts')
-        return replace(part, content=part.content + self.content_delta)
-
-
 @dataclass
 class ToolCallPartDelta:
     """A partial update (delta) for a `ToolCallPart` to modify tool name, arguments, or tool call ID."""
@@ -801,9 +759,7 @@ def _apply_to_part(self, part: ToolCallPart) -> ToolCallPart:
         return part
 
 
-ModelResponsePartDelta = Annotated[
-    Union[TextPartDelta, OutputPartDelta, ToolCallPartDelta], pydantic.Discriminator('part_delta_kind')
-]
+ModelResponsePartDelta = Annotated[Union[TextPartDelta, ToolCallPartDelta], pydantic.Discriminator('part_delta_kind')]
 """A partial update (delta) for any model response part."""
 
 

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -266,7 +266,7 @@ class ModelRequestParameters:
     output_mode: OutputMode | None
     output_object: OutputObjectDefinition | None
     output_tools: list[ToolDefinition]
-    allow_text_output: bool
+    require_tool_use: bool
 
 
 class Model(ABC):

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -21,7 +21,6 @@
     ModelResponse,
     ModelResponsePart,
     ModelResponseStreamEvent,
-    OutputPart,
     RetryPromptPart,
     SystemPromptPart,
     TextPart,
@@ -213,7 +212,7 @@ async def _messages_create(
         if not tools:
             tool_choice = None
         else:
-            if not model_request_parameters.allow_text_output:
+            if model_request_parameters.require_tool_use:
                 tool_choice = {'type': 'any'}
             else:
                 tool_choice = {'type': 'auto'}
@@ -321,7 +320,7 @@ async def _map_message(self, messages: list[ModelMessage]) -> tuple[str, list[Me
             elif isinstance(m, ModelResponse):
                 assistant_content_params: list[TextBlockParam | ToolUseBlockParam] = []
                 for response_part in m.parts:
-                    if isinstance(response_part, (TextPart, OutputPart)):
+                    if isinstance(response_part, TextPart):
                         assistant_content_params.append(TextBlockParam(text=response_part.content, type='text'))
                     else:
                         tool_use_block_param = ToolUseBlockParam(

pydantic_ai_slim/pydantic_ai/models/bedrock.py

@@ -304,7 +304,7 @@ async def _messages_create(
         support_tools_choice = self.model_name.startswith(('anthropic', 'us.anthropic'))
         if not tools or not support_tools_choice:
             tool_choice: ToolChoiceTypeDef = {}
-        elif not model_request_parameters.allow_text_output:
+        elif model_request_parameters.require_tool_use:
             tool_choice = {'any': {}}
         else:
             tool_choice = {'auto': {}}

pydantic_ai_slim/pydantic_ai/models/cohere.py

@@ -13,7 +13,6 @@
     ModelRequest,
     ModelResponse,
     ModelResponsePart,
-    OutputPart,
     RetryPromptPart,
     SystemPromptPart,
     TextPart,
@@ -206,7 +205,7 @@ def _map_messages(self, messages: list[ModelMessage]) -> list[ChatMessageV2]:
                 texts: list[str] = []
                 tool_calls: list[ToolCallV2] = []
                 for item in message.parts:
-                    if isinstance(item, (TextPart, OutputPart)):
+                    if isinstance(item, TextPart):
                         texts.append(item.content)
                     elif isinstance(item, ToolCallPart):
                         tool_calls.append(self._map_tool_call(item))

pydantic_ai_slim/pydantic_ai/models/function.py

@@ -21,7 +21,6 @@
     ModelRequest,
     ModelResponse,
     ModelResponseStreamEvent,
-    OutputPart,
     RetryPromptPart,
     SystemPromptPart,
     TextPart,
@@ -92,7 +91,7 @@ async def request(
     ) -> ModelResponse:
         agent_info = AgentInfo(
             model_request_parameters.function_tools,
-            model_request_parameters.allow_text_output,
+            not model_request_parameters.require_tool_use,
             model_request_parameters.output_tools,
             model_settings,
         )
@@ -121,7 +120,7 @@ async def request_stream(
     ) -> AsyncIterator[StreamedResponse]:
         agent_info = AgentInfo(
             model_request_parameters.function_tools,
-            model_request_parameters.allow_text_output,
+            not model_request_parameters.require_tool_use,
             model_request_parameters.output_tools,
             model_settings,
         )
@@ -267,7 +266,7 @@ def _estimate_usage(messages: Iterable[ModelMessage]) -> usage.Usage:
                     assert_never(part)
         elif isinstance(message, ModelResponse):
             for part in message.parts:
-                if isinstance(part, (TextPart, OutputPart)):
+                if isinstance(part, TextPart):
                     response_tokens += _estimate_string_tokens(part.content)
                 elif isinstance(part, ToolCallPart):
                     call = part

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -28,7 +28,6 @@
     ModelResponse,
     ModelResponsePart,
     ModelResponseStreamEvent,
-    OutputPart,
     RetryPromptPart,
     SystemPromptPart,
     TextPart,
@@ -183,7 +182,7 @@ def _customize_output_object_def(o: OutputObjectDefinition):
             if model_request_parameters.output_object
             else None,
             output_tools=[_customize_tool_def(tool) for tool in model_request_parameters.output_tools],
-            allow_text_output=model_request_parameters.allow_text_output,
+            require_tool_use=model_request_parameters.require_tool_use,
         )
 
     @property
@@ -205,7 +204,7 @@ def _get_tools(self, model_request_parameters: ModelRequestParameters) -> _Gemin
     def _get_tool_config(
         self, model_request_parameters: ModelRequestParameters, tools: _GeminiTools | None
     ) -> _GeminiToolConfig | None:
-        if model_request_parameters.allow_text_output:
+        if not model_request_parameters.require_tool_use:
             return None
         elif tools:
             return _tool_config([t['name'] for t in tools['function_declarations']])
@@ -559,7 +558,7 @@ def _content_model_response(m: ModelResponse) -> _GeminiContent:
     for item in m.parts:
         if isinstance(item, ToolCallPart):
             parts.append(_function_call_part_from_call(item))
-        elif isinstance(item, (TextPart, OutputPart)):
+        elif isinstance(item, TextPart):
             if item.content:
                 parts.append(_GeminiTextPart(text=item.content))
         else: