Update Tool dataclass to accept a custom function_schema argument

pydantic_ai_slim/pydantic_ai/_pydantic.py

@@ -6,7 +6,7 @@
 from __future__ import annotations as _annotations
 
 from inspect import Parameter, signature
-from typing import TYPE_CHECKING, Any, Callable, TypedDict, cast, get_origin
+from typing import TYPE_CHECKING, Any, Callable, cast, get_origin
 
 from pydantic import ConfigDict
 from pydantic._internal import _decorators, _generate_schema, _typing_extra
@@ -20,24 +20,12 @@
 from ._utils import check_object_json_schema, is_model_like
 
 if TYPE_CHECKING:
-    from .tools import DocstringFormat, ObjectJsonSchema
+    from .tools import DocstringFormat, FunctionSchema
 
 
 __all__ = ('function_schema',)
 
 
-class FunctionSchema(TypedDict):
-    """Internal information about a function schema."""
-
-    description: str
-    validator: SchemaValidator
-    json_schema: ObjectJsonSchema
-    # if not None, the function takes a single by that name (besides potentially `info`)
-    single_arg_name: str | None
-    positional_fields: list[str]
-    var_positional_field: str | None
-
-
 def function_schema(  # noqa: C901
     function: Callable[..., Any],
     takes_ctx: bool,
@@ -161,14 +149,14 @@ def function_schema(  # noqa: C901
         # and set it on the tool
         description = json_schema.pop('description', None)
 
-    return FunctionSchema(
-        description=description,
-        validator=schema_validator,
-        json_schema=check_object_json_schema(json_schema),
-        single_arg_name=single_arg_name,
-        positional_fields=positional_fields,
-        var_positional_field=var_positional_field,
-    )
+    return {
+        'description': description,
+        'validator': schema_validator,
+        'json_schema': check_object_json_schema(json_schema),
+        'single_arg_name': single_arg_name,
+        'positional_fields': positional_fields,
+        'var_positional_field': var_positional_field,
+    }
 
 
 def takes_ctx(function: Callable[..., Any]) -> bool:

pydantic_ai_slim/pydantic_ai/tools.py

@@ -4,7 +4,7 @@
 import inspect
 from collections.abc import Awaitable
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Union, cast
+from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, TypedDict, Union, cast
 
 from pydantic import ValidationError
 from pydantic_core import SchemaValidator
@@ -19,6 +19,7 @@
 __all__ = (
     'AgentDepsT',
     'DocstringFormat',
+    'FunctionSchema',
     'RunContext',
     'SystemPromptFunc',
     'ToolFuncContext',
@@ -35,6 +36,18 @@
 """Type variable for agent dependencies."""
 
 
+class FunctionSchema(TypedDict):
+    """Internal information about a function schema."""
+
+    description: str
+    validator: SchemaValidator
+    json_schema: ObjectJsonSchema
+    # if not None, the function takes a single by that name (besides potentially `info`)
+    single_arg_name: str | None
+    positional_fields: list[str]
+    var_positional_field: str | None
+
+
 @dataclasses.dataclass
 class RunContext(Generic[AgentDepsT]):
     """Information about the current call."""
@@ -149,6 +162,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 +185,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 +232,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 +243,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,51 @@ 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:
+        raise NotImplementedError
+
+    def y_tool(y: str) -> None:
+        raise NotImplementedError
+
+    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 agent._function_tools['x_tool']._parameters_json_schema == snapshot(
+        {
+            'additionalProperties': False,
+            'properties': {'y': {'title': 'Y', 'type': 'string'}},
+            'required': ['y'],
+            'type': 'object',
+        }
+    )
+
+
+def test_init_tool_ctx_with_function_schema():
+    def x_tool(ctx: RunContext[int], x: int) -> None:
+        raise NotImplementedError
+
+    def y_tool(ctx: RunContext[int], y: str) -> None:
+        raise NotImplementedError
+
+    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)
+
+    assert agent._function_tools['x_tool']._parameters_json_schema == snapshot(
+        {
+            'additionalProperties': False,
+            'properties': {'y': {'title': 'Y', 'type': 'string'}},
+            'required': ['y'],
+            'type': 'object',
+        }
+    )
+
+
 def ctx_tool(ctx: RunContext[int], x: int) -> int:
     return x + ctx.deps