From 017eac174c58a0cce1521b13ac2dd7d2018f0e2c Mon Sep 17 00:00:00 2001 From: adtyavrdhn Date: Wed, 10 Dec 2025 08:38:59 +0530 Subject: [PATCH 1/9] init --- pydantic_ai_slim/pydantic_ai/_agent_graph.py | 35 +++++++++--- pydantic_ai_slim/pydantic_ai/_run_context.py | 2 + pydantic_ai_slim/pydantic_ai/_tool_manager.py | 20 +++++++ .../pydantic_ai/agent/__init__.py | 6 ++ pydantic_ai_slim/pydantic_ai/tools.py | 3 + .../pydantic_ai/toolsets/abstract.py | 2 + .../pydantic_ai/toolsets/combined.py | 1 + .../pydantic_ai/toolsets/function.py | 15 +++++ tests/test_tools.py | 56 +++++++++++++++++++ 9 files changed, 133 insertions(+), 7 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_agent_graph.py b/pydantic_ai_slim/pydantic_ai/_agent_graph.py index 043c27c4f5..4400443920 100644 --- a/pydantic_ai_slim/pydantic_ai/_agent_graph.py +++ b/pydantic_ai_slim/pydantic_ai/_agent_graph.py @@ -95,6 +95,7 @@ class GraphAgentState: retries: int = 0 run_step: int = 0 run_id: str = dataclasses.field(default_factory=lambda: str(uuid.uuid4())) + tool_usage: dict[str, int] = dataclasses.field(default_factory=dict) def increment_retries( self, @@ -821,6 +822,7 @@ def build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT else DEFAULT_INSTRUMENTATION_VERSION, run_step=ctx.state.run_step, run_id=ctx.state.run_id, + tool_usage=ctx.state.tool_usage, ) validation_context = build_validation_context(ctx.deps.validation_context, run_context) run_context = replace(run_context, validation_context=validation_context) @@ -1020,14 +1022,33 @@ async def _call_tools( projected_usage.tool_calls += len(tool_calls) usage_limits.check_before_tool_call(projected_usage) + calls_to_run: list[_messages.ToolCallPart] = [] + + # For each tool, check how many calls are going to be made and if it is over the limit + tool_call_counts: defaultdict[str, int] = defaultdict(int) + for call in tool_calls: + tool_call_counts[call.tool_name] += 1 + for call in tool_calls: - yield _messages.FunctionToolCallEvent(call) + current_tool_use = tool_manager.get_current_use_of_tool(call.tool_name) + max_tool_use = tool_manager.get_max_use_of_tool(call.tool_name) + if max_tool_use is not None and current_tool_use + tool_call_counts[call.tool_name] > max_tool_use: + return_part = _messages.ToolReturnPart( + tool_name=call.tool_name, + content=f'Tool call limit reached for tool "{call.tool_name}".', + tool_call_id=call.tool_call_id, + ) + output_parts.append(return_part) + yield _messages.FunctionToolResultEvent(return_part) + else: + yield _messages.FunctionToolCallEvent(call) + calls_to_run.append(call) with tracer.start_as_current_span( 'running tools', attributes={ - 'tools': [call.tool_name for call in tool_calls], - 'logfire.msg': f'running {len(tool_calls)} tool{"" if len(tool_calls) == 1 else "s"}', + 'tools': [call.tool_name for call in calls_to_run], + 'logfire.msg': f'running {len(calls_to_run)} tool{"" if len(calls_to_run) == 1 else "s"}', }, ): @@ -1061,8 +1082,8 @@ async def handle_call_or_result( return _messages.FunctionToolResultEvent(tool_part, content=tool_user_content) - if tool_manager.should_call_sequentially(tool_calls): - for index, call in enumerate(tool_calls): + if tool_manager.should_call_sequentially(calls_to_run): + for index, call in enumerate(calls_to_run): if event := await handle_call_or_result( _call_tool(tool_manager, call, tool_call_results.get(call.tool_call_id)), index, @@ -1075,7 +1096,7 @@ async def handle_call_or_result( _call_tool(tool_manager, call, tool_call_results.get(call.tool_call_id)), name=call.tool_name, ) - for call in tool_calls + for call in calls_to_run ] pending = tasks @@ -1092,7 +1113,7 @@ async def handle_call_or_result( output_parts.extend([user_parts_by_index[k] for k in sorted(user_parts_by_index)]) _populate_deferred_calls( - tool_calls, deferred_calls_by_index, deferred_metadata_by_index, output_deferred_calls, output_deferred_metadata + calls_to_run, deferred_calls_by_index, deferred_metadata_by_index, output_deferred_calls, output_deferred_metadata ) diff --git a/pydantic_ai_slim/pydantic_ai/_run_context.py b/pydantic_ai_slim/pydantic_ai/_run_context.py index b605bd8b54..d2013d055e 100644 --- a/pydantic_ai_slim/pydantic_ai/_run_context.py +++ b/pydantic_ai_slim/pydantic_ai/_run_context.py @@ -48,6 +48,8 @@ class RunContext(Generic[RunContextAgentDepsT]): """Instrumentation settings version, if instrumentation is enabled.""" retries: dict[str, int] = field(default_factory=dict) """Number of retries for each tool so far.""" + tool_usage: dict[str, int] = field(default_factory=dict) + """Number of calls for each tool so far.""" tool_call_id: str | None = None """The ID of the tool call.""" tool_name: str | None = None diff --git a/pydantic_ai_slim/pydantic_ai/_tool_manager.py b/pydantic_ai_slim/pydantic_ai/_tool_manager.py index 9a9f93e1ff..8e6f5ad21d 100644 --- a/pydantic_ai_slim/pydantic_ai/_tool_manager.py +++ b/pydantic_ai_slim/pydantic_ai/_tool_manager.py @@ -161,6 +161,8 @@ async def _call_tool( partial_output=allow_partial, ) + self.ctx.tool_usage[name] = self.ctx.tool_usage.get(name, 0) + 1 + pyd_allow_partial = 'trailing-strings' if allow_partial else 'off' validator = tool.args_validator if isinstance(call.args, str): @@ -274,3 +276,21 @@ async def _call_function_tool( ) return tool_result + + def get_max_use_of_tool(self, tool_name: str) -> int | None: + """Get the maximum number of uses allowed for a given tool, or `None` if unlimited.""" + if self.tools is None: + raise ValueError('ToolManager has not been prepared for a run step yet') # pragma: no cover + + tool = self.tools.get(tool_name, None) + if tool is None: + return None + + return tool.max_uses + + def get_current_use_of_tool(self, tool_name: str) -> int: + """Get the current number of uses of a given tool.""" + if self.ctx is None: + raise ValueError('ToolManager has not been prepared for a run step yet') # pragma: no cover + + return self.ctx.tool_usage.get(tool_name, 0) \ No newline at end of file diff --git a/pydantic_ai_slim/pydantic_ai/agent/__init__.py b/pydantic_ai_slim/pydantic_ai/agent/__init__.py index 19edb4a619..aed4a0811a 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/agent/__init__.py @@ -1032,6 +1032,7 @@ def tool( sequential: bool = False, requires_approval: bool = False, metadata: dict[str, Any] | None = None, + max_uses: int | None = None, ) -> Callable[[ToolFuncContext[AgentDepsT, ToolParams]], ToolFuncContext[AgentDepsT, ToolParams]]: ... def tool( @@ -1050,6 +1051,7 @@ def tool( sequential: bool = False, requires_approval: bool = False, metadata: dict[str, Any] | None = None, + max_uses: int | None = None, ) -> Any: """Decorator to register a tool function which takes [`RunContext`][pydantic_ai.tools.RunContext] as its first argument. @@ -1119,6 +1121,7 @@ def tool_decorator( sequential=sequential, requires_approval=requires_approval, metadata=metadata, + max_uses=max_uses, ) return func_ @@ -1143,6 +1146,7 @@ def tool_plain( sequential: bool = False, requires_approval: bool = False, metadata: dict[str, Any] | None = None, + max_uses: int | None = None, ) -> Callable[[ToolFuncPlain[ToolParams]], ToolFuncPlain[ToolParams]]: ... def tool_plain( @@ -1161,6 +1165,7 @@ def tool_plain( sequential: bool = False, requires_approval: bool = False, metadata: dict[str, Any] | None = None, + max_uses: int | None = None, ) -> Any: """Decorator to register a tool function which DOES NOT take `RunContext` as an argument. @@ -1228,6 +1233,7 @@ def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams sequential=sequential, requires_approval=requires_approval, metadata=metadata, + max_uses=max_uses, ) return func_ diff --git a/pydantic_ai_slim/pydantic_ai/tools.py b/pydantic_ai_slim/pydantic_ai/tools.py index dcd860b019..a801fc886d 100644 --- a/pydantic_ai_slim/pydantic_ai/tools.py +++ b/pydantic_ai_slim/pydantic_ai/tools.py @@ -264,6 +264,7 @@ class Tool(Generic[ToolAgentDepsT]): function: ToolFuncEither[ToolAgentDepsT] takes_ctx: bool max_retries: int | None + max_uses: int | None name: str description: str | None prepare: ToolPrepareFunc[ToolAgentDepsT] | None @@ -286,6 +287,7 @@ def __init__( *, takes_ctx: bool | None = None, max_retries: int | None = None, + max_uses: int | None = None, name: str | None = None, description: str | None = None, prepare: ToolPrepareFunc[ToolAgentDepsT] | None = None, @@ -364,6 +366,7 @@ async def prep_my_tool( ) self.takes_ctx = self.function_schema.takes_ctx self.max_retries = max_retries + self.max_uses = max_uses self.name = name or function.__name__ self.description = description or self.function_schema.description self.prepare = prepare diff --git a/pydantic_ai_slim/pydantic_ai/toolsets/abstract.py b/pydantic_ai_slim/pydantic_ai/toolsets/abstract.py index 98d9cd224f..a75bcdbd30 100644 --- a/pydantic_ai_slim/pydantic_ai/toolsets/abstract.py +++ b/pydantic_ai_slim/pydantic_ai/toolsets/abstract.py @@ -52,6 +52,8 @@ class ToolsetTool(Generic[AgentDepsT]): """The tool definition for this tool, including the name, description, and parameters.""" max_retries: int """The maximum number of retries to attempt if the tool call fails.""" + max_uses: int + """The maximum number of uses allowed for this tool.""" args_validator: SchemaValidator | SchemaValidatorProt """The Pydantic Core validator for the tool's arguments. diff --git a/pydantic_ai_slim/pydantic_ai/toolsets/combined.py b/pydantic_ai_slim/pydantic_ai/toolsets/combined.py index e095e4aa1f..0c4e9bd98b 100644 --- a/pydantic_ai_slim/pydantic_ai/toolsets/combined.py +++ b/pydantic_ai_slim/pydantic_ai/toolsets/combined.py @@ -77,6 +77,7 @@ async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[ toolset=tool_toolset, tool_def=tool.tool_def, max_retries=tool.max_retries, + max_uses=tool.max_uses, args_validator=tool.args_validator, source_toolset=toolset, source_tool=tool, diff --git a/pydantic_ai_slim/pydantic_ai/toolsets/function.py b/pydantic_ai_slim/pydantic_ai/toolsets/function.py index e185ed0273..2c4f6aaebb 100644 --- a/pydantic_ai_slim/pydantic_ai/toolsets/function.py +++ b/pydantic_ai_slim/pydantic_ai/toolsets/function.py @@ -35,6 +35,7 @@ class FunctionToolset(AbstractToolset[AgentDepsT]): tools: dict[str, Tool[Any]] max_retries: int + max_uses: int | None _id: str | None docstring_format: DocstringFormat require_parameter_descriptions: bool @@ -45,6 +46,7 @@ def __init__( tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = [], *, max_retries: int = 1, + max_uses: int | None = None, docstring_format: DocstringFormat = 'auto', require_parameter_descriptions: bool = False, schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema, @@ -60,6 +62,8 @@ def __init__( tools: The tools to add to the toolset. max_retries: The maximum number of retries for each tool during a run. Applies to all tools, unless overridden when adding a tool. + max_uses: The maximum number of uses allowed for each tool during a run. + Applies to all tools, unless overridden when adding a tool. docstring_format: Format of tool docstring, see [`DocstringFormat`][pydantic_ai.tools.DocstringFormat]. Defaults to `'auto'`, such that the format is inferred from the structure of the docstring. Applies to all tools, unless overridden when adding a tool. @@ -81,6 +85,7 @@ def __init__( """ self.max_retries = max_retries self._id = id + self.max_uses = max_uses self.docstring_format = docstring_format self.require_parameter_descriptions = require_parameter_descriptions self.schema_generator = schema_generator @@ -111,6 +116,7 @@ def tool( name: str | None = None, description: str | None = None, retries: int | None = None, + max_uses: int | None = None, prepare: ToolPrepareFunc[AgentDepsT] | None = None, docstring_format: DocstringFormat | None = None, require_parameter_descriptions: bool | None = None, @@ -129,6 +135,7 @@ def tool( name: str | None = None, description: str | None = None, retries: int | None = None, + max_uses: int | None = None, prepare: ToolPrepareFunc[AgentDepsT] | None = None, docstring_format: DocstringFormat | None = None, require_parameter_descriptions: bool | None = None, @@ -205,6 +212,7 @@ def tool_decorator( name=name, description=description, retries=retries, + max_uses=max_uses, prepare=prepare, docstring_format=docstring_format, require_parameter_descriptions=require_parameter_descriptions, @@ -225,6 +233,7 @@ def add_function( name: str | None = None, description: str | None = None, retries: int | None = None, + max_uses: int | None = None, prepare: ToolPrepareFunc[AgentDepsT] | None = None, docstring_format: DocstringFormat | None = None, require_parameter_descriptions: bool | None = None, @@ -248,6 +257,7 @@ def add_function( description: The description of the tool, defaults to the function docstring. retries: The number of retries to allow for this tool, defaults to the agent's default retries, which defaults to 1. + max_uses: The maximum number of uses allowed for this tool during a run. Defaults to None (unlimited). prepare: custom method to prepare the tool definition for each step, return `None` to omit this tool from a given step. This is useful if you want to customise a tool at call time, or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc]. @@ -287,6 +297,7 @@ def add_function( name=name, description=description, max_retries=retries, + max_uses=max_uses, prepare=prepare, docstring_format=docstring_format, require_parameter_descriptions=require_parameter_descriptions, @@ -308,6 +319,8 @@ def add_tool(self, tool: Tool[AgentDepsT]) -> None: raise UserError(f'Tool name conflicts with existing tool: {tool.name!r}') if tool.max_retries is None: tool.max_retries = self.max_retries + if tool.max_uses is None: + tool.max_uses = self.max_uses if self.metadata is not None: tool.metadata = self.metadata | (tool.metadata or {}) self.tools[tool.name] = tool @@ -316,6 +329,7 @@ async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[ tools: dict[str, ToolsetTool[AgentDepsT]] = {} for original_name, tool in self.tools.items(): max_retries = tool.max_retries if tool.max_retries is not None else self.max_retries + max_uses = tool.max_uses if tool.max_uses is not None else self.max_uses run_context = replace( ctx, tool_name=original_name, @@ -337,6 +351,7 @@ async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[ toolset=self, tool_def=tool_def, max_retries=max_retries, + max_uses=max_uses, args_validator=tool.function_schema.validator, call_func=tool.function_schema.call, is_async=tool.function_schema.is_async, diff --git a/tests/test_tools.py b/tests/test_tools.py index bcdf537994..c577b35022 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -1309,6 +1309,62 @@ def infinite_retry_tool(ctx: RunContext[None]) -> int: assert call_last_attempt == snapshot([False, False, False, False, False, True]) +def test_tool_max_uses(): + agent = Agent(TestModel(), output_type=[str, DeferredToolRequests]) + + @agent.tool(max_uses=1) + def tool_with_max_use(ctx: RunContext[None]) -> str: + return 'Used' + + # Force the agent to use this tool now + + result = agent.run_sync('Hello') + assert result.output == snapshot('{"tool_with_max_use":"Used"}') + messages = result.all_messages() + assert messages == snapshot( + [ + ModelRequest( + parts=[ + UserPromptPart( + content='Hello', + timestamp=IsDatetime(), + ) + ], + run_id=IsStr(), + ), + ModelResponse( + parts=[ + ToolCallPart( + tool_name='tool_with_max_use', args={}, tool_call_id='pyd_ai_tool_call_id__tool_with_max_use' + ) + ], + usage=RequestUsage(input_tokens=51, output_tokens=2), + model_name='test', + timestamp=IsDatetime(), + run_id=IsStr(), + ), + ModelRequest( + parts=[ + ToolReturnPart( + tool_name='tool_with_max_use', + content='Used', + tool_call_id='pyd_ai_tool_call_id__tool_with_max_use', + timestamp=IsDatetime(), + ) + ], + run_id=IsStr(), + ), + ModelResponse( + parts=[TextPart(content='{"tool_with_max_use":"Used"}')], + usage=RequestUsage(input_tokens=52, output_tokens=6), + model_name='test', + timestamp=IsDatetime(), + run_id=IsStr(), + ), + ] + ) + + def test_tool_raises_call_deferred(): agent = Agent(TestModel(), output_type=[str, DeferredToolRequests]) From 190e23026b03cdeee57fe1c5382afdb931f79014 Mon Sep 17 00:00:00 2001 From: adtyavrdhn Date: Wed, 10 Dec 2025 14:53:37 +0530 Subject: [PATCH 2/9] docstrings --- pydantic_ai_slim/pydantic_ai/agent/__init__.py | 2 ++ pydantic_ai_slim/pydantic_ai/tools.py | 1 + pydantic_ai_slim/pydantic_ai/toolsets/function.py | 1 + tests/test_tools.py | 2 -- 4 files changed, 4 insertions(+), 2 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/agent/__init__.py b/pydantic_ai_slim/pydantic_ai/agent/__init__.py index aed4a0811a..9a8d8ac68c 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/agent/__init__.py @@ -1101,6 +1101,7 @@ async def spam(ctx: RunContext[str], y: float) -> float: requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False. See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info. metadata: Optional metadata for the tool. This is not sent to the model but can be used for filtering and tool behavior customization. + max_uses: Optional maximum number of times this tool can be used during a run. Defaults to None (unlimited). """ def tool_decorator( @@ -1215,6 +1216,7 @@ async def spam(ctx: RunContext[str]) -> float: requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False. See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info. metadata: Optional metadata for the tool. This is not sent to the model but can be used for filtering and tool behavior customization. + max_uses: Optional maximum number of times this tool can be used during a run. Defaults to None (unlimited). """ def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams]: diff --git a/pydantic_ai_slim/pydantic_ai/tools.py b/pydantic_ai_slim/pydantic_ai/tools.py index a801fc886d..bfccf18a97 100644 --- a/pydantic_ai_slim/pydantic_ai/tools.py +++ b/pydantic_ai_slim/pydantic_ai/tools.py @@ -339,6 +339,7 @@ async def prep_my_tool( takes_ctx: Whether the function takes a [`RunContext`][pydantic_ai.tools.RunContext] first argument, this is inferred if unset. max_retries: Maximum number of retries allowed for this tool, set to the agent default if `None`. + max_uses: The maximum number of uses allowed for this tool during a run. Defaults to None (unlimited). name: Name of the tool, inferred from the function if `None`. description: Description of the tool, inferred from the function if `None`. prepare: custom method to prepare the tool definition for each step, return `None` to omit this diff --git a/pydantic_ai_slim/pydantic_ai/toolsets/function.py b/pydantic_ai_slim/pydantic_ai/toolsets/function.py index 2c4f6aaebb..3a766283f7 100644 --- a/pydantic_ai_slim/pydantic_ai/toolsets/function.py +++ b/pydantic_ai_slim/pydantic_ai/toolsets/function.py @@ -181,6 +181,7 @@ async def spam(ctx: RunContext[str], y: float) -> float: description: The description of the tool,defaults to the function docstring. retries: The number of retries to allow for this tool, defaults to the agent's default retries, which defaults to 1. + max_uses: The maximum number of uses allowed for this tool during a run. Defaults to None (unlimited). prepare: custom method to prepare the tool definition for each step, return `None` to omit this tool from a given step. This is useful if you want to customise a tool at call time, or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc]. diff --git a/tests/test_tools.py b/tests/test_tools.py index c577b35022..a96af78cc9 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -1316,8 +1316,6 @@ def test_tool_max_uses(): def tool_with_max_use(ctx: RunContext[None]) -> str: return 'Used' - # Force the agent to use this tool now - result = agent.run_sync('Hello') assert result.output == snapshot('{"tool_with_max_use":"Used"}') messages = result.all_messages() From c9d7aede070c8eb4646a614decbaf828dc3c3060 Mon Sep 17 00:00:00 2001 From: adtyavrdhn Date: Wed, 10 Dec 2025 15:00:38 +0530 Subject: [PATCH 3/9] docstrings + args --- pydantic_ai_slim/pydantic_ai/agent/__init__.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/agent/__init__.py b/pydantic_ai_slim/pydantic_ai/agent/__init__.py index 9a8d8ac68c..2102f6206c 100644 --- a/pydantic_ai_slim/pydantic_ai/agent/__init__.py +++ b/pydantic_ai_slim/pydantic_ai/agent/__init__.py @@ -1024,6 +1024,7 @@ def tool( name: str | None = None, description: str | None = None, retries: int | None = None, + max_uses: int | None = None, prepare: ToolPrepareFunc[AgentDepsT] | None = None, docstring_format: DocstringFormat = 'auto', require_parameter_descriptions: bool = False, @@ -1032,7 +1033,6 @@ def tool( sequential: bool = False, requires_approval: bool = False, metadata: dict[str, Any] | None = None, - max_uses: int | None = None, ) -> Callable[[ToolFuncContext[AgentDepsT, ToolParams]], ToolFuncContext[AgentDepsT, ToolParams]]: ... def tool( @@ -1043,6 +1043,7 @@ def tool( name: str | None = None, description: str | None = None, retries: int | None = None, + max_uses: int | None = None, prepare: ToolPrepareFunc[AgentDepsT] | None = None, docstring_format: DocstringFormat = 'auto', require_parameter_descriptions: bool = False, @@ -1051,7 +1052,6 @@ def tool( sequential: bool = False, requires_approval: bool = False, metadata: dict[str, Any] | None = None, - max_uses: int | None = None, ) -> Any: """Decorator to register a tool function which takes [`RunContext`][pydantic_ai.tools.RunContext] as its first argument. @@ -1088,6 +1088,7 @@ async def spam(ctx: RunContext[str], y: float) -> float: description: The description of the tool, defaults to the function docstring. retries: The number of retries to allow for this tool, defaults to the agent's default retries, which defaults to 1. + max_uses: The maximum number of uses allowed for this tool during a run. Defaults to None (unlimited). prepare: custom method to prepare the tool definition for each step, return `None` to omit this tool from a given step. This is useful if you want to customise a tool at call time, or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc]. @@ -1101,7 +1102,6 @@ async def spam(ctx: RunContext[str], y: float) -> float: requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False. See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info. metadata: Optional metadata for the tool. This is not sent to the model but can be used for filtering and tool behavior customization. - max_uses: Optional maximum number of times this tool can be used during a run. Defaults to None (unlimited). """ def tool_decorator( @@ -1114,6 +1114,7 @@ def tool_decorator( name=name, description=description, retries=retries, + max_uses=max_uses, prepare=prepare, docstring_format=docstring_format, require_parameter_descriptions=require_parameter_descriptions, @@ -1122,7 +1123,6 @@ def tool_decorator( sequential=sequential, requires_approval=requires_approval, metadata=metadata, - max_uses=max_uses, ) return func_ @@ -1139,6 +1139,7 @@ def tool_plain( name: str | None = None, description: str | None = None, retries: int | None = None, + max_uses: int | None = None, prepare: ToolPrepareFunc[AgentDepsT] | None = None, docstring_format: DocstringFormat = 'auto', require_parameter_descriptions: bool = False, @@ -1147,7 +1148,6 @@ def tool_plain( sequential: bool = False, requires_approval: bool = False, metadata: dict[str, Any] | None = None, - max_uses: int | None = None, ) -> Callable[[ToolFuncPlain[ToolParams]], ToolFuncPlain[ToolParams]]: ... def tool_plain( @@ -1158,6 +1158,7 @@ def tool_plain( name: str | None = None, description: str | None = None, retries: int | None = None, + max_uses: int | None = None, prepare: ToolPrepareFunc[AgentDepsT] | None = None, docstring_format: DocstringFormat = 'auto', require_parameter_descriptions: bool = False, @@ -1166,7 +1167,6 @@ def tool_plain( sequential: bool = False, requires_approval: bool = False, metadata: dict[str, Any] | None = None, - max_uses: int | None = None, ) -> Any: """Decorator to register a tool function which DOES NOT take `RunContext` as an argument. @@ -1203,6 +1203,7 @@ async def spam(ctx: RunContext[str]) -> float: description: The description of the tool, defaults to the function docstring. retries: The number of retries to allow for this tool, defaults to the agent's default retries, which defaults to 1. + max_uses: The maximum number of uses allowed for this tool during a run. Defaults to None (unlimited). prepare: custom method to prepare the tool definition for each step, return `None` to omit this tool from a given step. This is useful if you want to customise a tool at call time, or omit it completely from a step. See [`ToolPrepareFunc`][pydantic_ai.tools.ToolPrepareFunc]. @@ -1216,7 +1217,6 @@ async def spam(ctx: RunContext[str]) -> float: requires_approval: Whether this tool requires human-in-the-loop approval. Defaults to False. See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info. metadata: Optional metadata for the tool. This is not sent to the model but can be used for filtering and tool behavior customization. - max_uses: Optional maximum number of times this tool can be used during a run. Defaults to None (unlimited). """ def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams]: @@ -1227,6 +1227,7 @@ def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams name=name, description=description, retries=retries, + max_uses=max_uses, prepare=prepare, docstring_format=docstring_format, require_parameter_descriptions=require_parameter_descriptions, @@ -1235,7 +1236,6 @@ def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams sequential=sequential, requires_approval=requires_approval, metadata=metadata, - max_uses=max_uses, ) return func_ From 909716be502471c1cf58a04176f39004c43ad435 Mon Sep 17 00:00:00 2001 From: adtyavrdhn Date: Wed, 10 Dec 2025 15:07:15 +0530 Subject: [PATCH 4/9] documentation --- docs/tools-advanced.md | 2 +- docs/tools.md | 2 +- pydantic_ai_slim/pydantic_ai/_tool_manager.py | 1 - 3 files changed, 2 insertions(+), 3 deletions(-) diff --git a/docs/tools-advanced.md b/docs/tools-advanced.md index b498f247f6..f304e56f0c 100644 --- a/docs/tools-advanced.md +++ b/docs/tools-advanced.md @@ -379,7 +379,7 @@ If a tool requires sequential/serial execution, you can pass the [`sequential`][ Async functions are run on the event loop, while sync functions are offloaded to threads. To get the best performance, _always_ use an async function _unless_ you're doing blocking I/O (and there's no way to use a non-blocking library instead) or CPU-bound work (like `numpy` or `scikit-learn` operations), so that simple functions are not offloaded to threads unnecessarily. !!! note "Limiting tool executions" - You can cap tool executions within a run using [`UsageLimits(tool_calls_limit=...)`](agents.md#usage-limits). The counter increments only after a successful tool invocation. Output tools (used for [structured output](output.md)) are not counted in the `tool_calls` metric. + You can cap the total number of tool executions within a run using [`UsageLimits(tool_calls_limit=...)`](agents.md#usage-limits). For finer control, you can limit how many times a *specific* tool can be called by setting the `max_uses` parameter when registering the tool (e.g., `@agent.tool(max_uses=3)` or `Tool(func, max_uses=3)`). Once a tool reaches its `max_uses` limit, it is automatically removed from the available tools for subsequent steps in the run. The `tool_calls` counter increments only after a successful tool invocation. Output tools (used for [structured output](output.md)) are not counted in the `tool_calls` metric. ## See Also diff --git a/docs/tools.md b/docs/tools.md index 40dcf5c810..7d5b3c3211 100644 --- a/docs/tools.md +++ b/docs/tools.md @@ -361,7 +361,7 @@ _(This example is complete, it can be run "as is")_ For more tool features and integrations, see: -- [Advanced Tool Features](tools-advanced.md) - Custom schemas, dynamic tools, tool execution and retries +- [Advanced Tool Features](tools-advanced.md) - Custom schemas, dynamic tools, tool execution, retries, and usage limits - [Toolsets](toolsets.md) - Managing collections of tools - [Builtin Tools](builtin-tools.md) - Native tools provided by LLM providers - [Common Tools](common-tools.md) - Ready-to-use tool implementations diff --git a/pydantic_ai_slim/pydantic_ai/_tool_manager.py b/pydantic_ai_slim/pydantic_ai/_tool_manager.py index 8e6f5ad21d..231b380d3e 100644 --- a/pydantic_ai_slim/pydantic_ai/_tool_manager.py +++ b/pydantic_ai_slim/pydantic_ai/_tool_manager.py @@ -287,7 +287,6 @@ def get_max_use_of_tool(self, tool_name: str) -> int | None: return None return tool.max_uses - def get_current_use_of_tool(self, tool_name: str) -> int: """Get the current number of uses of a given tool.""" if self.ctx is None: From 3880ab690eed249310f4c8371072b8f7d1c7c51a Mon Sep 17 00:00:00 2001 From: adtyavrdhn Date: Wed, 10 Dec 2025 15:33:04 +0530 Subject: [PATCH 5/9] type --- pydantic_ai_slim/pydantic_ai/_tool_manager.py | 17 ++++++++++++++--- .../pydantic_ai/toolsets/abstract.py | 2 +- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_tool_manager.py b/pydantic_ai_slim/pydantic_ai/_tool_manager.py index 231b380d3e..eac3abdc8a 100644 --- a/pydantic_ai_slim/pydantic_ai/_tool_manager.py +++ b/pydantic_ai_slim/pydantic_ai/_tool_manager.py @@ -66,11 +66,22 @@ async def for_run_step(self, ctx: RunContext[AgentDepsT]) -> ToolManager[AgentDe @property def tool_defs(self) -> list[ToolDefinition]: - """The tool definitions for the tools in this tool manager.""" - if self.tools is None: + """The tool definitions for the tools in this tool manager. + + Tools that have reached their `max_uses` limit are filtered out. + """ + if self.tools is None or self.ctx is None: raise ValueError('ToolManager has not been prepared for a run step yet') # pragma: no cover - return [tool.tool_def for tool in self.tools.values()] + result: list[ToolDefinition] = [] + for tool in self.tools.values(): + # Filter out tools that have reached their max_uses limit + if tool.max_uses is not None: + current_uses = self.ctx.tool_usage.get(tool.tool_def.name, 0) + if current_uses >= tool.max_uses: + continue + result.append(tool.tool_def) + return result def should_call_sequentially(self, calls: list[ToolCallPart]) -> bool: """Whether to require sequential tool calls for a list of tool calls.""" diff --git a/pydantic_ai_slim/pydantic_ai/toolsets/abstract.py b/pydantic_ai_slim/pydantic_ai/toolsets/abstract.py index a75bcdbd30..a99f0f49e2 100644 --- a/pydantic_ai_slim/pydantic_ai/toolsets/abstract.py +++ b/pydantic_ai_slim/pydantic_ai/toolsets/abstract.py @@ -52,7 +52,7 @@ class ToolsetTool(Generic[AgentDepsT]): """The tool definition for this tool, including the name, description, and parameters.""" max_retries: int """The maximum number of retries to attempt if the tool call fails.""" - max_uses: int + max_uses: int | None """The maximum number of uses allowed for this tool.""" args_validator: SchemaValidator | SchemaValidatorProt """The Pydantic Core validator for the tool's arguments. From 1d62583b049f20826f92c514614fe20544a23299 Mon Sep 17 00:00:00 2001 From: adtyavrdhn Date: Wed, 10 Dec 2025 15:53:27 +0530 Subject: [PATCH 6/9] fix --- pydantic_ai_slim/pydantic_ai/_agent_graph.py | 22 +++++++++++++++---- pydantic_ai_slim/pydantic_ai/_tool_manager.py | 3 ++- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/_agent_graph.py b/pydantic_ai_slim/pydantic_ai/_agent_graph.py index 4400443920..9650f7a2ab 100644 --- a/pydantic_ai_slim/pydantic_ai/_agent_graph.py +++ b/pydantic_ai_slim/pydantic_ai/_agent_graph.py @@ -1001,6 +1001,17 @@ async def process_tool_calls( # noqa: C901 output_final_result.append(final_result) +def _projection_count_of_tool_usage( + tool_call_counts: defaultdict[str, int], tool_calls: list[_messages.ToolCallPart] +) -> None: + """Populate a count of tool usage based on the provided tool calls for this run step. + + We will use this to make sure the calls do not exceed tool usage limits. + """ + for call in tool_calls: + tool_call_counts[call.tool_name] += 1 + + async def _call_tools( tool_manager: ToolManager[DepsT], tool_calls: list[_messages.ToolCallPart], @@ -1024,10 +1035,9 @@ async def _call_tools( calls_to_run: list[_messages.ToolCallPart] = [] - # For each tool, check how many calls are going to be made and if it is over the limit + # For each tool, check how many calls are going to be made tool_call_counts: defaultdict[str, int] = defaultdict(int) - for call in tool_calls: - tool_call_counts[call.tool_name] += 1 + _projection_count_of_tool_usage(tool_call_counts, tool_calls) for call in tool_calls: current_tool_use = tool_manager.get_current_use_of_tool(call.tool_name) @@ -1113,7 +1123,11 @@ async def handle_call_or_result( output_parts.extend([user_parts_by_index[k] for k in sorted(user_parts_by_index)]) _populate_deferred_calls( - calls_to_run, deferred_calls_by_index, deferred_metadata_by_index, output_deferred_calls, output_deferred_metadata + calls_to_run, + deferred_calls_by_index, + deferred_metadata_by_index, + output_deferred_calls, + output_deferred_metadata, ) diff --git a/pydantic_ai_slim/pydantic_ai/_tool_manager.py b/pydantic_ai_slim/pydantic_ai/_tool_manager.py index eac3abdc8a..a9959da8fe 100644 --- a/pydantic_ai_slim/pydantic_ai/_tool_manager.py +++ b/pydantic_ai_slim/pydantic_ai/_tool_manager.py @@ -298,9 +298,10 @@ def get_max_use_of_tool(self, tool_name: str) -> int | None: return None return tool.max_uses + def get_current_use_of_tool(self, tool_name: str) -> int: """Get the current number of uses of a given tool.""" if self.ctx is None: raise ValueError('ToolManager has not been prepared for a run step yet') # pragma: no cover - return self.ctx.tool_usage.get(tool_name, 0) \ No newline at end of file + return self.ctx.tool_usage.get(tool_name, 0) From 114d2ebf15ca11be1c8de8c3d66efb8e37c234bf Mon Sep 17 00:00:00 2001 From: adtyavrdhn Date: Wed, 10 Dec 2025 16:07:17 +0530 Subject: [PATCH 7/9] better test --- tests/test_tools.py | 78 +++++++++++++++++++++++++++++++++++++-------- 1 file changed, 65 insertions(+), 13 deletions(-) diff --git a/tests/test_tools.py b/tests/test_tools.py index a96af78cc9..8a744df002 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -23,6 +23,7 @@ PrefixedToolset, RetryPromptPart, RunContext, + RunUsage, TextPart, Tool, ToolCallPart, @@ -1310,14 +1311,40 @@ def infinite_retry_tool(ctx: RunContext[None]) -> int: def test_tool_max_uses(): - agent = Agent(TestModel(), output_type=[str, DeferredToolRequests]) + """Test that a tool with max_uses=2 can only be called twice, and the third call is rejected.""" + call_count = 0 + + def my_model(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + nonlocal call_count + call_count += 1 + + if call_count == 1: + # First round: call the tool twice (will succeed, uses up the limit) + return ModelResponse( + parts=[ + ToolCallPart(tool_name='tool_with_max_use', args={}, tool_call_id='call_1'), + ToolCallPart(tool_name='tool_with_max_use', args={}, tool_call_id='call_2'), + ] + ) + elif call_count == 2: + # Second round: try to call the tool again (should be rejected) + return ModelResponse( + parts=[ + ToolCallPart(tool_name='tool_with_max_use', args={}, tool_call_id='call_3'), + ] + ) + else: + # Third round: return final output + return ModelResponse(parts=[TextPart(content='Done')]) + + agent = Agent(FunctionModel(my_model), output_type=str) - @agent.tool(max_uses=1) + @agent.tool(max_uses=2) def tool_with_max_use(ctx: RunContext[None]) -> str: return 'Used' result = agent.run_sync('Hello') - assert result.output == snapshot('{"tool_with_max_use":"Used"}') + assert result.output == snapshot('Done') messages = result.all_messages() assert messages == snapshot( [ @@ -1332,12 +1359,11 @@ def tool_with_max_use(ctx: RunContext[None]) -> str: ), ModelResponse( parts=[ - ToolCallPart( - tool_name='tool_with_max_use', args={}, tool_call_id='pyd_ai_tool_call_id__tool_with_max_use' - ) + ToolCallPart(tool_name='tool_with_max_use', args={}, tool_call_id='call_1'), + ToolCallPart(tool_name='tool_with_max_use', args={}, tool_call_id='call_2'), ], - usage=RequestUsage(input_tokens=51, output_tokens=2), - model_name='test', + usage=RequestUsage(input_tokens=51, output_tokens=4), + model_name=IsStr(), timestamp=IsDatetime(), run_id=IsStr(), ), @@ -1346,16 +1372,42 @@ def tool_with_max_use(ctx: RunContext[None]) -> str: ToolReturnPart( tool_name='tool_with_max_use', content='Used', - tool_call_id='pyd_ai_tool_call_id__tool_with_max_use', + tool_call_id='call_1', timestamp=IsDatetime(), - ) + ), + ToolReturnPart( + tool_name='tool_with_max_use', + content='Used', + tool_call_id='call_2', + timestamp=IsDatetime(), + ), ], run_id=IsStr(), ), ModelResponse( - parts=[TextPart(content='{"tool_with_max_use":"Used"}')], - usage=RequestUsage(input_tokens=52, output_tokens=6), - model_name='test', + parts=[ + ToolCallPart(tool_name='tool_with_max_use', args={}, tool_call_id='call_3'), + ], + usage=RequestUsage(input_tokens=53, output_tokens=6), + model_name=IsStr(), + timestamp=IsDatetime(), + run_id=IsStr(), + ), + ModelRequest( + parts=[ + ToolReturnPart( + tool_name='tool_with_max_use', + content='Tool call limit reached for tool "tool_with_max_use".', + tool_call_id='call_3', + timestamp=IsDatetime(), + ), + ], + run_id=IsStr(), + ), + ModelResponse( + parts=[TextPart(content='Done')], + usage=RequestUsage(input_tokens=61, output_tokens=7), + model_name=IsStr(), timestamp=IsDatetime(), run_id=IsStr(), ), From 64f359d06af22c42060413ce92bbd85545ffb754 Mon Sep 17 00:00:00 2001 From: adtyavrdhn Date: Wed, 10 Dec 2025 16:07:52 +0530 Subject: [PATCH 8/9] better test --- tests/test_tools.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_tools.py b/tests/test_tools.py index 8a744df002..264893db6a 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -23,7 +23,6 @@ PrefixedToolset, RetryPromptPart, RunContext, - RunUsage, TextPart, Tool, ToolCallPart, From 8ba5cc815ef5bf054a0baa0952b883a50e73447a Mon Sep 17 00:00:00 2001 From: adtyavrdhn Date: Wed, 10 Dec 2025 18:00:56 +0530 Subject: [PATCH 9/9] passing max_uses --- pydantic_ai_slim/pydantic_ai/_output.py | 1 + pydantic_ai_slim/pydantic_ai/mcp.py | 1 + pydantic_ai_slim/pydantic_ai/toolsets/external.py | 1 + pydantic_ai_slim/pydantic_ai/toolsets/fastmcp.py | 1 + 4 files changed, 4 insertions(+) diff --git a/pydantic_ai_slim/pydantic_ai/_output.py b/pydantic_ai_slim/pydantic_ai/_output.py index e3adfbd190..df1a67fc12 100644 --- a/pydantic_ai_slim/pydantic_ai/_output.py +++ b/pydantic_ai_slim/pydantic_ai/_output.py @@ -972,6 +972,7 @@ async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[ toolset=self, tool_def=tool_def, max_retries=self.max_retries, + max_uses=None, args_validator=self.processors[tool_def.name].validator, ) for tool_def in self._tool_defs diff --git a/pydantic_ai_slim/pydantic_ai/mcp.py b/pydantic_ai_slim/pydantic_ai/mcp.py index 227b8e1399..cc682c6fd3 100644 --- a/pydantic_ai_slim/pydantic_ai/mcp.py +++ b/pydantic_ai_slim/pydantic_ai/mcp.py @@ -590,6 +590,7 @@ def tool_for_tool_def(self, tool_def: ToolDefinition) -> ToolsetTool[Any]: toolset=self, tool_def=tool_def, max_retries=self.max_retries, + max_uses=None, args_validator=TOOL_SCHEMA_VALIDATOR, ) diff --git a/pydantic_ai_slim/pydantic_ai/toolsets/external.py b/pydantic_ai_slim/pydantic_ai/toolsets/external.py index 9ec4a0e0c7..adf4080bd8 100644 --- a/pydantic_ai_slim/pydantic_ai/toolsets/external.py +++ b/pydantic_ai_slim/pydantic_ai/toolsets/external.py @@ -36,6 +36,7 @@ async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[ toolset=self, tool_def=replace(tool_def, kind='external'), max_retries=0, + max_uses=None, args_validator=TOOL_SCHEMA_VALIDATOR, ) for tool_def in self.tool_defs diff --git a/pydantic_ai_slim/pydantic_ai/toolsets/fastmcp.py b/pydantic_ai_slim/pydantic_ai/toolsets/fastmcp.py index 2d907266fd..4c5a72c21d 100644 --- a/pydantic_ai_slim/pydantic_ai/toolsets/fastmcp.py +++ b/pydantic_ai_slim/pydantic_ai/toolsets/fastmcp.py @@ -170,6 +170,7 @@ def tool_for_tool_def(self, tool_def: ToolDefinition) -> ToolsetTool[AgentDepsT] tool_def=tool_def, toolset=self, max_retries=self.max_retries, + max_uses=None, args_validator=TOOL_SCHEMA_VALIDATOR, )