Update Tool dataclass to accept a custom function_schema argument

pydantic_ai_slim/pydantic_ai/tools.py

@@ -149,6 +149,7 @@ class Tool(Generic[AgentDepsT]):
     max_retries: int | None
     name: str
     description: str
+    function_schema: _pydantic.FunctionSchema | None
     prepare: ToolPrepareFunc[AgentDepsT] | None
     docstring_format: DocstringFormat
     require_parameter_descriptions: bool
@@ -171,6 +172,7 @@ def __init__(
         max_retries: int | None = None,
         name: str | None = None,
         description: str | None = None,
+        function_schema: _pydantic.FunctionSchema | None = None,
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
@@ -217,6 +219,7 @@ async def prep_my_tool(
             max_retries: Maximum number of retries allowed for this tool, set to the agent default if `None`.
             name: Name of the tool, inferred from the function if `None`.
             description: Description of the tool, inferred from the function if `None`.
+            function_schema: Function schema 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
                 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].
@@ -227,7 +230,10 @@ async def prep_my_tool(
         if takes_ctx is None:
             takes_ctx = _pydantic.takes_ctx(function)
 
-        f = _pydantic.function_schema(function, takes_ctx, docstring_format, require_parameter_descriptions)
+        f = self.function_schema = function_schema or _pydantic.function_schema(
+            function, takes_ctx, docstring_format, require_parameter_descriptions
+        )
+
         self.function = function
         self.takes_ctx = takes_ctx
         self.max_retries = max_retries

tests/test_tools.py

@@ -9,7 +9,7 @@
 from pydantic import BaseModel, Field
 from pydantic_core import PydanticSerializationError
 
-from pydantic_ai import Agent, RunContext, Tool, UserError
+from pydantic_ai import Agent, RunContext, Tool, UserError, _pydantic
 from pydantic_ai.messages import (
     ModelMessage,
     ModelRequest,
@@ -345,6 +345,40 @@ def plain_tool(x: int) -> int:
     assert agent_infer._function_tools['plain_tool'].max_retries == 7
 
 
+def test_init_tool_with_function_schema():
+    def x_tool(x: int) -> None:
+        pass
+
+    def y_tool(y: str) -> None:
+        pass
+
+    y_fs = _pydantic.function_schema(
+        y_tool, takes_ctx=False, docstring_format='auto', require_parameter_descriptions=False
+    )
+    agent = Agent('test', tools=[Tool(x_tool, function_schema=y_fs)])
+
+    # make sure the function schema for y_tool is used instead of the default of x_tool.
+    assert 'x' not in agent._function_tools['x_tool']._parameters_json_schema['properties']
+    assert 'y' in agent._function_tools['x_tool']._parameters_json_schema['properties']
+
+
+def test_init_tool_ctx_with_function_schema():
+    def x_tool(ctx: RunContext[int], x: int) -> None:
+        pass
+
+    def y_tool(ctx: RunContext[int], y: str) -> None:
+        pass
+
+    y_fs = _pydantic.function_schema(
+        y_tool, takes_ctx=True, docstring_format='auto', require_parameter_descriptions=False
+    )
+    agent = Agent('test', tools=[Tool(x_tool, function_schema=y_fs, takes_ctx=True)], deps_type=int)
+
+    # make sure the function schema for y_tool is used instead of the default of x_tool.
+    assert 'x' not in agent._function_tools['x_tool']._parameters_json_schema['properties']
+    assert 'y' in agent._function_tools['x_tool']._parameters_json_schema['properties']
+
+
 def ctx_tool(ctx: RunContext[int], x: int) -> int:
     return x + ctx.deps