Pass just one toolset into the run

Author: DouweM

docs/agents.md

@@ -808,7 +808,7 @@ with capture_run_messages() as messages:  # (2)!
         result = agent.run_sync('Please get me the volume of a box with size 6.')
     except UnexpectedModelBehavior as e:
         print('An error occurred:', e)
-        #> An error occurred: Tool exceeded max retries count of 1
+        #> An error occurred: Tool 'calc_volume' exceeded max retries count of 1
         print('cause:', repr(e.__cause__))
         #> cause: ModelRetry('Please try again.')
         print('messages:', messages)

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -17,7 +17,7 @@
 
 from pydantic_ai._function_schema import _takes_ctx as is_takes_ctx  # type: ignore
 from pydantic_ai._utils import is_async_callable, run_in_executor
-from pydantic_ai.toolset import AbstractToolset, CombinedToolset, RunToolset
+from pydantic_ai.toolset import AbstractToolset, RunToolset
 from pydantic_graph import BaseNode, Graph, GraphRunContext
 from pydantic_graph.nodes import End, NodeRunEndT
 
@@ -107,7 +107,6 @@ class GraphAgentDeps(Generic[DepsT, OutputDataT]):
     get_instructions: Callable[[RunContext[DepsT]], Awaitable[str | None]]
 
     output_schema: _output.OutputSchema[OutputDataT]
-    output_toolset: RunToolset[DepsT]
     output_validators: list[_output.OutputValidator[DepsT, OutputDataT]]
 
     history_processors: Sequence[HistoryProcessor[DepsT]]
@@ -249,11 +248,7 @@ async def _prepare_request_parameters(
 ) -> models.ModelRequestParameters:
     """Build tools and create an agent model."""
     run_context = build_run_context(ctx)
-    ctx.deps.toolset = toolset = await ctx.deps.toolset.prepare_for_run(run_context)
-    ctx.deps.output_toolset = output_toolset = await ctx.deps.output_toolset.prepare_for_run(run_context)
-
-    # This will raise errors for any name conflicts
-    CombinedToolset[DepsT]([output_toolset, toolset])
+    ctx.deps.toolset = await ctx.deps.toolset.prepare_for_run(run_context)
 
     output_schema = ctx.deps.output_schema
     output_object = None
