Actually use Agent.prepare_output_tools

Author: DouweM

pydantic_ai_slim/pydantic_ai/agent.py

@@ -733,12 +733,12 @@ async def main():
 
         output_toolset = self._output_toolset
         if output_schema != self._output_schema or output_validators:
-            output_toolset = output_schema.toolset
+            output_toolset = cast(OutputToolset[AgentDepsT], output_schema.toolset)
             if output_toolset:
                 output_toolset.max_retries = self._max_result_retries
                 output_toolset.output_validators = output_validators
-                if self._prepare_output_tools:
-                    output_toolset = PreparedToolset(output_toolset, self._prepare_output_tools)
+        if output_toolset and self._prepare_output_tools:
+            output_toolset = PreparedToolset(output_toolset, self._prepare_output_tools)
 
         # Build the graph
         graph: Graph[_agent_graph.GraphAgentState, _agent_graph.GraphAgentDeps[AgentDepsT, Any], FinalResult[Any]] = (

tests/test_streaming.py

@@ -1109,7 +1109,7 @@ async def prepare_tool(ctx: RunContext[None], tool_def: ToolDefinition) -> ToolD
 
     @agent.tool_plain(prepare=prepare_tool)
     def my_tool(x: int) -> int:
-        return x + 1
+        return x + 1  # pragma: no cover
 
     async with agent.run_stream('Hello') as result:
         assert not result.is_complete
@@ -1143,7 +1143,7 @@ async def prepare_tool(ctx: RunContext[None], tool_def: ToolDefinition) -> ToolD
 
     @agent.tool_plain(prepare=prepare_tool)
     def my_tool(x: int) -> int:
-        return x + 1
+        return x + 1  # pragma: no cover
 
     outputs: list[str | DeferredToolCalls] = []
     events: list[Any] = []