More flexible tool exception handling

[strawgate]

Description

I think function Toolset is great for taking python code and making it runnable by an Agent.

An exception raised during a function call ends the agent run unless it's a model retry.

This means that often times users have to modify their functions in offer to use them with a function Toolset

We should think about making the default behavior that it catches exceptions from the function call and raises a model retry and perhaps have an option on the function tool/toolset to do something else (maybe a new exception like ModelFunctionToolError)

References

No response

[DouweM]

@strawgate What kinds of exceptions are you thinking of? Sharing the error message with the model and telling it to retry may make sense in some cases, but really all it can do is change its arguments or just try again verbatim, and the latter would be better handled by retrying the tool call on the Python rather than LLM side (using tenacity like in https://ai.pydantic.dev/retries/). Other errors may also be unrecoverable.

Having some way to configure how to handle specific types of exceptions at the tool or FunctionToolset level makes sense to me, as an alternative to having to repeat except statements in every tool function, but I don't think treating every exception as "model retryable" by default would make sense. I'd imagine a dict of exception classes (or tuples thereof) pointing to ModelRetry or a tenacity.AsyncRetrying, or some other classes/string literals codifying exception handling behavior we may like. If you want to retry all errors, you'd end up with something like on_exception={Exception: ModelRetry} (all names TBD).

How does that sound to you?

[jcaplan]

I don't think treating every exception as "model retryable" by default would make sense

I definitely think it can make sense depending on the tools you're using.

I decided to switch over from openai agents sdk because on the surface they're quite similar so not a lot of code to rework, but overall there's a few key things in pydantic-ai I like better. but this issue is really hurting me right now.

[DouweM]

I definitely think it can make sense depending on the tools you're using.

@jcaplan I agree, I'm saying it probably shouldn't be the default for Pydantic AI, but it should definitely be something users should be able to configure.

Note that you can already implement this using a custom WrapperToolset: https://ai.pydantic.dev/toolsets/#changing-tool-execution. Something like:

class RetryingToolset(WrapperToolset):
    async def call_tool(self, name: str, tool_args: dict[str, Any], ctx: RunContext, tool: ToolsetTool) -> Any:
        try:
            result = await super().call_tool(name, tool_args, ctx, tool)
        except ModelRetry:
            raise
        except Exception as e:
            raise ModelRetry(str(e)) from e
        else:
            return result

[jcaplan]

That's really helpful thanks!

[jcaplan]

seems like each tool has its own tool count though

from typing import Any
from pydantic_ai import Agent, RunContext, ModelRetry
from pydantic_ai.toolsets.wrapper import WrapperToolset
from pydantic_ai.toolsets.abstract import ToolsetTool
from pydantic_ai.toolsets import FunctionToolset


class RetryingToolset(WrapperToolset):
    async def call_tool(self, name: str, tool_args: dict[str, Any], ctx: RunContext, tool: ToolsetTool) -> Any:
        try:
            result = await super().call_tool(name, tool_args, ctx, tool)
        except ModelRetry:
            raise
        except Exception as e:
            raise ModelRetry(str(e)) from e
        else:
            return result


@pytest.mark.asyncio
async def test_tool_retry():
    count = 0
    def tool() -> str:
        """Test tool for retrying until success

        Raises:
            ValueError: If the tool fails

        Returns:
            str: The result of the tool
        """
        nonlocal count
        count += 1
        if count < 3:
            raise ValueError("Simulated failure")

        return "Success"

    toolset = RetryingToolset(FunctionToolset(tools=[tool]))
    agent = Agent(
        name="retry_agent",
        toolsets=[toolset],
        retries=5,
    )

    await agent.run("call the tool until it returns success (up to 5 times)")

gives an error:

            if current_retry == max_retries:
>               raise UnexpectedModelBehavior(f'Tool {name!r} exceeded max retries count of {max_retries}') from e
E               pydantic_ai.exceptions.UnexpectedModelBehavior: Tool 'tool' exceeded max retries count of 1

.venv/lib/python3.12/site-packages/pydantic_ai/_tool_manager.py:156: UnexpectedModelBehavior

[jcaplan]

it seems that the toolset has a default retry of 1 which overrides what I set for the agent. that's a bit surprising (to me at least)

[jcaplan]

lets say i'm running a bunch of tools in parallel and if one of them throws an exception, and the LLM should still partially complete whatever it's working on with the successful tool results. I didn't try yet but I think that still won't work here? how can the llm proceed even once max_retries is exceeded?

[DouweM]

lets say i'm running a bunch of tools in parallel and if one of them throws an exception, and the LLM should still partially complete whatever it's working on with the successful tool results. I didn't try yet but I think that still won't work here? how can the llm proceed even once max_retries is exceeded?

@jcaplan If you don't want a failed to tool to be retried, but just to send a message to the model like "Couldn't generate result", you can return that string to the model directly, or have your RetryingToolset return text instead of raise ModelRetry. If you want the tool to be retried once, and then say something like "Retries exhausted" to the model before the retry count is about to be hit, you can compare ctx.retry against tool.max_retries in that call_tool function. We're discussing letting that done from individual tool functions as well in #2927.