Let ApprovalRequiredToolset function be dynamic on tool_args

Author: DouweM

pydantic_ai_slim/pydantic_ai/toolsets/abstract.py

@@ -177,7 +177,9 @@ def renamed(self, name_map: dict[str, str]) -> RenamedToolset[AgentDepsT]:
 
     def approval_required(
         self,
-        approval_required_func: Callable[[RunContext[AgentDepsT], ToolDefinition], bool] = lambda ctx, tool_def: True,
+        approval_required_func: Callable[[RunContext[AgentDepsT], ToolDefinition, dict[str, Any]], bool] = (
+            lambda ctx, tool_def, tool_args: True
+        ),
     ) -> ApprovalRequiredToolset[AgentDepsT]:
         """TODO: Docstring."""
         from .approval_required import ApprovalRequiredToolset

pydantic_ai_slim/pydantic_ai/toolsets/approval_required.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
-from dataclasses import dataclass, replace
-from typing import Callable
+from dataclasses import dataclass
+from typing import Any, Callable
+
+from pydantic_ai.exceptions import ApprovalRequired
 
 from .._run_context import AgentDepsT, RunContext
 from ..tools import ToolDefinition
@@ -13,12 +15,14 @@
 class ApprovalRequiredToolset(WrapperToolset[AgentDepsT]):
     """TODO: Docstring."""
 
-    approval_required_func: Callable[[RunContext[AgentDepsT], ToolDefinition], bool] = lambda ctx, tool_def: True
+    approval_required_func: Callable[[RunContext[AgentDepsT], ToolDefinition, dict[str, Any]], bool] = (
+        lambda ctx, tool_def, tool_args: True
+    )
+
+    async def call_tool(
+        self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT]
+    ) -> Any:
+        if not ctx.tool_call_approved and self.approval_required_func(ctx, tool.tool_def, tool_args):
+            raise ApprovalRequired
 
-    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
-        return {
-            name: replace(tool, tool_def=replace(tool.tool_def, kind='unapproved'))
-            if not ctx.tool_call_approved and self.approval_required_func(ctx, tool.tool_def)
-            else tool
-            for name, tool in (await super().get_tools(ctx)).items()
-        }
+        return await super().call_tool(name, tool_args, ctx, tool)

tests/test_tools.py

@@ -1693,7 +1693,7 @@ def foo(x: int) -> int:
     def bar(x: int) -> int:
         return x * 3
 
-    toolset = toolset.approval_required(lambda ctx, tool_def: tool_def.name == 'foo')
+    toolset = toolset.approval_required(lambda ctx, tool_def, tool_args: tool_def.name == 'foo')
 
     agent = Agent(FunctionModel(llm), toolsets=[toolset], output_type=[str, DeferredToolRequests])