Skip to content

Commit 9c51f80

Browse files
DouweMadtyavrdhnclaude
authored
feat: add return_schema and function_signature to ToolDefinition (#4964)
Co-authored-by: Aditya Vardhan <[email protected]> Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
1 parent 0c0cb3b commit 9c51f80

40 files changed

Lines changed: 4008 additions & 64 deletions

‎docs/api/function_signature.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
# `pydantic_ai.function_signature`
2+
3+
::: pydantic_ai.function_signature

‎docs/api/toolsets.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,11 @@
99
- ApprovalRequiredToolset
1010
- FilteredToolset
1111
- FunctionToolset
12+
- IncludeReturnSchemasToolset
1213
- DeferredLoadingToolset
1314
- PrefixedToolset
1415
- RenamedToolset
16+
- SetMetadataToolset
1517
- PreparedToolset
1618
- WrapperToolset
1719
- ToolsetFunc

‎docs/capabilities.md‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@ Pydantic AI ships with several capabilities that cover common needs:
2727
| [`PrefixTools`][pydantic_ai.capabilities.PrefixTools] | Wraps a capability and prefixes its tool names | Yes |
2828
| [`BuiltinTool`][pydantic_ai.capabilities.BuiltinTool] | Registers a [builtin tool](builtin-tools.md) with the agent | Yes |
2929
| [`Toolset`][pydantic_ai.capabilities.Toolset] | Wraps an [`AbstractToolset`][pydantic_ai.toolsets.AbstractToolset] | — |
30+
| [`IncludeToolReturnSchemas`][pydantic_ai.capabilities.IncludeToolReturnSchemas] | Includes return type schemas in tool definitions sent to the model | Yes |
31+
| [`SetToolMetadata`][pydantic_ai.capabilities.SetToolMetadata] | Merges metadata key-value pairs onto selected tools | Yes |
3032
| [`HistoryProcessor`][pydantic_ai.capabilities.HistoryProcessor] | Wraps a [history processor](message-history.md#processing-message-history) | — |
3133
| [`ThreadExecutor`][pydantic_ai.capabilities.ThreadExecutor] | Uses a custom thread executor for [sync functions](tools-advanced.md#thread-executor-for-long-running-servers) | — |
3234

@@ -217,6 +219,117 @@ Every [`AbstractCapability`][pydantic_ai.capabilities.AbstractCapability] has a
217219
MCP(url='https://mcp.example.com/api').prefix_tools('mcp')
218220
```
219221

222+
### IncludeToolReturnSchemas
223+
224+
[`IncludeToolReturnSchemas`][pydantic_ai.capabilities.IncludeToolReturnSchemas] includes return type schemas in tool definitions sent to the model. For models that natively support return schemas (e.g. Google Gemini), the schema is passed as a structured field in the API request. For other models, it is injected into the tool description as JSON text.
225+
226+
```python {title="include_return_schemas.py" lint="skip"}
227+
from pydantic_ai import Agent
228+
from pydantic_ai.capabilities import IncludeToolReturnSchemas
229+
from pydantic_ai.models.test import TestModel
230+
231+
232+
test_model = TestModel()
233+
agent = Agent(test_model, capabilities=[IncludeToolReturnSchemas()])
234+
235+
236+
@agent.tool_plain
237+
def get_temperature(city: str) -> float:
238+
"""Get the temperature for a city."""
239+
return 21.0
240+
241+
242+
result = agent.run_sync('What is the temperature in Paris?')
243+
params = test_model.last_model_request_parameters
244+
assert params is not None
245+
td = params.function_tools[0]
246+
assert td.include_return_schema is True
247+
```
248+
249+
_(This example is complete, it can be run "as is")_
250+
251+
Use the `tools` parameter to select which tools should include return schemas. It accepts a list of tool names, a metadata dict for matching, or a callable predicate:
252+
253+
```python {title="include_return_schemas_selective.py" lint="skip"}
254+
from pydantic_ai import Agent
255+
from pydantic_ai.capabilities import IncludeToolReturnSchemas
256+
from pydantic_ai.models.test import TestModel
257+
258+
259+
test_model = TestModel()
260+
agent = Agent(
261+
test_model,
262+
capabilities=[IncludeToolReturnSchemas(tools=['get_temperature'])],
263+
)
264+
265+
266+
@agent.tool_plain
267+
def get_temperature(city: str) -> float:
268+
"""Get the temperature for a city."""
269+
return 21.0
270+
271+
272+
@agent.tool_plain
273+
def get_greeting(name: str) -> str:
274+
"""Get a greeting."""
275+
return f'Hello, {name}!'
276+
277+
278+
result = agent.run_sync('Hello')
279+
params = test_model.last_model_request_parameters
280+
assert params is not None
281+
temp_tool = next(t for t in params.function_tools if t.name == 'get_temperature')
282+
greet_tool = next(t for t in params.function_tools if t.name == 'get_greeting')
283+
assert temp_tool.include_return_schema is True
284+
assert greet_tool.include_return_schema is None
285+
```
286+
287+
_(This example is complete, it can be run "as is")_
288+
289+
The same effect can be achieved at the toolset level using [`.include_return_schemas()`][pydantic_ai.toolsets.AbstractToolset.include_return_schemas] — see [toolset composition](toolsets.md#including-return-schemas).
290+
291+
### SetToolMetadata
292+
293+
[`SetToolMetadata`][pydantic_ai.capabilities.SetToolMetadata] merges metadata key-value pairs onto selected tools. This is useful for tagging tools with configuration that other capabilities or custom logic can inspect:
294+
295+
```python {title="set_tool_metadata.py" lint="skip"}
296+
from pydantic_ai import Agent
297+
from pydantic_ai.capabilities import SetToolMetadata
298+
from pydantic_ai.models.test import TestModel
299+
300+
301+
test_model = TestModel()
302+
agent = Agent(
303+
test_model,
304+
capabilities=[SetToolMetadata(tools=['search'], sensitive=True)],
305+
)
306+
307+
308+
@agent.tool_plain
309+
def search(query: str) -> str:
310+
"""Search for information."""
311+
return f'Results for: {query}'
312+
313+
314+
@agent.tool_plain
315+
def greet(name: str) -> str:
316+
"""Greet someone."""
317+
return f'Hello, {name}!'
318+
319+
320+
result = agent.run_sync('Search for pydantic')
321+
params = test_model.last_model_request_parameters
322+
assert params is not None
323+
search_tool = next(t for t in params.function_tools if t.name == 'search')
324+
greet_tool = next(t for t in params.function_tools if t.name == 'greet')
325+
assert search_tool.metadata is not None and search_tool.metadata.get('sensitive') is True
326+
assert greet_tool.metadata is None or greet_tool.metadata.get('sensitive') is None
327+
```
328+
329+
_(This example is complete, it can be run "as is")_
330+
331+
The same effect can be achieved at the toolset level using [`.with_metadata()`][pydantic_ai.toolsets.AbstractToolset.with_metadata] — see [toolset composition](toolsets.md#setting-tool-metadata).
332+
220333
## Building custom capabilities
221334

222335
To build your own capability, subclass [`AbstractCapability`][pydantic_ai.capabilities.AbstractCapability] and override the methods you need. There are two categories: **configuration methods** that are called at agent construction (except [`get_wrapper_toolset`][pydantic_ai.capabilities.AbstractCapability.get_wrapper_toolset] which is called per-run), and **lifecycle hooks** that fire during each run.

‎docs/toolsets.md‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -526,6 +526,67 @@ mcp = MCPServerHTTP('http://localhost:8000/mcp')
526526
agent = Agent('openai:gpt-5.2', toolsets=[mcp.defer_loading()])
527527
```
528528

529+
### Including Return Schemas
530+
531+
[`IncludeReturnSchemasToolset`][pydantic_ai.toolsets.IncludeReturnSchemasToolset] wraps a toolset and sets `include_return_schema=True` on all its tools, causing the model to receive return type information. For models that natively support return schemas (e.g. Google Gemini), the schema is passed as a structured API field. For other models, it is injected into the tool description as JSON text.
532+
533+
To easily chain different modifications, you can also call [`.include_return_schemas()`][pydantic_ai.toolsets.AbstractToolset.include_return_schemas] on any toolset instead of directly constructing an `IncludeReturnSchemasToolset`.
534+
535+
```python {title="include_return_schemas_toolset.py"}
536+
from pydantic_ai import Agent, FunctionToolset
537+
from pydantic_ai.models.test import TestModel
538+
539+
540+
def get_temperature(city: str) -> float:
541+
"""Get the temperature for a city."""
542+
return 21.0
543+
544+
545+
toolset = FunctionToolset(tools=[get_temperature])
546+
547+
test_model = TestModel()
548+
agent = Agent(test_model, toolsets=[toolset.include_return_schemas()])
549+
result = agent.run_sync('What is the temperature?')
550+
params = test_model.last_model_request_parameters
551+
assert params is not None
552+
assert params.function_tools[0].include_return_schema is True
553+
```
554+
555+
_(This example is complete, it can be run "as is")_
556+
557+
This is the toolset-level equivalent of the [`IncludeToolReturnSchemas`][pydantic_ai.capabilities.IncludeToolReturnSchemas] capability, which applies across all toolsets or a selected subset.
558+
559+
### Setting Tool Metadata
560+
561+
[`SetMetadataToolset`][pydantic_ai.toolsets.SetMetadataToolset] wraps a toolset and merges metadata key-value pairs onto all its tools. This is useful for tagging tools with configuration that other capabilities or custom logic can inspect.
562+
563+
To easily chain different modifications, you can also call [`.with_metadata()`][pydantic_ai.toolsets.AbstractToolset.with_metadata] on any toolset instead of directly constructing a `SetMetadataToolset`.
564+
565+
```python {title="set_metadata_toolset.py"}
566+
from pydantic_ai import Agent, FunctionToolset
567+
from pydantic_ai.models.test import TestModel
568+
569+
570+
def search(query: str) -> str:
571+
"""Search for information."""
572+
return f'Results for: {query}'
573+
574+
575+
toolset = FunctionToolset(tools=[search])
576+
577+
test_model = TestModel()
578+
agent = Agent(test_model, toolsets=[toolset.with_metadata(sensitive=True)])
579+
result = agent.run_sync('Search for something')
580+
params = test_model.last_model_request_parameters
581+
assert params is not None
582+
assert params.function_tools[0].metadata is not None
583+
assert params.function_tools[0].metadata['sensitive'] is True
584+
```
585+
586+
_(This example is complete, it can be run "as is")_
587+
588+
This is the toolset-level equivalent of the [`SetToolMetadata`][pydantic_ai.capabilities.SetToolMetadata] capability, which applies across all toolsets or a selected subset.
589+
529590
### Changing Tool Execution
530591

531592
[`WrapperToolset`][pydantic_ai.toolsets.WrapperToolset] wraps another toolset and delegates all responsibility to it.

‎mkdocs.yml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,7 @@ nav:
152152
- api/exceptions.md
153153
- api/ext.md
154154
- api/format_prompt.md
155+
- api/function_signature.md
155156
- api/mcp.md
156157
- api/messages.md
157158
- api/models/anthropic.md

‎pydantic_ai_slim/pydantic_ai/__init__.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,9 +140,11 @@
140140
ExternalToolset,
141141
FilteredToolset,
142142
FunctionToolset,
143+
IncludeReturnSchemasToolset,
143144
PrefixedToolset,
144145
PreparedToolset,
145146
RenamedToolset,
147+
SetMetadataToolset,
146148
ToolsetFunc,
147149
ToolsetTool,
148150
WrapperToolset,
@@ -268,9 +270,11 @@
268270
'ExternalToolset',
269271
'FilteredToolset',
270272
'FunctionToolset',
273+
'IncludeReturnSchemasToolset',
271274
'PrefixedToolset',
272275
'PreparedToolset',
273276
'RenamedToolset',
277+
'SetMetadataToolset',
274278
'ToolsetFunc',
275279
'ToolsetTool',
276280
'WrapperToolset',

‎pydantic_ai_slim/pydantic_ai/_agent_graph.py‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,12 +17,12 @@
1717

1818
from pydantic_ai._history_processor import HistoryProcessor
1919
from pydantic_ai._instrumentation import DEFAULT_INSTRUMENTATION_VERSION
20-
from pydantic_ai._tool_manager import ToolManager, ValidatedToolCall
2120
from pydantic_ai._utils import dataclasses_no_defaults_repr, get_union_args, now_utc
2221
from pydantic_ai._uuid import uuid7
2322
from pydantic_ai.builtin_tools import AbstractBuiltinTool
2423
from pydantic_ai.capabilities.abstract import AbstractCapability
2524
from pydantic_ai.models import ModelRequestContext
25+
from pydantic_ai.tool_manager import ToolManager, ValidatedToolCall
2626
from pydantic_graph import BaseNode, GraphRunContext
2727
from pydantic_graph.beta import Graph, GraphBuilder
2828
from pydantic_graph.nodes import End, NodeRunEndT
@@ -1239,6 +1239,7 @@ def build_run_context(ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT
12391239
run_step=ctx.state.run_step,
12401240
run_id=ctx.state.run_id,
12411241
metadata=ctx.state.metadata,
1242+
tool_manager=ctx.deps.tool_manager,
12421243
)
12431244
validation_context = build_validation_context(ctx.deps.validation_context, run_context)
12441245
run_context = replace(run_context, validation_context=validation_context)
@@ -1687,6 +1688,7 @@ async def _call_tool(
16871688
validated = None
16881689
call = tool_call
16891690

1691+
tool_result: Any
16901692
try:
16911693
if tool_call_result is None or isinstance(tool_call_result, ToolApproved):
16921694
if validated is not None:
@@ -1717,14 +1719,16 @@ async def _call_tool(
17171719
return e.tool_retry, None
17181720

17191721
if isinstance(tool_result, _messages.ToolReturn):
1720-
tool_return = tool_result
1721-
elif isinstance(tool_result, list) and any(isinstance(i, _messages.ToolReturn) for i in tool_result): # pyright: ignore[reportUnknownVariableType]
1722+
tool_return = cast(_messages.ToolReturn[Any], tool_result)
1723+
elif isinstance(tool_result, list) and any(
1724+
isinstance(i, _messages.ToolReturn) for i in cast(list[Any], tool_result)
1725+
):
17221726
raise exceptions.UserError(
17231727
f'The return value of tool {call.tool_name!r} contains invalid nested `ToolReturn` objects. '
17241728
f'`ToolReturn` should be used directly.'
17251729
)
17261730
else:
1727-
tool_return = _messages.ToolReturn(return_value=tool_result) # pyright: ignore[reportUnknownArgumentType]
1731+
tool_return = _messages.ToolReturn[Any](return_value=cast(Any, tool_result))
17281732

17291733
return_part = _messages.ToolReturnPart(
17301734
tool_name=call.tool_name,

0 commit comments

Comments
 (0)