Add per-tool usage limits

[samvanstroud]

Description

Pydantic AI supports per-agent tool call limits, and per-tool retry limits, but it's currently not possible to specify per-tool usage limits, which seems like a gap. This would be useful in cases where some tools are much more expensive than others, and so their usage needs to be more tightly controlled.

For example, it would be great to be able to do

@agent.tool(max_uses=2)
async def expensive_tool(ctx: RunContext[Deps], input: str) -> str:
    ...

@agent.tool(max_uses=10)
async def cheap_tool(ctx: RunContext[Deps], input: str) -> str:
    ...

The implementation should closely follow the tool max_retries and agent usage_limits setting:

• The initial system prompt should include information about the max uses for each tool.
• Calling a tool too many times should raise a UsageLimitExceeded or similar exception.
• The agent should be allowed to call other tools according to usage_limits and any other per-tool usage limits

This should prevent some boilerplate in adding usage counts and limits to the deps and checking this in the tool definition, as in:

@agent.tool
async def my_tool(ctx: RunContext[Deps], input: str) -> str:
    if ctx.deps.tool_uses >= ctx.deps.tool_use_limit:
        return "You have already used the my_tool tool the maximum number of allowed times."

    # tool logic
    ...

    ctx.deps.tool_uses += 1

Given potential issues with parallel tools (#3120), I think it's worth a robust pydantic-ai side implementation here.

[DouweM]

@samvanstroud Can you please expand on what type of limits you'd like to put, and what you'd like to happen when the limit is hit? Specific use case examples help a lot.

[chmielcode]

I hope I found the right issue.

@DouweM To have a one-shot tool that is hidden from the model after being called once I keep a one_shot_tool_calls set in run context and check if the tool name is already in it. If it is, it's tool_def is filtered out in prepare function.

I do not like it, because I need to add a technical variable (one_shot_tool_calls) to the run context and the tool has to add itself to the one_shot_tool_calls.

Some alternative would be to get tool calls from message history, but since I do not see a separate "message history within this run", I would have to add another technical variable e.g. run_start_timestamp to count only current tool calls on subsequent runs extending the same message history.

Another problem is when we are within one run step and the model does parallel function calling to the one-shot tool. How to stop the model technically (not by prompting or returning "do not call me again" - this works, but it is not real [and nice] control) from reusing this tool? Sequential=true appears too harsh since it makes other tool calls withing the same step sequential as well (I think).

To get some control we could have @agent.tool(per_run_call_limit=N) to auto-hide tool after it's used N times or raise UnexpectedModelBehavior* if N was exceeded within one run step.

*this is the interesting part - if there was something like ModelPleaseDont exception or other reliable way to stop model from calling a tool within the same step and not break the run (use previous call output) at the same time, solutions above would have become a little cleaner. ModelRetry("Do not retry.") appears wrong and I believe it is not to be used this way.

This or maybe I am missing something obvious.

[samvanstroud]

Thanks @DouweM, I added a bit more detail to the original comment. I think @chmielcode has a similar request.

Let me know what you think.

[DouweM]

@samvanstroud

• The initial system prompt should include information about the max uses for each tool.

Would you expect Pydantic AI to insert this into the tool description that's passed to the model automatically? We currently just use the tool function's docstring or @agent.tool(description=...) argument, and never insert our own text into the system prompt/instructions unless Prompted Output is used, so I'm hesitant to start doing this now.

I'm also not sure what the best way to communicate this information to the model would be:

• "This tool can only be used N times." would rely on the model to know how often it's already used it, but the history may be composed of multiple runs while the limit is only enforced on the current run, so it could be overcounting and use the tool less times than it's allowed to.
• "This tool can only be used M more times." or "This tool can only be used N times and has already been used M times." would result in the prompt cache being invalidated when the tool is used, which is not ideal.

A better approach may be to have the tool when output "This tool can only be called M more times." after its real result, so that it works like a counter that increases during the run of the conversation. But that'd definitely be something for the user to do, I think, not Pydantic AI itself.

• Calling a tool too many times should raise a UsageLimitExceeded or similar exception.

Note that @chmielcode wants the tool to be disabled after its limit is hit, which would be less likely to blow up the agent run as the model can't stubbornly decide to call the tool anyway. Would that work for you or do you prefer the error?

This should prevent some boilerplate in adding usage counts and limits to the deps and checking this in the tool definition, as in:

Note that you could implement the usage tracking (and automatic disabling after X calls) using a custom WrapperToolset: https://ai.pydantic.dev/toolsets/#changing-tool-execution

---

@chmielcode

Some alternative would be to get tool calls from message history, but since I do not see a separate "message history within this run", I would have to add another technical variable e.g. run_start_timestamp to count only current tool calls on subsequent runs extending the same message history.

This approach will get a lot easier to implement once we get #3171 in the next few days!

Another problem is when we are within one run step and the model does parallel function calling to the one-shot tool. How to stop the model technically (not by prompting or returning "do not call me again" - this works, but it is not real [and nice] control) from reusing this tool? Sequential=true appears too harsh since it makes other tool calls withing the same step sequential as well (I think).

If the model calls the limited tool multiple times in parallel and one or more of those calls would be over the limit, I think our only option is to return a tool result like "This tool call was not executed because the tool was already used X times." Then on the next run step the tool would not be available at all anymore.

*this is the interesting part - if there was something like ModelPleaseDont exception or other reliable way to stop model from calling a tool within the same step and not break the run (use previous call output) at the same time, solutions above would have become a little cleaner. ModelRetry("Do not retry.") appears wrong and I believe it is not to be used this way.

This is related to #2586 (comment). Some new exception that results in the tool being disabled for future steps would be great.

[samvanstroud]

Thanks @DouweM for the detailed reply.

Regarding informing the model about usage information, what is the approach used for the existing UsageLimit.tool_calls_limit? Can we use the same approach here?

Note that @chmielcode wants the tool to be disabled after its limit is hit, which would be less likely to blow up the agent run as the model can't stubbornly decide to call the tool anyway. Would that work for you or do you prefer the error?

This sounds fine to me.

[DouweM]

Regarding informing the model about usage information, what is the approach used for the existing UsageLimit.tool_calls_limit? Can we use the same approach here?

@samvanstroud The model is not informed of that limit at all, so we could definitely do the same thing here 😄 If you want the model to be aware of the limit you'd have to tell it yourself.