This is an automated email from the ASF dual-hosted git repository.
Lee-W pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 1530df852ba Rename `HookToolset` `tool_name_prefix` to `tool_prefix`
(#74447)
1530df852ba is described below
commit 1530df852ba964a20de939933f1e3ae7b42b9fb0
Author: Wei Lee <[email protected]>
AuthorDate: Mon Oct 12 15:38:21 2026 +0900
Rename `HookToolset` `tool_name_prefix` to `tool_prefix` (#74447)
---
providers/common/ai/docs/agent_security.rst | 2 +-
providers/common/ai/docs/stability.rst | 5 +-
providers/common/ai/docs/toolsets/hook.rst | 11 ++--
.../common/ai/example_dags/example_agent.py | 4 +-
.../common/ai/example_dags/example_hook_toolset.py | 2 +-
.../airflow/providers/common/ai/tools/__init__.py | 8 +--
.../airflow/providers/common/ai/toolsets/hook.py | 40 ++++++++++++---
.../unit/common/ai/tools/test__from_toolset.py | 2 +-
.../ai/tests/unit/common/ai/tools/test_tools.py | 5 +-
.../ai/tests/unit/common/ai/toolsets/test_hook.py | 58 ++++++++++++++++++++--
10 files changed, 111 insertions(+), 26 deletions(-)
diff --git a/providers/common/ai/docs/agent_security.rst
b/providers/common/ai/docs/agent_security.rst
index d5e1a07ae22..58a26dc17b9 100644
--- a/providers/common/ai/docs/agent_security.rst
+++ b/providers/common/ai/docs/agent_security.rst
@@ -255,7 +255,7 @@ database user with the minimum privileges required.
HookToolset(
s3_hook,
allowed_methods=["list_keys", "read_key"],
- tool_name_prefix="s3_",
+ tool_prefix="s3",
)
# Bad: exposes delete and write operations
diff --git a/providers/common/ai/docs/stability.rst
b/providers/common/ai/docs/stability.rst
index 8f3682b60b4..15d352b9a15 100644
--- a/providers/common/ai/docs/stability.rst
+++ b/providers/common/ai/docs/stability.rst
@@ -77,8 +77,9 @@ Pydantic AI toolsets, but the Pydantic AI class they inherit
from can change.
statement, the model gets the error and can correct it, up to the
tool's retry
limit, after which the task fails.
* - :class:`~airflow.providers.common.ai.toolsets.hook.HookToolset`
- - Exposes exactly the hook methods in ``allowed_methods``, each named
after its method
- with ``tool_name_prefix`` in front, and raises an error when the
toolset is created
+ - Exposes exactly the hook methods in ``allowed_methods``, each named
+ ``<tool_prefix>_<method>`` when ``tool_prefix`` is set, which must be a
valid Python
+ identifier, and raises an error when the toolset is created
if a listed method does not exist on the hook. ``pinned_arguments`` is
experimental; see below.
* - :class:`~airflow.providers.common.ai.toolsets.mcp.MCPToolset` and
diff --git a/providers/common/ai/docs/toolsets/hook.rst
b/providers/common/ai/docs/toolsets/hook.rst
index f44b72f62cc..4ef43cc701a 100644
--- a/providers/common/ai/docs/toolsets/hook.rst
+++ b/providers/common/ai/docs/toolsets/hook.rst
@@ -32,7 +32,7 @@ list; there is no auto-discovery.
toolset = HookToolset(
http_hook,
allowed_methods=["run"],
- tool_name_prefix="http_",
+ tool_prefix="http",
)
For each listed method, the introspection engine:
@@ -168,8 +168,11 @@ Parameters
- ``hook``: An instantiated Airflow Hook. Its connection ID is templated.
- ``allowed_methods``: Method names to expose as tools. Required. Methods
are validated with ``hasattr`` + ``callable`` at instantiation time.
-- ``tool_name_prefix``: Optional prefix prepended to each tool name
- (e.g. ``"s3_"`` produces ``"s3_list_keys"``).
+- ``tool_prefix``: Optional prefix joined to each tool name with an underscore
+ (e.g. ``"s3"`` produces ``"s3_list_keys"``). Must be a valid Python
identifier.
+ ``tool_name_prefix`` is deprecated and concatenated as given (``"s3_"``
produces
+ ``"s3_list_keys"``); passing both raises ``ValueError`` unless
``tool_name_prefix``
+ equals ``"<tool_prefix>_"``.
- ``pinned_arguments``: Arguments fixed by the Dag author rather than chosen
by the
model. See above.
- ``max_retries``: How many times the model may correct a call with invalid
arguments,
@@ -218,7 +221,7 @@ reflection-based adapter, so the work is choosing the
method list.
HookToolset(
s3_hook,
allowed_methods=["list_keys", "read_key"],
- tool_name_prefix="s3_",
+ tool_prefix="s3",
)
**Credentials and where it runs.** The hook instance is yours, so the
credential
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent.py
index 2eb36275f53..b6bbd55502c 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent.py
@@ -114,7 +114,7 @@ def example_agent_operator_hook():
HookToolset(
http_hook,
allowed_methods=["run"],
- tool_name_prefix="http_",
+ tool_prefix="http",
)
],
)
@@ -150,7 +150,7 @@ def example_agent_self_hosted():
HookToolset(
http_hook,
allowed_methods=["run"],
- tool_name_prefix="http_",
+ tool_prefix="http",
)
],
)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_hook_toolset.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_hook_toolset.py
index 88dc99aa108..adbe6752b7f 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_hook_toolset.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_hook_toolset.py
@@ -45,7 +45,7 @@ if S3Hook is not None:
# Fix the bucket: the model never sees the argument and
cannot set it.
pinned_arguments={"bucket_name": "acme-reports"},
# Name the tools s3_list_keys and s3_read_key.
- tool_name_prefix="s3_",
+ tool_prefix="s3",
# Let the model correct an invalid call up to 2 times.
max_retries=2,
)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/tools/__init__.py
b/providers/common/ai/src/airflow/providers/common/ai/tools/__init__.py
index b6317aeb691..dc7673e88f6 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/tools/__init__.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/tools/__init__.py
@@ -154,7 +154,8 @@ def collect_tools(sources: Iterable[ToolProvider |
AirflowTool]) -> list[Airflow
Flatten toolsets and individual tools, as a framework adapter receives
them, into tools.
:raises ValueError: when two tools share a name. Frameworks route a call
by name, so one
- of them would silently never run; give one toolset a tool name prefix.
+ of them would silently never run; give the toolsets different
``tool_prefix`` values,
+ where they take one.
"""
tools: list[AirflowTool] = []
for source in sources:
@@ -165,8 +166,7 @@ def collect_tools(sources: Iterable[ToolProvider |
AirflowTool]) -> list[Airflow
names = [tool.name for tool in tools]
if duplicates := sorted({name for name in names if names.count(name) > 1}):
raise ValueError(
- f"More than one tool is named {', '.join(duplicates)}. Give the
toolsets different tool "
- "name prefixes where they take one (HookToolset's
tool_name_prefix, the tool_prefix of "
- "SandboxToolset), so the model can call each of them."
+ f"More than one tool is named {', '.join(duplicates)}. "
+ "Give the toolsets different tool_prefix values, where they take
one, so the model can call each of them."
)
return tools
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py
b/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py
index 66814946586..b6402796658 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py
@@ -22,12 +22,14 @@ import copy
import inspect
import re
import types
+import warnings
from typing import TYPE_CHECKING, Any, Union, get_args, get_origin,
get_type_hints
from pydantic_ai.exceptions import ModelRetry
from pydantic_ai.tools import ToolDefinition
from pydantic_ai.toolsets.abstract import ToolsetTool
+from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.common.ai.utils.tool_definition import (
build_args_validator,
return_schema_kwargs,
@@ -70,8 +72,12 @@ class HookToolset(AirflowToolset):
is not modified; each task instance gets a copy.
:param allowed_methods: Method names to expose as tools. Required —
auto-discovery is intentionally not supported for safety.
- :param tool_name_prefix: Optional prefix prepended to each tool name
- (e.g. ``"s3_"`` → ``"s3_list_keys"``).
+ :param tool_prefix: Optional prefix joined to each tool name with an
underscore
+ (e.g. ``"s3"`` → ``"s3_list_keys"``). Must be a valid Python
identifier.
+ :param tool_name_prefix: Deprecated. Concatenated to each tool name as
given
+ (``"s3_"`` → ``"s3_list_keys"``, ``"s3"`` → ``"s3list_keys"``) and not
validated.
+ Passing it with ``tool_prefix`` raises ``ValueError`` unless it equals
+ ``"<tool_prefix>_"``.
:param pinned_arguments: Experimental. Arguments the Dag author fixes,
such as the bucket a
storage hook may use: ``{"bucket_name": "reports"}``. Each is left out
of the
arguments the model sees, refused if the model supplies it anyway, and
passed to
@@ -95,10 +101,32 @@ class HookToolset(AirflowToolset):
hook: BaseHook,
*,
allowed_methods: list[str],
- tool_name_prefix: str = "",
+ tool_prefix: str = "",
+ tool_name_prefix: str | None = None,
pinned_arguments: dict[str, Any] | None = None,
max_retries: int | None = None,
) -> None:
+ if tool_prefix and not tool_prefix.isidentifier():
+ # The prefixed names are rendered as Python function signatures
under
+ # code mode, so a name that is not an identifier breaks there.
+ raise ValueError(f"tool_prefix must be a valid Python identifier,
got {tool_prefix!r}.")
+ if tool_name_prefix is not None:
+ if tool_prefix and f"{tool_prefix}_" != tool_name_prefix:
+ raise ValueError(
+ f"tool_prefix={tool_prefix!r} and
tool_name_prefix={tool_name_prefix!r} conflict. "
+ f"Pass only tool_prefix; tool_name_prefix is its
deprecated alias and must equal "
+ f"{tool_prefix + '_'!r} when both are given."
+ )
+ warnings.warn(
+ f"Parameter `tool_name_prefix` is deprecated. Use "
+ f"`tool_prefix={tool_name_prefix.removesuffix('_')!r}`
instead; "
+ f"HookToolset adds the '_' itself.",
+ category=AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
+ name_prefix = tool_name_prefix
+ else:
+ name_prefix = f"{tool_prefix}_" if tool_prefix else ""
self._max_retries = validate_max_retries(max_retries)
if not allowed_methods:
raise ValueError("allowed_methods must be a non-empty list.")
@@ -135,7 +163,7 @@ class HookToolset(AirflowToolset):
self._hook = hook
self._allowed_methods = allowed_methods
- self._tool_name_prefix = tool_name_prefix
+ self._name_prefix = name_prefix
# The attribute holding the hook's connection ID, e.g.
``postgres_conn_id``. Some hooks
# name one attribute in conn_name_attr but keep the ID in ``conn_id``
(WasbHook,
# KubernetesHook), so fall back to that.
@@ -170,7 +198,7 @@ class HookToolset(AirflowToolset):
tools: dict[str, ToolsetTool[Any]] = {}
for method_name in self._allowed_methods:
method = getattr(self._hook, method_name)
- tool_name = f"{self._tool_name_prefix}{method_name}" if
self._tool_name_prefix else method_name
+ tool_name = f"{self._name_prefix}{method_name}"
json_schema = _build_json_schema_from_signature(method)
description = _extract_description(method)
@@ -228,7 +256,7 @@ class HookToolset(AirflowToolset):
ctx: RunContext[Any],
tool: ToolsetTool[Any],
) -> Any:
- method_name = name.removeprefix(self._tool_name_prefix) if
self._tool_name_prefix else name
+ method_name = name.removeprefix(self._name_prefix)
method: Callable[..., Any] = getattr(self._hook, method_name)
# The framework bridges validate arguments without
args_validator_func, so check again here.
self._refuse_pinned(ctx, **tool_args)
diff --git
a/providers/common/ai/tests/unit/common/ai/tools/test__from_toolset.py
b/providers/common/ai/tests/unit/common/ai/tools/test__from_toolset.py
index e4d8dd0f2e7..76f5030eaec 100644
--- a/providers/common/ai/tests/unit/common/ai/tools/test__from_toolset.py
+++ b/providers/common/ai/tests/unit/common/ai/tools/test__from_toolset.py
@@ -365,7 +365,7 @@ class _FakeHook:
class TestHookToolsetAirflowTools:
def test_calls_the_hook_method(self):
- ts = HookToolset(_FakeHook(), allowed_methods=["list_keys"],
tool_name_prefix="s3_")
+ ts = HookToolset(_FakeHook(), allowed_methods=["list_keys"],
tool_prefix="s3")
tool = _by_name(ts.airflow_tools())["s3_list_keys"]
result = asyncio.run(tool.call({"bucket": "raw", "prefix": "2026/"}))
diff --git a/providers/common/ai/tests/unit/common/ai/tools/test_tools.py
b/providers/common/ai/tests/unit/common/ai/tools/test_tools.py
index b51a7b97bd5..35d1975382c 100644
--- a/providers/common/ai/tests/unit/common/ai/tools/test_tools.py
+++ b/providers/common/ai/tests/unit/common/ai/tools/test_tools.py
@@ -106,7 +106,10 @@ class TestCollectTools:
def test_refuses_two_tools_with_the_same_name(self):
"""Frameworks route a call by name, so one of the two would never
run."""
- with pytest.raises(ValueError, match="More than one tool is named
query"):
+ with pytest.raises(
+ ValueError,
+ match="More than one tool is named query. .*tool_prefix values,
where they take one, so",
+ ):
collect_tools([_Provider(_named("query")),
_Provider(_named("query"))])
diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py
b/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py
index d3154e27bf9..974b6c1d4a9 100644
--- a/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py
+++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py
@@ -17,6 +17,7 @@
from __future__ import annotations
import asyncio
+import contextlib
import re
import threading
from unittest.mock import MagicMock
@@ -28,6 +29,7 @@ from pydantic_ai.messages import ModelResponse,
RetryPromptPart, TextPart, ToolC
from pydantic_ai.models.function import FunctionModel
from pydantic_core import ValidationError
+from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.common.ai.toolsets.hook import (
HookToolset,
_build_json_schema_from_signature,
@@ -166,12 +168,52 @@ class TestHookToolsetGetTools:
# prefix has a default, so it's not required
assert "prefix" not in tool_def.parameters_json_schema.get("required",
[])
- def test_tool_name_prefix(self):
+ def test_tool_prefix(self):
hook = _FakeHook()
- ts = HookToolset(hook, allowed_methods=["list_keys"],
tool_name_prefix="s3_")
+ ts = HookToolset(hook, allowed_methods=["list_keys"], tool_prefix="s3")
tools = asyncio.run(ts.get_tools(ctx=MagicMock()))
assert "s3_list_keys" in tools
+ @pytest.mark.parametrize("tool_prefix", ["my-hook", "1abc", "has space"])
+ def test_tool_prefix_must_be_identifier(self, tool_prefix):
+ with pytest.raises(ValueError, match="tool_prefix must be a valid
Python identifier"):
+ HookToolset(_FakeHook(), allowed_methods=["list_keys"],
tool_prefix=tool_prefix)
+
+ @pytest.mark.parametrize(
+ ("tool_name_prefix", "expected"),
+ [("s3_", "s3_list_keys"), ("s3", "s3list_keys")],
+ )
+ def
test_tool_name_prefix_is_a_deprecated_alias_concatenated_as_given(self,
tool_name_prefix, expected):
+ hook = _FakeHook()
+ with pytest.warns(AirflowProviderDeprecationWarning,
match="tool_name_prefix"):
+ ts = HookToolset(hook, allowed_methods=["list_keys"],
tool_name_prefix=tool_name_prefix)
+ tools = asyncio.run(ts.get_tools(ctx=MagicMock()))
+ assert expected in tools
+
+ def test_tool_name_prefix_warning_suggests_migrated_value(self):
+ with pytest.warns(
+ AirflowProviderDeprecationWarning,
+ match=r"Use `tool_prefix='s3'` instead; HookToolset adds the '_'
itself",
+ ):
+ HookToolset(_FakeHook(), allowed_methods=["list_keys"],
tool_name_prefix="s3_")
+
+ def test_tool_prefix_and_tool_name_prefix_that_agree_still_warn(self):
+ with pytest.warns(AirflowProviderDeprecationWarning,
match="tool_name_prefix"):
+ HookToolset(_FakeHook(), allowed_methods=["list_keys"],
tool_prefix="s3", tool_name_prefix="s3_")
+
+ @pytest.mark.parametrize(
+ ("tool_prefix", "tool_name_prefix"),
+ [("a", "b_"), ("s3", "s3")],
+ )
+ def test_conflicting_prefixes_raise(self, tool_prefix, tool_name_prefix):
+ with pytest.raises(ValueError, match="tool_prefix.*tool_name_prefix"):
+ HookToolset(
+ _FakeHook(),
+ allowed_methods=["list_keys"],
+ tool_prefix=tool_prefix,
+ tool_name_prefix=tool_name_prefix,
+ )
+
def test_description_from_docstring(self):
hook = _FakeHook()
ts = HookToolset(hook, allowed_methods=["list_keys"])
@@ -254,9 +296,17 @@ class TestHookToolsetCallTool:
)
assert "data/file1.txt" in result
- def test_dispatches_with_prefix(self):
+ @pytest.mark.parametrize(
+ ("kwargs", "expected_warning"),
+ [
+ ({"tool_prefix": "storage"}, None),
+ ({"tool_name_prefix": "storage_"},
AirflowProviderDeprecationWarning),
+ ],
+ )
+ def test_dispatches_with_prefix(self, kwargs, expected_warning):
hook = _FakeHook()
- ts = HookToolset(hook, allowed_methods=["read_file"],
tool_name_prefix="storage_")
+ with pytest.warns(expected_warning) if expected_warning else
contextlib.nullcontext():
+ ts = HookToolset(hook, allowed_methods=["read_file"], **kwargs)
tools = asyncio.run(ts.get_tools(ctx=MagicMock()))
result = asyncio.run(