@@ -263,10 +258,18 @@ async def _prepare_request_parameters(
     # ToolOrTextOutputSchema, NativeOutputSchema, and PromptedOutputSchema all inherit from TextOutputSchema
     allow_text_output = isinstance(output_schema, _output.TextOutputSchema)
 
+    function_tools: list[ToolDefinition] = []
+    output_tools: list[ToolDefinition] = []
+    for tool_def in ctx.deps.toolset.tool_defs:
+        if tool_def.kind == 'output':
+            output_tools.append(tool_def)
+        else:
+            function_tools.append(tool_def)
+
     return models.ModelRequestParameters(
-        function_tools=toolset.tool_defs,
+        function_tools=function_tools,
         output_mode=output_schema.mode,
-        output_tools=output_toolset.tool_defs,
+        output_tools=output_tools,
         output_object=output_object,
         allow_text_output=allow_text_output,
     )
@@ -487,7 +490,7 @@ async def _handle_tool_calls(  # noqa: C901
         final_result: result.FinalResult[NodeRunEndT] | None = None
         parts: list[_messages.ModelRequestPart] = []
 
-        toolset = CombinedToolset([ctx.deps.toolset, ctx.deps.output_toolset])
+        toolset = ctx.deps.toolset
 
         unknown_calls: list[_messages.ToolCallPart] = []
         tool_calls_by_kind: dict[ToolKind, list[_messages.ToolCallPart]] = defaultdict(list)

pydantic_ai_slim/pydantic_ai/agent.py

@@ -686,11 +686,8 @@ async def main():
             run_step=state.run_step,
         )
 
-        run_toolset = await self._toolset.prepare_for_run(run_context)
-        run_output_toolset = await output_toolset.prepare_for_run(run_context)
-
-        # This will raise errors for any name conflicts
-        CombinedToolset([run_output_toolset, run_toolset])
+        toolset = CombinedToolset([output_toolset, self._toolset])
+        run_toolset = await toolset.prepare_for_run(run_context)
 
         model_settings = merge_model_settings(self.model_settings, model_settings)
         usage_limits = usage_limits or _usage.UsageLimits()
@@ -738,7 +735,6 @@ async def get_instructions(run_context: RunContext[AgentDepsT]) -> str | None:
             end_strategy=self.end_strategy,
             output_schema=output_schema,
             output_validators=output_validators,
-            output_toolset=run_output_toolset,
             history_processors=self.history_processors,
             toolset=run_toolset,
             tracer=tracer,

pydantic_ai_slim/pydantic_ai/toolset.py

@@ -368,6 +368,11 @@ def tool_names(self) -> list[str]:
     def get_tool_args_validator(self, ctx: RunContext[AgentDepsT], name: str) -> SchemaValidator:
         return self._toolset_for_tool_name(name).get_tool_args_validator(ctx, name)
 
+    def validate_tool_args(
+        self, ctx: RunContext[AgentDepsT], name: str, args: str | dict[str, Any] | None, allow_partial: bool = False
+    ) -> dict[str, Any]:
+        return self._toolset_for_tool_name(name).validate_tool_args(ctx, name, args, allow_partial)
+
     def max_retries_for_tool(self, name: str) -> int:
         return self._toolset_for_tool_name(name).max_retries_for_tool(name)
 
@@ -678,9 +683,9 @@ def _on_error(self, name: str, e: Exception) -> Never:
         max_retries = self.max_retries_for_tool(name)
         current_retry = self._retries.get(name, 0)
         if current_retry == max_retries:
-            raise UnexpectedModelBehavior(f'Tool exceeded max retries count of {max_retries}') from e
+            raise UnexpectedModelBehavior(f'Tool {name!r} exceeded max retries count of {max_retries}') from e
         else:
-            self._retries[name] = current_retry + 1
+            self._retries[name] = current_retry + 1  # TODO: Reset on successful call!
             raise e
 
     def _validate_tool_name(self, name: str) -> None:

tests/models/test_model_test.py

@@ -157,7 +157,7 @@ def validate_output(ctx: RunContext[None], output: OutputModel) -> OutputModel:
         call_count += 1
         raise ModelRetry('Fail')
 
-    with pytest.raises(UnexpectedModelBehavior, match='Exceeded maximum retries'):
+    with pytest.raises(UnexpectedModelBehavior, match="Tool 'final_result' exceeded max retries count of 2"):
         agent.run_sync('Hello', model=TestModel())
 
     assert call_count == 3
@@ -200,7 +200,7 @@ class ResultModel(BaseModel):
 
     agent = Agent('test', output_type=ResultModel, retries=2)
 
-    with pytest.raises(UnexpectedModelBehavior, match='Exceeded maximum retries'):
+    with pytest.raises(UnexpectedModelBehavior, match="Tool 'final_result' exceeded max retries count of 2"):
         agent.run_sync('Hello', model=TestModel(custom_output_args={'foo': 'a', 'bar': 1}))
 
 

tests/test_agent.py

@@ -2293,7 +2293,7 @@ def another_tool(y: int) -> int:
                             tool_name='another_tool', content=2, tool_call_id=IsStr(), timestamp=IsNow(tz=timezone.utc)
                         ),
                         RetryPromptPart(
-                            content="Unknown tool name: 'unknown_tool'. Available tools: regular_tool, another_tool, final_result",
+                            content="Unknown tool name: 'unknown_tool'. Available tools: final_result, regular_tool, another_tool",
                             tool_name='unknown_tool',
                             tool_call_id=IsStr(),
                             timestamp=IsDatetime(),
@@ -2380,7 +2380,7 @@ def another_tool(y: int) -> int:  # pragma: no cover
                         ),
                         RetryPromptPart(
                             tool_name='unknown_tool',
-                            content="Unknown tool name: 'unknown_tool'. Available tools: regular_tool, another_tool, final_result",
+                            content="Unknown tool name: 'unknown_tool'. Available tools: final_result, regular_tool, another_tool",
                             timestamp=IsNow(tz=timezone.utc),
                             tool_call_id=IsStr(),
                         ),

tests/test_tools.py

@@ -1171,7 +1171,7 @@ def infinite_retry_tool(ctx: RunContext[None]) -> int:
         call_retries.append(ctx.retry)
         raise ModelRetry('Please try again.')
 
-    with pytest.raises(UnexpectedModelBehavior, match='Tool exceeded max retries count of 5'):
+    with pytest.raises(UnexpectedModelBehavior, match="Tool 'infinite_retry_tool' exceeded max retries count of 5"):
         agent.run_sync('Begin infinite retry loop!')
 
     # There are extra 0s here because the toolset is prepared once ahead of the graph run, before the user prompt part is added in.