Add DeferredToolset

Author: DouweM

pydantic_ai_slim/pydantic_ai/toolset.py

@@ -319,6 +319,30 @@ async def call_tool(
         return output
 
 
+class DeferredToolset(AbstractToolset[AgentDepsT]):
+    """A toolset that holds deferred tool."""
+
+    _tool_defs: list[ToolDefinition]
+
+    def __init__(self, tool_defs: list[ToolDefinition]):
+        self._tool_defs = tool_defs
+
+    @property
+    def tool_defs(self) -> list[ToolDefinition]:
+        return [replace(tool_def, kind='deferred') for tool_def in self._tool_defs]
+
+    def _get_tool_args_validator(self, ctx: RunContext[AgentDepsT], name: str) -> SchemaValidator:
+        raise NotImplementedError('Deferred tools cannot be validated')
+
+    def _max_retries_for_tool(self, name: str) -> int:
+        raise NotImplementedError('Deferred tools cannot be retried')
+
+    async def call_tool(
+        self, ctx: RunContext[AgentDepsT], name: str, tool_args: dict[str, Any], *args: Any, **kwargs: Any
+    ) -> Any:
+        raise NotImplementedError('Deferred tools cannot be called')
+
+
 @dataclass(init=False)
 class CombinedToolset(AbstractToolset[AgentDepsT]):
     """A toolset that combines multiple toolsets."""

tests/test_tools.py

@@ -18,6 +18,7 @@
 from pydantic_ai.models.test import TestModel
 from pydantic_ai.output import DeferredToolCalls, ToolOutput
 from pydantic_ai.tools import ToolDefinition
+from pydantic_ai.toolset import DeferredToolset
 
 from .conftest import IsStr
 
@@ -1181,14 +1182,16 @@ def infinite_retry_tool(ctx: RunContext[None]) -> int:
 
 
 def test_deferred_tool():
-    agent = Agent(TestModel(), output_type=[str, DeferredToolCalls])
-
-    async def prepare_tool(ctx: RunContext[None], tool_def: ToolDefinition) -> ToolDefinition:
-        return replace(tool_def, kind='deferred')
-
-    @agent.tool_plain(prepare=prepare_tool)
-    def my_tool(x: int) -> int:
-        return x + 1
+    deferred_toolset = DeferredToolset(
+        [
+            ToolDefinition(
+                name='my_tool',
+                description='',
+                parameters_json_schema={'type': 'object', 'properties': {'x': {'type': 'integer'}}, 'required': ['x']},
+            ),
+        ]
+    )
+    agent = Agent(TestModel(), output_type=[str, DeferredToolCalls], toolsets=[deferred_toolset])
 
     result = agent.run_sync('Hello')
     assert result.output == snapshot(
@@ -1199,10 +1202,9 @@ def my_tool(x: int) -> int:
                     name='my_tool',
                     description='',
                     parameters_json_schema={
-                        'additionalProperties': False,
+                        'type': 'object',
                         'properties': {'x': {'type': 'integer'}},
                         'required': ['x'],
-                        'type': 'object',
                     },
                     kind='deferred',
                 )