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 7a41261811b Add toolset as a provider module category (#70122)
7a41261811b is described below
commit 7a41261811bdd8985284d7aa87b0524ee1b323ab
Author: Wei Lee <[email protected]>
AuthorDate: Tue Jul 21 17:59:22 2026 +0800
Add toolset as a provider module category (#70122)
---
airflow-core/src/airflow/provider.yaml.schema.json | 24 +++++++++++++++
dev/registry/registry_tools/types.py | 8 +++++
providers/common/ai/provider.yaml | 11 +++++++
registry/src/_data/types.json | 5 +++
registry/src/css/main.css | 1 +
registry/src/css/tokens.css | 1 +
.../in_container/run_provider_yaml_files_check.py | 36 ++++++++++++++++------
7 files changed, 76 insertions(+), 10 deletions(-)
diff --git a/airflow-core/src/airflow/provider.yaml.schema.json
b/airflow-core/src/airflow/provider.yaml.schema.json
index b9d12bc2975..4f3b53f1077 100644
--- a/airflow-core/src/airflow/provider.yaml.schema.json
+++ b/airflow-core/src/airflow/provider.yaml.schema.json
@@ -373,6 +373,30 @@
]
}
},
+ "toolsets": {
+ "type": "array",
+ "items": {
+ "type": "object",
+ "properties": {
+ "integration-name": {
+ "type": "string",
+ "description": "Integration name. It must have a
matching item in the 'integration' section of any provider."
+ },
+ "python-modules": {
+ "description": "List of python modules containing the
toolsets.",
+ "type": "array",
+ "items": {
+ "type": "string"
+ }
+ }
+ },
+ "additionalProperties": false,
+ "required": [
+ "integration-name",
+ "python-modules"
+ ]
+ }
+ },
"connection-types": {
"type": "array",
"description": "Array of connection types mapped to hook class
names",
diff --git a/dev/registry/registry_tools/types.py
b/dev/registry/registry_tools/types.py
index 7b266d15da2..859e70da1d5 100644
--- a/dev/registry/registry_tools/types.py
+++ b/dev/registry/registry_tools/types.py
@@ -102,6 +102,13 @@ MODULE_TYPES: dict[str, dict] = {
"label": "Decorators",
"icon": "@",
},
+ "toolset": {
+ "yaml_key": "toolsets",
+ "level": "module",
+ "suffixes": ["Toolset"],
+ "label": "Toolsets",
+ "icon": "U",
+ },
}
# Runtime base class imports for issubclass checks (extract_parameters.py).
@@ -113,6 +120,7 @@ BASE_CLASS_IMPORTS: list[tuple[str, str]] = [
("hook", "airflow.sdk.bases.hook.BaseHook"),
("bundle", "airflow.dag_processing.bundles.base.BaseDagBundle"),
("operator", "airflow.sdk.bases.operator.BaseOperator"),
+ ("toolset", "pydantic_ai.toolsets.abstract.AbstractToolset"),
]
# Derived lookups used by extraction scripts.
diff --git a/providers/common/ai/provider.yaml
b/providers/common/ai/provider.yaml
index e5c26e60a55..f7f8ef3f2ea 100644
--- a/providers/common/ai/provider.yaml
+++ b/providers/common/ai/provider.yaml
@@ -451,6 +451,17 @@ operators:
- airflow.providers.common.ai.operators.llamaindex_embedding
- airflow.providers.common.ai.operators.llamaindex_retrieval
+toolsets:
+ - integration-name: Common AI
+ python-modules:
+ - airflow.providers.common.ai.toolsets.hook
+ - airflow.providers.common.ai.toolsets.sql
+ - airflow.providers.common.ai.toolsets.datafusion
+ - airflow.providers.common.ai.toolsets.logging
+ - airflow.providers.common.ai.toolsets.mcp
+ - airflow.providers.common.ai.toolsets.skills
+ - airflow.providers.common.ai.toolsets.langchain_bridge
+
task-decorators:
- class-name: airflow.providers.common.ai.decorators.agent.agent_task
name: agent
diff --git a/registry/src/_data/types.json b/registry/src/_data/types.json
index 000e26a172f..52e5657d499 100644
--- a/registry/src/_data/types.json
+++ b/registry/src/_data/types.json
@@ -53,5 +53,10 @@
"id": "decorator",
"label": "Decorators",
"icon": "@"
+ },
+ {
+ "id": "toolset",
+ "label": "Toolsets",
+ "icon": "U"
}
]
diff --git a/registry/src/css/main.css b/registry/src/css/main.css
index 2f8f4d4cc6a..c733a4dd5f1 100644
--- a/registry/src/css/main.css
+++ b/registry/src/css/main.css
@@ -3065,6 +3065,7 @@ main {
.tab-icon.executor { background: rgb(from var(--color-executor) r g b / 0.2);
color: var(--color-executor); }
.tab-icon.bundle { background: rgb(from var(--color-bundle) r g b / 0.2);
color: var(--color-bundle); }
.tab-icon.decorator { background: rgb(from var(--color-decorator) r g b /
0.2); color: var(--color-decorator); }
+.tab-icon.toolset { background: rgb(from var(--color-toolset) r g b / 0.2);
color: var(--color-toolset); }
/* Modules Layout (sidebar + content) */
.modules-layout {
diff --git a/registry/src/css/tokens.css b/registry/src/css/tokens.css
index dcaa444bf4f..30f885b4927 100644
--- a/registry/src/css/tokens.css
+++ b/registry/src/css/tokens.css
@@ -65,6 +65,7 @@
--color-executor: #f97316;
--color-bundle: #0ea5e9;
--color-decorator: #d946ef;
+ --color-toolset: #84cc16;
/* Additional Colors */
--color-green-400: #4ade80;
diff --git a/scripts/in_container/run_provider_yaml_files_check.py
b/scripts/in_container/run_provider_yaml_files_check.py
index 9518e75ad65..cefdaea8ad5 100755
--- a/scripts/in_container/run_provider_yaml_files_check.py
+++ b/scripts/in_container/run_provider_yaml_files_check.py
@@ -75,6 +75,17 @@ KNOWN_DEPRECATED_CLASSES = [
"airflow.providers.amazon.aws.hooks.kinesis.FirehoseHook",
]
+# AbstractToolset subclasses that are internal implementation detail --
auto-applied by
+# a provider's own runtime rather than constructed by users -- and are
intentionally not
+# part of the registry's public "toolsets" module category, so they are exempt
from
+# check_all_provider_classes_are_registered's registration requirement.
Contrast with
+# e.g. LoggingToolset, which is documented in the toolsets how-to guide and is
registered.
+INTERNAL_UNREGISTERED_TOOLSET_CLASSES = {
+ # Wraps a toolset with per-step result caching for durable execution;
applied
+ # automatically by AgentOperator, not part of the public toolsets how-to
guide.
+ "airflow.providers.common.ai.durable.caching_toolset.CachingToolset",
+}
+
if __name__ != "__main__":
raise SystemExit(
"This file is intended to be executed as an executable program. You
cannot use it as a module."
@@ -368,14 +379,14 @@ def check_integration_duplicates(yaml_files: dict[str,
dict]) -> tuple[int, int]
return num_integrations, num_errors
-@run_check("Checking completeness of list of {sensors, hooks, operators,
triggers, bundles}")
+@run_check("Checking completeness of list of {sensors, hooks, operators,
triggers, bundles, toolsets}")
def check_correctness_of_list_of_sensors_operators_hook_trigger_modules(
yaml_files: dict[str, dict],
) -> tuple[int, int]:
num_errors = 0
num_modules = 0
for (yaml_file_path, provider_data), resource_type in itertools.product(
- yaml_files.items(), ["sensors", "operators", "hooks", "triggers",
"bundles"]
+ yaml_files.items(), ["sensors", "operators", "hooks", "triggers",
"bundles", "toolsets"]
):
expected_modules, provider_package, resource_data = parse_module_data(
provider_data, resource_type, yaml_file_path
@@ -407,14 +418,14 @@ def
check_correctness_of_list_of_sensors_operators_hook_trigger_modules(
return num_modules, num_errors
-@run_check("Checking for duplicates in list of {sensors, hooks, operators,
triggers, bundles}")
+@run_check("Checking for duplicates in list of {sensors, hooks, operators,
triggers, bundles, toolsets}")
def check_duplicates_in_integrations_names_of_hooks_sensors_operators(
yaml_files: dict[str, dict],
) -> tuple[int, int]:
num_errors = 0
num_integrations = 0
for (yaml_file_path, provider_data), resource_type in itertools.product(
- yaml_files.items(), ["sensors", "operators", "hooks", "triggers",
"bundles"]
+ yaml_files.items(), ["sensors", "operators", "hooks", "triggers",
"bundles", "toolsets"]
):
resource_data = provider_data.get(resource_type, [])
count_integrations = Counter(r.get("integration-name", "") for r in
resource_data)
@@ -595,18 +606,20 @@ def
check_hook_classes_with_conn_type_are_registered(yaml_files: dict[str, dict]
@run_check(
- "Checking that all provider Hook/Operator/Sensor/Trigger/Executor/Notifier"
+ "Checking that all provider
Hook/Operator/Sensor/Trigger/Executor/Notifier/Toolset"
" classes are registered in provider.yaml"
)
def check_all_provider_classes_are_registered(yaml_files: dict[str, dict]) ->
tuple[int, int]:
"""
Walk all provider source files, find
Hook/Operator/Sensor/Trigger/Executor/Notifier/
- SecretsBackend/AuthManager/LoggingHandler/DagBundle/DBManager subclasses,
and verify
- they are registered in the appropriate provider.yaml section.
+ SecretsBackend/AuthManager/LoggingHandler/DagBundle/DBManager/Toolset
subclasses, and
+ verify they are registered in the appropriate provider.yaml section.
This catches classes placed in non-standard directories or modules that
were missed
when updating provider.yaml.
"""
+ from pydantic_ai.toolsets.abstract import AbstractToolset
+
from airflow.api_fastapi.auth.managers.base_auth_manager import
BaseAuthManager
from airflow.dag_processing.bundles.base import BaseDagBundle
from airflow.executors.base_executor import BaseExecutor
@@ -632,6 +645,7 @@ def check_all_provider_classes_are_registered(yaml_files:
dict[str, dict]) -> tu
(FileTaskHandler, "logging"),
(BaseDagBundle, "bundles"),
(BaseDBManager, "db-managers"),
+ (AbstractToolset, "toolsets"),
]
# Resource types where registration is by class path (not module)
@@ -656,7 +670,7 @@ def check_all_provider_classes_are_registered(yaml_files:
dict[str, dict]) -> tu
# Collect all modules registered in provider.yaml across all resource
types
registered_modules: set[str] = set()
- for resource_type in ("hooks", "operators", "sensors", "triggers",
"bundles"):
+ for resource_type in ("hooks", "operators", "sensors", "triggers",
"bundles", "toolsets"):
for entry in provider_data.get(resource_type, []):
registered_modules.update(entry.get("python-modules", []))
for entry in provider_data.get("transfers", []):
@@ -733,8 +747,10 @@ def check_all_provider_classes_are_registered(yaml_files:
dict[str, dict]) -> tu
for base_class, resource_type in base_class_resource_map:
if issubclass(obj, base_class) and obj is not base_class:
- num_checks += 1
full_class_name = f"{module_name}.{attr_name}"
+ if full_class_name in
INTERNAL_UNREGISTERED_TOOLSET_CLASSES:
+ break
+ num_checks += 1
# Executors and notifications are registered by class
path;
# other types are registered by module path.
if resource_type in class_level_resource_types:
@@ -889,7 +905,7 @@ def check_invalid_integration(yaml_files: dict[str, dict])
-> tuple[int, int]:
num_errors = 0
num_integrations = len(all_integration_names)
for (yaml_file_path, provider_data), resource_type in itertools.product(
- yaml_files.items(), ["sensors", "operators", "hooks", "triggers",
"bundles"]
+ yaml_files.items(), ["sensors", "operators", "hooks", "triggers",
"bundles", "toolsets"]
):
resource_data = provider_data.get(resource_type, [])
current_names = {r["integration-name"] for r in resource_data}
