Skip to content

Add support for specifying tool name when using decorator. #1087

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 10, 2025
Merged

Add support for specifying tool name when using decorator. #1087

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f618b39
Add support for specifying tool name when using decorator.
phemmer Mar 10, 2025
9f4d035
Adjust test
Kludex Mar 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 18 additions & 4 deletions pydantic_ai_slim/pydantic_ai/agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -922,6 +922,7 @@ def tool(
self,
/,
*,
name: str | None = None,
retries: int | None = None,
prepare: ToolPrepareFunc[AgentDepsT] | None = None,
docstring_format: DocstringFormat = 'auto',
Expand All @@ -933,6 +934,7 @@ def tool(
func: ToolFuncContext[AgentDepsT, ToolParams] | None = None,
/,
*,
name: str | None = None,
retries: int | None = None,
prepare: ToolPrepareFunc[AgentDepsT] | None = None,
docstring_format: DocstringFormat = 'auto',
Expand Down Expand Up @@ -969,6 +971,7 @@ async def spam(ctx: RunContext[str], y: float) -> float:

Args:
func: The tool function to register.
name: The name of the tool, defaults to the function name.
retries: The number of retries to allow for this tool, defaults to the agent's default retries,
which defaults to 1.
prepare: custom method to prepare the tool definition for each step, return `None` to omit this
Expand All @@ -984,13 +987,17 @@ def tool_decorator(
func_: ToolFuncContext[AgentDepsT, ToolParams],
) -> ToolFuncContext[AgentDepsT, ToolParams]:
# noinspection PyTypeChecker
self._register_function(func_, True, retries, prepare, docstring_format, require_parameter_descriptions)
self._register_function(
func_, True, name, retries, prepare, docstring_format, require_parameter_descriptions
)
return func_

return tool_decorator
else:
# noinspection PyTypeChecker
self._register_function(func, True, retries, prepare, docstring_format, require_parameter_descriptions)
self._register_function(
func, True, name, retries, prepare, docstring_format, require_parameter_descriptions
)
return func

@overload
Expand All @@ -1001,6 +1008,7 @@ def tool_plain(
self,
/,
*,
name: str | None = None,
retries: int | None = None,
prepare: ToolPrepareFunc[AgentDepsT] | None = None,
docstring_format: DocstringFormat = 'auto',
Expand All @@ -1012,6 +1020,7 @@ def tool_plain(
func: ToolFuncPlain[ToolParams] | None = None,
/,
*,
name: str | None = None,
retries: int | None = None,
prepare: ToolPrepareFunc[AgentDepsT] | None = None,
docstring_format: DocstringFormat = 'auto',
Expand Down Expand Up @@ -1048,6 +1057,7 @@ async def spam(ctx: RunContext[str]) -> float:

Args:
func: The tool function to register.
name: The name of the tool, defaults to the function name.
retries: The number of retries to allow for this tool, defaults to the agent's default retries,
which defaults to 1.
prepare: custom method to prepare the tool definition for each step, return `None` to omit this
Expand All @@ -1062,19 +1072,22 @@ async def spam(ctx: RunContext[str]) -> float:
def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams]:
# noinspection PyTypeChecker
self._register_function(
func_, False, retries, prepare, docstring_format, require_parameter_descriptions
func_, False, name, retries, prepare, docstring_format, require_parameter_descriptions
)
return func_

return tool_decorator
else:
self._register_function(func, False, retries, prepare, docstring_format, require_parameter_descriptions)
self._register_function(
func, False, name, retries, prepare, docstring_format, require_parameter_descriptions
)
return func

def _register_function(
self,
func: ToolFuncEither[AgentDepsT, ToolParams],
takes_ctx: bool,
name: str | None,
retries: int | None,
prepare: ToolPrepareFunc[AgentDepsT] | None,
docstring_format: DocstringFormat,
Expand All @@ -1085,6 +1098,7 @@ def _register_function(
tool = Tool[AgentDepsT](
func,
takes_ctx=takes_ctx,
name=name,
max_retries=retries_,
prepare=prepare,
docstring_format=docstring_format,
Expand Down
22 changes: 22 additions & 0 deletions tests/test_tools.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -500,6 +500,28 @@ def foobar(ctx: RunContext[int], x: int, y: str) -> str:
assert r.data == snapshot('success (no tool calls)')


def test_plain_tool_name():
agent = Agent(FunctionModel(get_json_schema))

def my_tool(arg: str) -> str: ...

agent.tool_plain(name='foo_tool')(my_tool)
result = agent.run_sync('Hello')
json_schema = json.loads(result.data)
assert json_schema['name'] == 'foo_tool'


def test_tool_name():
agent = Agent(FunctionModel(get_json_schema))

def my_tool(ctx: RunContext, arg: str) -> str: ...

agent.tool(name='foo_tool')(my_tool)
result = agent.run_sync('Hello')
json_schema = json.loads(result.data)
assert json_schema['name'] == 'foo_tool'


def test_dynamic_tool_use_messages():
async def repeat_call_foobar(_messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
if info.function_tools:
Expand Down