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, cast
+from typing import TYPE_CHECKING, Any, Callable, cast, get_origin
 
 from pydantic import ConfigDict
 from pydantic._internal import _decorators, _generate_schema, _typing_extra
@@ -15,30 +15,17 @@
 from pydantic.json_schema import GenerateJsonSchema
 from pydantic.plugin._schema_validator import create_schema_validator
 from pydantic_core import SchemaValidator, core_schema
-from typing_extensions import TypedDict, get_origin
 
 from ._griffe import doc_descriptions
 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,
@@ -162,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

@@ -9,7 +9,7 @@
 from pydantic import ValidationError
 from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
 from pydantic_core import SchemaValidator, core_schema
-from typing_extensions import Concatenate, ParamSpec, TypeAlias, TypeVar
+from typing_extensions import Concatenate, ParamSpec, TypeAlias, TypedDict, TypeVar
 
 from . import _pydantic, _utils, messages as _messages, models
 from .exceptions import ModelRetry, UnexpectedModelBehavior
@@ -20,6 +20,7 @@
 __all__ = (
     'AgentDepsT',
     'DocstringFormat',
+    'FunctionSchema',
     'RunContext',
     'SystemPromptFunc',
     'ToolFuncContext',
@@ -36,6 +37,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."""
@@ -168,6 +181,7 @@ class Tool(Generic[AgentDepsT]):
     max_retries: int | None
     name: str
     description: str
+    function_schema: FunctionSchema | None
     prepare: ToolPrepareFunc[AgentDepsT] | None
     docstring_format: DocstringFormat
     require_parameter_descriptions: bool
@@ -190,6 +204,7 @@ def __init__(
         max_retries: int | None = None,
         name: str | None = None,
         description: str | None = None,
+        function_schema: FunctionSchema | None = None,
         prepare: ToolPrepareFunc[AgentDepsT] | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
@@ -237,6 +252,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].
@@ -248,7 +264,7 @@ async def prep_my_tool(
         if takes_ctx is None:
             takes_ctx = _pydantic.takes_ctx(function)
 
-        f = _pydantic.function_schema(
+        f = self.function_schema = function_schema or _pydantic.function_schema(
             function, takes_ctx, docstring_format, require_parameter_descriptions, schema_generator
         )
         self.function = function

tests/test_tools.py

@@ -10,7 +10,7 @@
 from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
 from pydantic_core import PydanticSerializationError, core_schema
 
-from pydantic_ai import Agent, RunContext, Tool, UserError
+from pydantic_ai import Agent, RunContext, Tool, UserError, _pydantic
 from pydantic_ai.messages import (
     ModelMessage,
     ModelRequest,
@@ -21,7 +21,7 @@
 )
 from pydantic_ai.models.function import AgentInfo, FunctionModel
 from pydantic_ai.models.test import TestModel
-from pydantic_ai.tools import ToolDefinition
+from pydantic_ai.tools import GenerateToolJsonSchema, ToolDefinition
 
 
 def test_tool_no_ctx():
@@ -346,6 +346,59 @@ 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,
+        schema_generator=GenerateToolJsonSchema,
+    )
+    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': {'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,
+        schema_generator=GenerateToolJsonSchema,
+    )
+    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': {'type': 'string'}},
+            'required': ['y'],
+            'type': 'object',
+        }
+    )
+
+
 def ctx_tool(ctx: RunContext[int], x: int) -> int:
     return x + ctx.deps