This is an automated email from the ASF dual-hosted git repository.
gopidesupavan 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 eddaa67f5bc Add AWS services toolset for agents to access 1000+ APIs
(#70087)
eddaa67f5bc is described below
commit eddaa67f5bc77e8ba63e645ae73948f25bbde8f6
Author: GPK <[email protected]>
AuthorDate: Mon Jul 20 18:42:13 2026 +0100
Add AWS services toolset for agents to access 1000+ APIs (#70087)
* Add AWS services toolset for agents to access 1000+ APIs
* Fix basic tests
* Fix basic tests
* Move to helpers
* resolve comments
* Add more tests
* Update description
* parametrize tests
---
dev/breeze/tests/test_selective_checks.py | 16 +-
providers/common/ai/docs/examples.rst | 4 +-
providers/common/ai/docs/index.rst | 5 +-
providers/common/ai/docs/toolsets.rst | 74 ++-
providers/common/ai/pyproject.toml | 9 +
.../common/ai/example_dags/example_aws_toolset.py | 54 +++
.../providers/common/ai/toolsets/__init__.py | 10 +-
.../airflow/providers/common/ai/toolsets/aws.py | 509 +++++++++++++++++++++
.../ai/tests/unit/common/ai/toolsets/test_aws.py | 359 +++++++++++++++
uv.lock | 12 +-
10 files changed, 1036 insertions(+), 16 deletions(-)
diff --git a/dev/breeze/tests/test_selective_checks.py
b/dev/breeze/tests/test_selective_checks.py
index 158bef04a31..196744ea2cd 100644
--- a/dev/breeze/tests/test_selective_checks.py
+++ b/dev/breeze/tests/test_selective_checks.py
@@ -1109,7 +1109,7 @@ def assert_outputs_are_printed(expected_outputs:
dict[str, str], stderr: str):
pytest.param(
("providers/amazon/src/airflow/providers/amazon/provider.yaml",),
{
- "selected-providers-list-as-string": "amazon apache.hive
cncf.kubernetes "
+ "selected-providers-list-as-string": "amazon apache.hive
cncf.kubernetes common.ai "
"common.compat common.messaging common.sql databricks exasol
ftp google http imap microsoft.azure "
"mongo mysql openlineage postgres salesforce ssh teradata",
"all-python-versions":
f"['{DEFAULT_PYTHON_MAJOR_MINOR_VERSION}']",
@@ -1136,9 +1136,9 @@ def assert_outputs_are_printed(expected_outputs:
dict[str, str], stderr: str):
{
"description": "amazon...google",
"test_types": "Providers[amazon]
Providers[apache.hive,cncf.kubernetes,"
-
"common.compat,common.messaging,common.sql,databricks,exasol,ftp,http,imap,"
-
"microsoft.azure,mongo,mysql,openlineage,postgres,salesforce,ssh,teradata] "
- "Providers[google]",
+
"common.ai,common.compat,common.messaging,common.sql,databricks,exasol,ftp,"
+
"http,imap,microsoft.azure,mongo,mysql,openlineage,postgres,salesforce,ssh,"
+ "teradata] Providers[google]",
}
]
),
@@ -1181,7 +1181,7 @@ def assert_outputs_are_printed(expected_outputs:
dict[str, str], stderr: str):
pytest.param(
("providers/amazon/src/airflow/providers/amazon/file.py",),
{
- "selected-providers-list-as-string": "amazon apache.hive
cncf.kubernetes "
+ "selected-providers-list-as-string": "amazon apache.hive
cncf.kubernetes common.ai "
"common.compat common.messaging common.sql databricks exasol
ftp google http imap microsoft.azure "
"mongo mysql openlineage postgres salesforce ssh teradata",
"all-python-versions":
f"['{DEFAULT_PYTHON_MAJOR_MINOR_VERSION}']",
@@ -1205,9 +1205,9 @@ def assert_outputs_are_printed(expected_outputs:
dict[str, str], stderr: str):
{
"description": "amazon...google",
"test_types": "Providers[amazon]
Providers[apache.hive,cncf.kubernetes,"
-
"common.compat,common.messaging,common.sql,databricks,exasol,ftp,http,imap,"
-
"microsoft.azure,mongo,mysql,openlineage,postgres,salesforce,ssh,teradata] "
- "Providers[google]",
+
"common.ai,common.compat,common.messaging,common.sql,databricks,exasol,ftp,"
+
"http,imap,microsoft.azure,mongo,mysql,openlineage,postgres,salesforce,ssh,"
+ "teradata] Providers[google]",
}
]
),
diff --git a/providers/common/ai/docs/examples.rst
b/providers/common/ai/docs/examples.rst
index a74550d5e60..973a4642d2f 100644
--- a/providers/common/ai/docs/examples.rst
+++ b/providers/common/ai/docs/examples.rst
@@ -69,7 +69,9 @@ Agents & tools
`example_agent_capabilities.py
<https://github.com/apache/airflow/blob/providers-common-ai/|version|/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_capabilities.py>`__).
* - :ref:`Toolsets <howto/toolsets>`
- Loading ``SKILL.md`` Agent Skills
- (`example_agent_skills.py
<https://github.com/apache/airflow/blob/providers-common-ai/|version|/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_skills.py>`__)
+ (`example_agent_skills.py
<https://github.com/apache/airflow/blob/providers-common-ai/|version|/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_skills.py>`__),
+ allow-listed AWS API access via ``AWSToolset``
+ (`example_aws_toolset.py
<https://github.com/apache/airflow/blob/providers-common-ai/|version|/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_aws_toolset.py>`__),
and exposing an Airflow toolset to a LangChain agent, the reverse bridge
(`example_langchain_toolset_bridge.py
<https://github.com/apache/airflow/blob/providers-common-ai/|version|/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_langchain_toolset_bridge.py>`__).
* - :doc:`connections/mcp`
diff --git a/providers/common/ai/docs/index.rst
b/providers/common/ai/docs/index.rst
index 9fce0fd08e0..51081846490 100644
--- a/providers/common/ai/docs/index.rst
+++ b/providers/common/ai/docs/index.rst
@@ -217,12 +217,13 @@ You can install such cross-provider dependencies when
installing from PyPI. For
.. code-block:: bash
- pip install apache-airflow-providers-common-ai[common.sql]
+ pip install apache-airflow-providers-common-ai[amazon]
============================================================================================================
==============
Dependent package
Extra
============================================================================================================
==============
+`apache-airflow-providers-amazon
<https://airflow.apache.org/docs/apache-airflow-providers-amazon>`_
``amazon``
`apache-airflow-providers-common-sql
<https://airflow.apache.org/docs/apache-airflow-providers-common-sql>`_
``common.sql``
`apache-airflow-providers-git
<https://airflow.apache.org/docs/apache-airflow-providers-git>`_
``git``
============================================================================================================
==============
@@ -252,12 +253,14 @@ Extra Dependencies
``avro`` ``fastavro>=1.10.0; python_version < "3.14"``,
``fastavro>=1.12.1; python_version >= "3.14"``
``parquet`` ``pyarrow>=18.0.0; python_version < '3.14'``,
``pyarrow>=22.0.0; python_version >= '3.14'``
``sql`` ``apache-airflow-providers-common-sql``, ``sqlglot>=30.0.0``
+``aws`` ``apache-airflow-providers-amazon>=9.0.0``
``common.sql`` ``apache-airflow-providers-common-sql``
``langchain`` ``langchain>=1.0.0``
``llamaindex`` ``dataclasses-json>=0.6.7``, ``llama-index-core>=0.13.0``,
``llama-index-embeddings-openai>=0.6.0``, ``llama-index-llms-openai>=0.6.0``
``pdf`` ``pypdf>=4.0.0``
``docx`` ``python-docx>=1.0.0``
``git`` ``apache-airflow-providers-git``
+``amazon`` ``apache-airflow-providers-amazon``
==============
=======================================================================================================================================
Downloading official packages
diff --git a/providers/common/ai/docs/toolsets.rst
b/providers/common/ai/docs/toolsets.rst
index 5c46ee45616..8d19e4547b6 100644
--- a/providers/common/ai/docs/toolsets.rst
+++ b/providers/common/ai/docs/toolsets.rst
@@ -24,17 +24,20 @@ Airflow's 350+ provider hooks already have typed methods,
rich docstrings,
and managed credentials. Toolsets expose them as pydantic-ai tools so that
LLM agents can call them during multi-turn reasoning.
-Three toolsets are included:
+Four toolsets are included:
+- :class:`~airflow.providers.common.ai.toolsets.aws.AWSToolset` — configured
+ AWS services toolset for agent access to AWS APIs through Airflow-managed
+ AWS connections.
- :class:`~airflow.providers.common.ai.toolsets.hook.HookToolset` — generic
adapter for any Airflow Hook.
-- :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` — curated
- 4-tool database toolset.
- :class:`~airflow.providers.common.ai.toolsets.mcp.MCPToolset` — connect to
`MCP servers <https://modelcontextprotocol.io/>`__ configured via Airflow
connections.
+- :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` — curated
+ 4-tool database toolset.
-All three implement pydantic-ai's
+All four implement pydantic-ai's
`AbstractToolset <https://ai.pydantic.dev/toolsets/>`__ interface and can be
passed to any pydantic-ai ``Agent``, including via
:class:`~airflow.providers.common.ai.operators.agent.AgentOperator`.
@@ -263,6 +266,69 @@ Parameters
support DDL for in-memory tables; this guard blocks those by default.
- ``max_rows``: Maximum rows returned from the ``query`` tool. Default ``50``.
+``AWSToolset``
+--------------
+
+Curated toolset that gives an agent allow-listed access to AWS APIs with
+three tools:
+
+.. list-table::
+ :header-rows: 1
+ :widths: 20 50
+
+ * - Tool
+ - Description
+ * - ``list_aws_operations``
+ - Lists the operations the toolset allows, grouped by service
+ * - ``describe_aws_operation``
+ - Returns an operation's parameter and response shapes (from the
+ botocore service model), so the agent can check what a call expects
+ before making it
+ * - ``call_aws``
+ - Executes an allowed operation and returns the response as JSON,
+ aggregating paginated results
+
+Credentials, region, and session configuration come from the Airflow
+connection (``aws_conn_id``) via the amazon provider — the model cannot
+override them through tool arguments. Requires the ``aws`` extra::
+
+ pip install "apache-airflow-providers-common-ai[aws]"
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_aws_toolset.py
+ :language: python
+ :start-after: [START howto_operator_agent_aws]
+ :end-before: [END howto_operator_agent_aws]
+
+Access is deny-by-default: ``allowed_actions`` is required and supports ``*``
+wildcards in the operation part only (``"athena:*"``, ``"s3:List*"``).
+Operations that return credentials or decrypted secrets — ``sts:AssumeRole``,
+``secretsmanager:GetSecretValue``, ``kms:Decrypt``, and similar — are never
+matched by a wildcard and must be listed verbatim to be callable.
+
+.. warning::
+ ``allowed_actions`` bounds what the agent can ask for, not what the
+ credentials can do. Point ``aws_conn_id`` at a least-privilege IAM role
+ scoped to the same operations.
+
+Parameters
+^^^^^^^^^^
+
+- ``aws_conn_id``: Airflow connection ID for AWS credentials. Default
+ ``aws_default``.
+- ``allowed_actions``: Operations the agent may call, in
+ ``"<service>:<Operation>"`` form using boto3 service names and API
+ operation names. Required — there is no auto-discovery. Matching is case-
+ and underscore-insensitive, so ``"s3:list_buckets"`` equals
+ ``"s3:ListBuckets"``. Validated against botocore's bundled service
+ definitions at instantiation time (local metadata only — no network, no
+ credentials).
+- ``region_name``: AWS region for API calls. Default ``None`` — use the
+ region configured on the connection.
+- ``max_items``: Upper bound on items aggregated from paginated operations.
+ Default ``1000``.
+- ``max_output_bytes``: Upper bound on the serialized response returned to
+ the agent; larger payloads are clipped and flagged. Default ``65536``.
+
``LoggingToolset``
------------------
diff --git a/providers/common/ai/pyproject.toml
b/providers/common/ai/pyproject.toml
index 53ba8065819..2c1f13eec2a 100644
--- a/providers/common/ai/pyproject.toml
+++ b/providers/common/ai/pyproject.toml
@@ -108,6 +108,11 @@ dependencies = [
"apache-airflow-providers-common-sql",
"sqlglot>=30.0.0",
]
+# AWSToolset: allow-listed AWS API access for agents. The amazon provider
+# supplies credential resolution (AwsBaseHook) and brings boto3/botocore.
+"aws" = [
+ "apache-airflow-providers-amazon>=9.0.0",
+]
"common.sql" = [
"apache-airflow-providers-common-sql"
]
@@ -125,12 +130,16 @@ dependencies = [
"git" = [
"apache-airflow-providers-git"
]
+"amazon" = [
+ "apache-airflow-providers-amazon"
+]
[dependency-groups]
dev = [
"apache-airflow",
"apache-airflow-task-sdk",
"apache-airflow-devel-common",
+ "apache-airflow-providers-amazon",
"apache-airflow-providers-common-compat",
"apache-airflow-providers-common-sql",
"apache-airflow-providers-git",
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_aws_toolset.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_aws_toolset.py
new file mode 100644
index 00000000000..682208447ef
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_aws_toolset.py
@@ -0,0 +1,54 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Example Dag: agent with allow-listed AWS access via AWSToolset."""
+
+from __future__ import annotations
+
+from airflow.providers.common.ai.operators.agent import AgentOperator
+from airflow.providers.common.ai.toolsets.aws import AWSToolset
+from airflow.providers.common.compat.sdk import dag
+
+
+# [START howto_operator_agent_aws]
+@dag(tags=["example"])
+def example_agent_aws_toolset():
+ AgentOperator(
+ task_id="s3_auditor",
+ prompt="Which buckets exist, and roughly how much data is in
'data-lake-raw'?",
+ llm_conn_id="pydanticai_default",
+ system_prompt=(
+ "You are an AWS operations assistant. Discover what you are
allowed "
+ "to call, check parameter shapes before calling, and answer with "
+ "concrete numbers."
+ ),
+ toolsets=[
+ AWSToolset(
+ aws_conn_id="aws_default",
+ allowed_actions=[
+ "s3:ListBuckets",
+ "s3:ListObjectsV2",
+ "s3:GetBucketLocation",
+ ],
+ region_name="us-east-1",
+ )
+ ],
+ )
+
+
+# [END howto_operator_agent_aws]
+
+example_agent_aws_toolset()
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py
b/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py
index 6c30fa4a733..e36690b036d 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/__init__.py
@@ -20,7 +20,7 @@ from __future__ import annotations
from airflow.providers.common.ai.toolsets.hook import HookToolset
-__all__ = ["HookToolset", "MCPToolset", "SQLToolset",
"airflow_toolset_to_langchain_tools"]
+__all__ = ["AWSToolset", "HookToolset", "MCPToolset", "SQLToolset",
"airflow_toolset_to_langchain_tools"]
def __getattr__(name: str):
@@ -46,4 +46,12 @@ def __getattr__(name: str):
raise AirflowOptionalProviderFeatureException(e)
return MCPToolset
+ if name == "AWSToolset":
+ try:
+ from airflow.providers.common.ai.toolsets.aws import AWSToolset
+ except ImportError as e:
+ from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+
+ raise AirflowOptionalProviderFeatureException() from e
+ return AWSToolset
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/toolsets/aws.py
b/providers/common/ai/src/airflow/providers/common/ai/toolsets/aws.py
new file mode 100644
index 00000000000..b37adec7497
--- /dev/null
+++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/aws.py
@@ -0,0 +1,509 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Curated AWS toolset exposing allow-listed boto3 operations as pydantic-ai
tools."""
+
+from __future__ import annotations
+
+import json
+import re
+from fnmatch import fnmatchcase
+from functools import cache
+from typing import TYPE_CHECKING, Any
+
+from pydantic_ai.exceptions import ModelRetry
+from pydantic_ai.tools import ToolDefinition
+from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
+from pydantic_core import SchemaValidator, core_schema
+
+from airflow.providers.common.ai.utils.tool_definition import
return_schema_kwargs
+
+try:
+ import botocore.session
+ from botocore import xform_name
+ from botocore.response import StreamingBody
+
+ from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+except ImportError as e:
+ from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+
+ raise AirflowOptionalProviderFeatureException() from e
+
+if TYPE_CHECKING:
+ from botocore.client import BaseClient
+ from botocore.model import ServiceModel, Shape
+ from pydantic_ai._run_context import RunContext
+
+_PASSTHROUGH_VALIDATOR = SchemaValidator(core_schema.any_schema())
+
+
+def _normalize_action(action: str) -> str:
+ """Comparison key for actions: case- and underscore-insensitive."""
+ return action.replace("_", "").casefold()
+
+
+# High-confidence operations that can put credentials, tokens, passwords,
+# decrypted secret material, or plaintext key material into the agent's
context.
+# A wildcard in ``allowed_actions`` never matches these -- each must be listed
+# verbatim to be callable. This is a defense-in-depth guard, not an exhaustive
+# AWS security boundary; least-privilege IAM on the connection's role remains
+# the real boundary.
+_CREDENTIAL_RETURNING_ACTIONS = frozenset(
+ _normalize_action(action)
+ for action in (
+ "sts:AssumeRole",
+ "sts:AssumeRoleWithSAML",
+ "sts:AssumeRoleWithWebIdentity",
+ "sts:GetFederationToken",
+ "sts:GetSessionToken",
+ "secretsmanager:GetSecretValue",
+ "secretsmanager:BatchGetSecretValue",
+ "ssm:GetParameter",
+ "ssm:GetParameters",
+ "ssm:GetParametersByPath",
+ "kms:Decrypt",
+ "kms:GenerateDataKey",
+ "kms:GenerateDataKeyPair",
+ "ecr:GetAuthorizationToken",
+ "ecr-public:GetAuthorizationToken",
+ "iam:CreateAccessKey",
+ "iam:CreateLoginProfile",
+ "iam:UpdateLoginProfile",
+ "iam:CreateServiceSpecificCredential",
+ "iam:ResetServiceSpecificCredential",
+ "cognito-identity:GetCredentialsForIdentity",
+ "cognito-identity:GetOpenIdToken",
+ "cognito-identity:GetOpenIdTokenForDeveloperIdentity",
+ "redshift:GetClusterCredentials",
+ "redshift:GetClusterCredentialsWithIAM",
+ "redshift-serverless:GetCredentials",
+ "lightsail:GetInstanceAccessDetails",
+ "lightsail:GetRelationalDatabaseMasterUserPassword",
+ )
+)
+
+_SHAPE_DEPTH = 3
+
+# JSON Schemas for the three AWS tools.
+_LIST_OPERATIONS_SCHEMA: dict[str, Any] = {
+ "type": "object",
+ "properties": {
+ "service": {
+ "type": "string",
+ "description": "Optional service name filter, e.g. 's3'.",
+ },
+ },
+}
+
+_DESCRIBE_OPERATION_SCHEMA: dict[str, Any] = {
+ "type": "object",
+ "properties": {
+ "service": {"type": "string", "description": "AWS service name, e.g.
's3'."},
+ "operation": {
+ "type": "string",
+ "description": "API operation name, e.g. 'ListObjectsV2'.",
+ },
+ },
+ "required": ["service", "operation"],
+}
+
+_CALL_AWS_SCHEMA: dict[str, Any] = {
+ "type": "object",
+ "properties": {
+ "service": {"type": "string", "description": "AWS service name, e.g.
's3'."},
+ "operation": {
+ "type": "string",
+ "description": "API operation name, e.g. 'ListObjectsV2'.",
+ },
+ "parameters": {
+ "type": "object",
+ "description": (
+ "Operation parameters exactly as the API expects them (see
describe_aws_operation)."
+ ),
+ },
+ },
+ "required": ["service", "operation"],
+}
+
+
+class AWSToolset(AbstractToolset[Any]):
+ """
+ Curated toolset that gives an LLM agent allow-listed access to AWS APIs.
+
+ Exposes three tools:
+
+ - ``list_aws_operations`` -- the operations this toolset allows, grouped by
+ service (the allow-list intersected with the botocore service model).
+ - ``describe_aws_operation`` -- an operation's parameter and response
shapes,
+ derived from the botocore service model, so the agent can check what a
+ call expects *before* making it.
+ - ``call_aws`` -- execute an allowed operation and return the response as
+ JSON (paginated operations are aggregated up to ``max_items``).
+
+ Credentials, region, and session configuration come exclusively from the
+ Airflow connection (resolved lazily through the amazon provider's
+ ``AwsBaseHook``); none of them are tool arguments, so the model cannot
+ steer them.
+
+ When a call fails, AWS's own error message is returned to the agent as a
+ retry (:class:`pydantic_ai.ModelRetry`) so the model can correct its
+ parameters within the run. pydantic-ai bounds this by the tool's
+ ``max_retries``, so an unrecoverable error -- bad credentials, missing IAM
+ permissions -- exhausts the retries and fails the task for Airflow to
retry.
+
+ :param aws_conn_id: Airflow connection ID for AWS credentials.
+ :param allowed_actions: Operations the agent may call, in
+ ``"<service>:<Operation>"`` form using boto3 service names and API
+ operation names -- e.g. ``["s3:ListBuckets", "s3:GetObject",
+ "athena:*"]``. Required -- access is deny-by-default and there is
+ deliberately no auto-discovery. The operation part accepts ``*``/``?``
+ wildcards; the service part must be explicit. Matching is case- and
+ underscore-insensitive, so ``"s3:list_buckets"`` and
+ ``"s3:ListBuckets"`` are equivalent.
+
+ Operations that return credentials or decrypted secrets (for example
+ ``sts:AssumeRole``, ``secretsmanager:GetSecretValue``, ``kms:Decrypt``)
+ are never matched by a wildcard -- each must be listed verbatim to be
+ callable.
+
+ .. note::
+ This is an application-level guardrail: it bounds what the agent
+ can *ask for*, not what the credentials can *do*. Pair it with a
+ least-privilege IAM role on ``aws_conn_id`` for a hard guarantee.
+
+ :param region_name: AWS region for API calls. ``None`` (default) uses the
+ region configured on the connection. Deliberately not exposed as a tool
+ argument.
+ :param max_items: Upper bound on items aggregated from paginated operations
+ (applied via the paginator's ``MaxItems``). Default ``1000``.
+ :param max_output_bytes: Upper bound on the serialized response returned to
+ the agent. Larger payloads are clipped and flagged with
+ ``"truncated": true``. Default ``65536``.
+ """
+
+ def __init__(
+ self,
+ aws_conn_id: str = "aws_default",
+ *,
+ allowed_actions: list[str],
+ region_name: str | None = None,
+ max_items: int = 1000,
+ max_output_bytes: int = 65536,
+ ) -> None:
+ if not allowed_actions:
+ raise ValueError("allowed_actions must be a non-empty list.")
+
+ # Validation below uses botocore's bundled service definitions only --
+ # local file reads, no credentials, no network -- so it is safe to run
+ # at Dag parse time.
+ patterns: dict[str, list[str]] = {}
+ literal_actions: set[str] = set()
+ for action in allowed_actions:
+ service, sep, operation = action.partition(":")
+ if not sep or not service or not operation:
+ raise ValueError(
+ f"Invalid action {action!r}: expected
'<service>:<Operation>', e.g. 's3:ListBuckets'."
+ )
+ if any(ch in service for ch in "*?["):
+ raise ValueError(f"Invalid action {action!r}: wildcards are
not allowed in the service part.")
+ if service not in _get_available_services():
+ raise ValueError(
+ f"Unknown AWS service {service!r} in action {action!r}. "
+ "Use the boto3 service name, e.g. 's3' or
'secretsmanager'."
+ )
+ key = _normalize_action(action)
+ if any(ch in operation for ch in "*?["):
+ model = _load_service_model(service)
+ matched_operations = [
+ name
+ for name in model.operation_names
+ if self._is_pattern_match_allowed(key, service, name)
+ ]
+ if not matched_operations:
+ raise ValueError(
+ f"Action pattern {action!r} does not match any
non-sensitive operation "
+ f"for service {service!r}."
+ )
+ patterns.setdefault(service, []).append(key)
+ else:
+ if _resolve_operation_name(_load_service_model(service),
operation) is None:
+ raise ValueError(
+ f"Unknown operation {operation!r} for service
{service!r} in action "
+ f"{action!r}. Use the API operation name, e.g.
'ListBuckets'."
+ )
+ patterns.setdefault(service, []).append(key)
+ literal_actions.add(key)
+
+ self._aws_conn_id = aws_conn_id
+ self._region_name = region_name
+ self._max_items = max_items
+ self._max_output_bytes = max_output_bytes
+ self._patterns = patterns
+ self._literal_actions = frozenset(literal_actions)
+ # Populated lazily at call time, so Dag parsing never creates AWS
clients.
+ self._clients: dict[str, BaseClient] = {}
+
+ @property
+ def id(self) -> str:
+ return f"aws-{self._aws_conn_id}"
+
+ # ------------------------------------------------------------------
+ # AbstractToolset interface
+ # ------------------------------------------------------------------
+
+ async def get_tools(self, ctx: RunContext[Any]) -> dict[str,
ToolsetTool[Any]]:
+ tools: dict[str, ToolsetTool[Any]] = {}
+
+ for name, description, schema in (
+ (
+ "list_aws_operations",
+ "List the AWS operations this toolset allows, grouped by
service.",
+ _LIST_OPERATIONS_SCHEMA,
+ ),
+ (
+ "describe_aws_operation",
+ "Show an operation's parameters and response shape before
calling it.",
+ _DESCRIBE_OPERATION_SCHEMA,
+ ),
+ (
+ "call_aws",
+ "Execute an allowed AWS operation and return the response as
JSON.",
+ _CALL_AWS_SCHEMA,
+ ),
+ ):
+ # sequential=True because all tools share per-service boto3 clients
+ # with synchronous I/O -- they must not run concurrently.
+ # return_schema is "string": every tool returns a JSON-encoded
string,
+ # so code mode renders `-> str` instead of `-> Any`.
+ tool_def = ToolDefinition(
+ name=name,
+ description=description,
+ parameters_json_schema=schema,
+ sequential=True,
+ **return_schema_kwargs({"type": "string"}),
+ )
+ tools[name] = ToolsetTool(
+ toolset=self,
+ tool_def=tool_def,
+ max_retries=1,
+ args_validator=_PASSTHROUGH_VALIDATOR,
+ )
+ return tools
+
+ async def call_tool(
+ self,
+ name: str,
+ tool_args: dict[str, Any],
+ ctx: RunContext[Any],
+ tool: ToolsetTool[Any],
+ ) -> Any:
+ if name not in ("list_aws_operations", "describe_aws_operation",
"call_aws"):
+ raise ValueError(f"Unknown tool: {name!r}")
+ try:
+ if name == "list_aws_operations":
+ return self._list_operations(tool_args.get("service"))
+ if name == "describe_aws_operation":
+ return self._describe_operation(tool_args["service"],
tool_args["operation"])
+ return self._call_aws(
+ tool_args["service"], tool_args["operation"],
tool_args.get("parameters") or {}
+ )
+ except Exception as e:
+ # Hand AWS's own error back to the agent as a retry so it can
correct
+ # its parameters within the run. pydantic-ai bounds this by the
tool's
+ # max_retries, so an unrecoverable error (bad credentials, missing
IAM
+ # permissions) exhausts the budget and fails the task for Airflow
to
+ # retry, rather than being silently worked around.
+ raise ModelRetry(
+ f"The {name} tool failed: {e}\n"
+ "Use list_aws_operations to see what you may call and "
+ "describe_aws_operation to check the expected parameters, then
try again."
+ ) from e
+
+ # ------------------------------------------------------------------
+ # Authorization
+ # ------------------------------------------------------------------
+
+ def _is_action_allowed(self, *, service: str, operation: str) -> bool:
+ key = _normalize_action(f"{service}:{operation}")
+ if key in _CREDENTIAL_RETURNING_ACTIONS:
+ return key in self._literal_actions
+ return any(fnmatchcase(key, pattern) for pattern in
self._patterns.get(service, ()))
+
+ def _is_pattern_match_allowed(self, pattern: str, service: str, operation:
str) -> bool:
+ key = _normalize_action(f"{service}:{operation}")
+ return key not in _CREDENTIAL_RETURNING_ACTIONS and fnmatchcase(key,
pattern)
+
+ def _resolve_allowed_service(self, service: str) -> str:
+ """Validate a tool-supplied service name against the allow-list."""
+ if service not in self._patterns:
+ raise ValueError(
+ f"Service {service!r} is not in this toolset's allowed
actions. "
+ f"Allowed services: {', '.join(sorted(self._patterns))}."
+ )
+ return service
+
+ def _resolve_allowed_operation(self, service: str, operation: str) ->
tuple[str, ServiceModel]:
+ model = _load_service_model(self._resolve_allowed_service(service))
+ canonical = _resolve_operation_name(model, operation)
+ if canonical is None:
+ raise ValueError(f"Unknown operation {operation!r} for service
{service!r}.")
+ if not self._is_action_allowed(service=service, operation=canonical):
+ raise ValueError(f"Operation {service}:{canonical} is not in this
toolset's allowed actions.")
+ return canonical, model
+
+ # ------------------------------------------------------------------
+ # Tool implementations
+ # ------------------------------------------------------------------
+
+ def _list_operations(self, service: str | None) -> str:
+ services = [self._resolve_allowed_service(service)] if service else
sorted(self._patterns)
+ listing = {
+ svc: sorted(
+ name
+ for name in _load_service_model(svc).operation_names
+ if self._is_action_allowed(service=svc, operation=name)
+ )
+ for svc in services
+ }
+ return json.dumps(listing)
+
+ def _describe_operation(self, service: str, operation: str) -> str:
+ canonical, model = self._resolve_allowed_operation(service, operation)
+ op = model.operation_model(canonical)
+ return json.dumps(
+ {
+ "service": service,
+ "operation": canonical,
+ "documentation": _strip_html(op.documentation)[:600],
+ "input": _describe_shape(op.input_shape),
+ "output": _describe_shape(op.output_shape),
+ }
+ )
+
+ def _call_aws(self, service: str, operation: str, parameters: dict[str,
Any]) -> str:
+ canonical, _ = self._resolve_allowed_operation(service, operation)
+ client = self._get_client(service)
+ method_name = xform_name(canonical)
+ if client.can_paginate(method_name):
+ pages = client.get_paginator(method_name).paginate(
+ **parameters, PaginationConfig={"MaxItems": self._max_items}
+ )
+ response = pages.build_full_result()
+ else:
+ response = getattr(client, method_name)(**parameters)
+ return self._serialize_response(response)
+
+ # ------------------------------------------------------------------
+ # Lazy client resolution and serialization
+ # ------------------------------------------------------------------
+
+ def _get_client(self, service: str) -> BaseClient:
+ if service not in self._clients:
+ hook = AwsBaseHook(
+ aws_conn_id=self._aws_conn_id, client_type=service,
region_name=self._region_name
+ )
+ self._clients[service] = hook.get_conn()
+ return self._clients[service]
+
+ def _serialize_response(self, response: Any) -> str:
+ if isinstance(response, dict):
+ response.pop("ResponseMetadata", None)
+ payload = json.dumps(_to_jsonable(response,
max_bytes=self._max_output_bytes), default=str)
+ if len(payload) > self._max_output_bytes:
+ return json.dumps(
+ {
+ "truncated": True,
+ "max_output_bytes": self._max_output_bytes,
+ "data": payload[: self._max_output_bytes],
+ }
+ )
+ return payload
+
+
+# ---------------------------------------------------------------------------
+# Private botocore helpers
+# ---------------------------------------------------------------------------
+
+
+@cache
+def _get_available_services() -> frozenset[str]:
+ return frozenset(botocore.session.get_session().get_available_services())
+
+
+@cache
+def _load_service_model(service: str) -> ServiceModel:
+ return botocore.session.get_session().get_service_model(service)
+
+
+def _resolve_operation_name(model: ServiceModel, operation: str) -> str | None:
+ """Resolve a case/underscore-insensitive operation name to its canonical
API name."""
+ wanted = _normalize_action(operation)
+ for name in model.operation_names:
+ if _normalize_action(name) == wanted:
+ return name
+ return None
+
+
+def _describe_shape(shape: Shape | None, depth: int = _SHAPE_DEPTH) -> Any:
+ """Render a botocore shape as a compact JSON-friendly summary."""
+ if shape is None:
+ return None
+ type_name = shape.type_name
+ if type_name == "structure":
+ if depth <= 0:
+ return type_name
+ required = set(shape.required_members)
+ return {
+ name: (
+ {"type": _describe_shape(member, depth - 1), "required": True}
+ if name in required
+ else {"type": _describe_shape(member, depth - 1)}
+ )
+ for name, member in shape.members.items()
+ }
+ if type_name == "list":
+ return [_describe_shape(shape.member, depth - 1)] if depth > 0 else
type_name
+ if type_name == "map":
+ return {"<key>": _describe_shape(shape.value, depth - 1)} if depth > 0
else type_name
+ enum = getattr(shape, "enum", None)
+ if enum:
+ return f"{type_name} (one of: {', '.join(enum)})"
+ return type_name
+
+
+_HTML_TAG_RE = re.compile(r"<[^>]+>")
+
+
+def _strip_html(text: str | None) -> str:
+ return " ".join(_HTML_TAG_RE.sub(" ", text or "").split())
+
+
+def _to_jsonable(value: Any, *, max_bytes: int) -> Any:
+ """Make a boto3 response JSON-serializable; datetimes are handled by
``default=str``."""
+ if isinstance(value, StreamingBody):
+ # Read one byte past the cap so truncation is detectable.
+ raw = value.read(max_bytes + 1)
+ text = raw[:max_bytes].decode("utf-8", errors="replace")
+ if len(raw) > max_bytes:
+ return {"truncated": True, "data": text}
+ return text
+ if isinstance(value, bytes):
+ return value.decode("utf-8", errors="replace")
+ if isinstance(value, dict):
+ return {k: _to_jsonable(v, max_bytes=max_bytes) for k, v in
value.items()}
+ if isinstance(value, list):
+ return [_to_jsonable(v, max_bytes=max_bytes) for v in value]
+ return value
diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_aws.py
b/providers/common/ai/tests/unit/common/ai/toolsets/test_aws.py
new file mode 100644
index 00000000000..08d29f12b90
--- /dev/null
+++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_aws.py
@@ -0,0 +1,359 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import asyncio
+import json
+from datetime import datetime, timezone
+from unittest.mock import MagicMock, patch
+
+import pytest
+from botocore.exceptions import ClientError
+from pydantic_ai.exceptions import ModelRetry
+
+from airflow.providers.common.ai.toolsets.aws import (
+ AWSToolset,
+ _get_available_services,
+ _load_service_model,
+ _resolve_operation_name,
+)
+from airflow.providers.common.ai.utils.tool_definition import
_SUPPORTS_RETURN_SCHEMA
+
+
+def _make_toolset(**kwargs):
+ kwargs.setdefault("allowed_actions", ["s3:List*", "s3:GetBucketLocation"])
+ return AWSToolset("aws_test", **kwargs)
+
+
+def _make_mock_client():
+ client = MagicMock()
+ client.can_paginate.return_value = False
+ return client
+
+
+def _call(ts, name, args):
+ return asyncio.run(ts.call_tool(name, args, ctx=MagicMock(),
tool=MagicMock()))
+
+
+class TestAWSToolsetInit:
+ def test_id_includes_conn_id(self):
+ ts = _make_toolset()
+ assert ts.id == "aws-aws_test"
+
+ def test_rejects_empty_allowed_actions(self):
+ with pytest.raises(ValueError, match="non-empty"):
+ AWSToolset("aws_test", allowed_actions=[])
+
+ @pytest.mark.parametrize(
+ ("action", "match"),
+ [
+ ("s3", "expected"),
+ ("s3:", "expected"),
+ (":ListBuckets", "expected"),
+ ("*:GetObject", "wildcards are not allowed in the service part"),
+ ("nosuchservice:DoThing", "Unknown AWS service"),
+ ("s3:NoSuchOperation", "Unknown operation"),
+ ],
+ )
+ def test_rejects_invalid_action(self, action, match):
+ with pytest.raises(ValueError, match=match):
+ AWSToolset("aws_test", allowed_actions=[action])
+
+ def test_accepts_snake_case_operation(self):
+ AWSToolset("aws_test", allowed_actions=["s3:list_buckets"])
+
+ def test_rejects_wildcard_action_that_matches_no_operations(self):
+ with pytest.raises(ValueError, match="does not match any"):
+ AWSToolset("aws_test", allowed_actions=["s3:ListObjectz*"])
+
+ def
test_rejects_wildcard_action_that_only_matches_sensitive_operations(self):
+ with pytest.raises(ValueError, match="does not match any"):
+ AWSToolset("aws_test", allowed_actions=["kms:Decrypt*"])
+
+ @patch("airflow.providers.common.ai.toolsets.aws._get_available_services",
autospec=True)
+ def test_rejects_service_missing_from_available_services(self,
mock_get_available_services):
+ mock_get_available_services.return_value = frozenset({"ec2"})
+
+ with pytest.raises(ValueError, match="Unknown AWS service 's3'"):
+ AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"])
+
+
+class TestAWSToolsetGetTools:
+ def test_returns_three_tools(self):
+ ts = _make_toolset()
+ tools = asyncio.run(ts.get_tools(ctx=MagicMock()))
+ assert set(tools.keys()) == {"list_aws_operations",
"describe_aws_operation", "call_aws"}
+
+ def test_tool_definitions_have_descriptions(self):
+ ts = _make_toolset()
+ tools = asyncio.run(ts.get_tools(ctx=MagicMock()))
+ for tool in tools.values():
+ assert tool.tool_def.description
+
+ @pytest.mark.skipif(
+ not _SUPPORTS_RETURN_SCHEMA, reason="pydantic-ai too old for
ToolDefinition.return_schema"
+ )
+ def test_tools_declare_string_return_schema(self):
+ ts = _make_toolset()
+ tools = asyncio.run(ts.get_tools(ctx=MagicMock()))
+ for tool in tools.values():
+ assert tool.tool_def.return_schema == {"type": "string"}
+
+
+class TestAWSToolsetActionMatching:
+ @pytest.mark.parametrize(
+ ("actions", "service", "operation", "allowed"),
+ [
+ (["s3:ListBuckets"], "s3", "ListBuckets", True),
+ (["s3:list_buckets"], "s3", "ListBuckets", True),
+ (["s3:List*"], "s3", "ListObjectsV2", True),
+ (["s3:List*"], "s3", "GetObject", False),
+ (["s3:*"], "ec2", "DescribeInstances", False),
+ (["secretsmanager:*"], "secretsmanager", "GetSecretValue", False),
+ (["secretsmanager:GetSecretValue"], "secretsmanager",
"GetSecretValue", True),
+ (["kms:*"], "kms", "Decrypt", False),
+ (["kms:Decrypt"], "kms", "Decrypt", True),
+ ],
+ )
+ def test_allow_list_matching(self, actions, service, operation, allowed):
+ ts = AWSToolset("aws_test", allowed_actions=actions)
+ assert ts._is_action_allowed(service=service, operation=operation) is
allowed
+
+
+class TestAWSToolsetPrivateHelpers:
+ def test_is_action_allowed_requires_keyword_arguments(self):
+ ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"])
+
+ with pytest.raises(TypeError, match="positional"):
+ ts._is_action_allowed("s3", "ListBuckets") # type: ignore[misc]
+
+ def
test_pattern_match_allows_normal_operations_but_not_sensitive_operations(self):
+ ts = AWSToolset("aws_test", allowed_actions=["s3:List*", "kms:List*"])
+
+ assert ts._is_pattern_match_allowed("s3:list*", "s3", "ListBuckets")
is True
+ assert ts._is_pattern_match_allowed("kms:decrypt*", "kms", "Decrypt")
is False
+
+ def test_resolve_allowed_service_returns_configured_service(self):
+ ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"])
+
+ assert ts._resolve_allowed_service("s3") == "s3"
+
+ def test_resolve_allowed_service_rejects_service_outside_allow_list(self):
+ ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"])
+
+ with pytest.raises(ValueError, match="Service 'ec2' is not in this
toolset's allowed actions"):
+ ts._resolve_allowed_service("ec2")
+
+ def test_resolve_allowed_operation_returns_canonical_name_and_model(self):
+ ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"])
+
+ operation, model = ts._resolve_allowed_operation("s3", "list_buckets")
+
+ assert operation == "ListBuckets"
+ assert "ListBuckets" in model.operation_names
+
+ def test_resolve_allowed_operation_rejects_unknown_operation(self):
+ ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"])
+
+ with pytest.raises(ValueError, match="Unknown operation
'NoSuchOperation' for service 's3'"):
+ ts._resolve_allowed_operation("s3", "NoSuchOperation")
+
+ def test_resolve_allowed_operation_rejects_disallowed_operation(self):
+ ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"])
+
+ with pytest.raises(
+ ValueError, match="Operation s3:GetObject is not in this toolset's
allowed actions"
+ ):
+ ts._resolve_allowed_operation("s3", "GetObject")
+
+ def
test_resolve_allowed_operation_rejects_service_outside_allow_list(self):
+ ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"])
+
+ with pytest.raises(ValueError, match="Service 'ec2' is not in this
toolset's allowed actions"):
+ ts._resolve_allowed_operation("ec2", "DescribeInstances")
+
+ def test_get_available_services_reads_botocore_catalog(self):
+ assert "s3" in _get_available_services()
+
+ def test_load_service_model_and_resolve_operation_name(self):
+ model = _load_service_model("s3")
+
+ assert _resolve_operation_name(model, "list_buckets") == "ListBuckets"
+ assert _resolve_operation_name(model, "NoSuchOperation") is None
+
+
+class TestAWSToolsetListOperations:
+ def test_lists_only_allowed_operations(self):
+ ts = AWSToolset("aws_test", allowed_actions=["s3:List*"])
+ listing = json.loads(_call(ts, "list_aws_operations", {}))
+ assert "ListBuckets" in listing["s3"]
+ assert "GetObject" not in listing["s3"]
+
+ def test_sensitive_operations_hidden_behind_wildcard(self):
+ ts = AWSToolset("aws_test", allowed_actions=["secretsmanager:*"])
+ listing = json.loads(_call(ts, "list_aws_operations", {}))
+ assert "ListSecrets" in listing["secretsmanager"]
+ assert "GetSecretValue" not in listing["secretsmanager"]
+
+ def test_verbatim_sensitive_operation_is_listed(self):
+ ts = AWSToolset("aws_test", allowed_actions=["secretsmanager:*",
"secretsmanager:GetSecretValue"])
+ listing = json.loads(_call(ts, "list_aws_operations", {}))
+ assert "GetSecretValue" in listing["secretsmanager"]
+
+ def test_service_outside_allow_list_raises_model_retry(self):
+ ts = _make_toolset()
+ with pytest.raises(ModelRetry, match="not in this toolset's allowed
actions"):
+ _call(ts, "list_aws_operations", {"service": "ec2"})
+
+
+class TestAWSToolsetDescribeOperation:
+ def test_returns_input_shape_with_required_members(self):
+ ts = AWSToolset("aws_test", allowed_actions=["s3:ListObjectsV2"])
+ described = json.loads(
+ _call(ts, "describe_aws_operation", {"service": "s3", "operation":
"ListObjectsV2"})
+ )
+ assert described["operation"] == "ListObjectsV2"
+ assert described["input"]["Bucket"]["required"] is True
+ assert "Contents" in described["output"]
+
+ def test_resolves_snake_case_operation_name(self):
+ ts = AWSToolset("aws_test", allowed_actions=["s3:ListObjectsV2"])
+ described = json.loads(
+ _call(ts, "describe_aws_operation", {"service": "s3", "operation":
"list_objects_v2"})
+ )
+ assert described["operation"] == "ListObjectsV2"
+
+ @pytest.mark.parametrize(
+ ("service", "operation", "match"),
+ [
+ ("s3", "GetObject", "not in this toolset's allowed actions"),
+ ("ec2", "DescribeInstances", "Service 'ec2' is not in this
toolset's allowed actions"),
+ ("s3", "NoSuchOperation", "Unknown operation 'NoSuchOperation' for
service 's3'"),
+ ],
+ ids=["disallowed_operation", "service_outside_allow_list",
"unknown_operation"],
+ )
+ def test_model_supplied_invalid_operation_raises_model_retry(self,
service, operation, match):
+ ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"])
+
+ with pytest.raises(ModelRetry, match=match):
+ _call(ts, "describe_aws_operation", {"service": service,
"operation": operation})
+
+
+class TestAWSToolsetCallAws:
+ def test_executes_operation_and_returns_json(self):
+ ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"])
+ client = _make_mock_client()
+ client.list_buckets.return_value = {
+ "Buckets": [
+ {
+ "Name": "data-lake-raw",
+ "CreationDate": datetime(2026, 1, 1, tzinfo=timezone.utc),
+ }
+ ],
+ "ResponseMetadata": {"RequestId": "abc123"},
+ }
+ ts._clients["s3"] = client
+
+ result = json.loads(_call(ts, "call_aws", {"service": "s3",
"operation": "ListBuckets"}))
+ client.list_buckets.assert_called_once_with()
+ assert result["Buckets"][0]["Name"] == "data-lake-raw"
+ # datetime serialized via default=str, ResponseMetadata stripped.
+ assert "2026-01-01" in result["Buckets"][0]["CreationDate"]
+ assert "ResponseMetadata" not in result
+
+ def test_snake_case_operation_resolves(self):
+ ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"])
+ client = _make_mock_client()
+ client.list_buckets.return_value = {"Buckets": []}
+ ts._clients["s3"] = client
+
+ result = json.loads(_call(ts, "call_aws", {"service": "s3",
"operation": "list_buckets"}))
+ assert result == {"Buckets": []}
+
+ def test_uses_paginator_when_available(self):
+ ts = AWSToolset("aws_test", allowed_actions=["s3:ListObjectsV2"],
max_items=250)
+ client = _make_mock_client()
+ client.can_paginate.return_value = True
+ paginator = client.get_paginator.return_value
+ paginator.paginate.return_value.build_full_result.return_value =
{"Contents": [{"Key": "a.parquet"}]}
+ ts._clients["s3"] = client
+
+ result = json.loads(
+ _call(
+ ts,
+ "call_aws",
+ {"service": "s3", "operation": "ListObjectsV2", "parameters":
{"Bucket": "b"}},
+ )
+ )
+ paginator.paginate.assert_called_once_with(Bucket="b",
PaginationConfig={"MaxItems": 250})
+ assert result["Contents"] == [{"Key": "a.parquet"}]
+
+ @pytest.mark.parametrize(
+ ("service", "operation", "match"),
+ [
+ ("s3", "DeleteBucket", "not in this toolset's allowed actions"),
+ ("ec2", "DescribeInstances", "Service 'ec2' is not in this
toolset's allowed actions"),
+ ("s3", "NoSuchOperation", "Unknown operation 'NoSuchOperation' for
service 's3'"),
+ ],
+ ids=["disallowed_operation", "service_outside_allow_list",
"unknown_operation"],
+ )
+ def
test_model_supplied_invalid_operation_raises_model_retry_without_client(
+ self, service, operation, match
+ ):
+ ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"])
+
+ with pytest.raises(ModelRetry, match=match):
+ _call(ts, "call_aws", {"service": service, "operation": operation})
+ assert ts._clients == {}
+
+ def test_client_error_raises_model_retry(self):
+ ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"])
+ client = _make_mock_client()
+ client.list_buckets.side_effect = ClientError(
+ {"Error": {"Code": "AccessDenied", "Message": "not authorized"}},
"ListBuckets"
+ )
+ ts._clients["s3"] = client
+
+ with pytest.raises(ModelRetry, match="AccessDenied"):
+ _call(ts, "call_aws", {"service": "s3", "operation":
"ListBuckets"})
+
+ def test_large_response_is_truncated(self):
+ ts = AWSToolset("aws_test", allowed_actions=["s3:ListBuckets"],
max_output_bytes=200)
+ client = _make_mock_client()
+ client.list_buckets.return_value = {"Buckets": [{"Name": "x" * 50}] *
50}
+ ts._clients["s3"] = client
+
+ result = json.loads(_call(ts, "call_aws", {"service": "s3",
"operation": "ListBuckets"}))
+ assert result["truncated"] is True
+ assert len(result["data"]) == 200
+
+ @patch("airflow.providers.common.ai.toolsets.aws.AwsBaseHook",
autospec=True)
+ def test_client_resolved_via_amazon_hook(self, mock_hook_cls):
+ client = _make_mock_client()
+ client.list_buckets.return_value = {"Buckets": []}
+ mock_hook_cls.return_value.get_conn.return_value = client
+
+ ts = AWSToolset("aws_prod", allowed_actions=["s3:ListBuckets"],
region_name="eu-west-1")
+ _call(ts, "call_aws", {"service": "s3", "operation": "ListBuckets"})
+ mock_hook_cls.assert_called_once_with(
+ aws_conn_id="aws_prod", client_type="s3", region_name="eu-west-1"
+ )
+
+ def test_unknown_tool_raises_value_error(self):
+ ts = _make_toolset()
+ with pytest.raises(ValueError, match="Unknown tool"):
+ _call(ts, "use_aws", {})
diff --git a/uv.lock b/uv.lock
index 976c0d847e1..ea072289887 100644
--- a/uv.lock
+++ b/uv.lock
@@ -4369,12 +4369,18 @@ dependencies = [
]
[package.optional-dependencies]
+amazon = [
+ { name = "apache-airflow-providers-amazon" },
+]
anthropic = [
{ name = "pydantic-ai-slim", extra = ["anthropic"] },
]
avro = [
{ name = "fastavro" },
]
+aws = [
+ { name = "apache-airflow-providers-amazon" },
+]
bedrock = [
{ name = "pydantic-ai-slim", extra = ["bedrock"] },
]
@@ -4430,6 +4436,7 @@ sql = [
dev = [
{ name = "apache-airflow" },
{ name = "apache-airflow-devel-common" },
+ { name = "apache-airflow-providers-amazon" },
{ name = "apache-airflow-providers-common-compat" },
{ name = "apache-airflow-providers-common-sql", extra = ["datafusion"] },
{ name = "apache-airflow-providers-git" },
@@ -4450,6 +4457,8 @@ docs = [
[package.metadata]
requires-dist = [
{ name = "apache-airflow", editable = "." },
+ { name = "apache-airflow-providers-amazon", marker = "extra == 'amazon'",
editable = "providers/amazon" },
+ { name = "apache-airflow-providers-amazon", marker = "extra == 'aws'",
editable = "providers/amazon" },
{ name = "apache-airflow-providers-common-compat", editable =
"providers/common/compat" },
{ name = "apache-airflow-providers-common-sql", marker = "extra ==
'common-sql'", editable = "providers/common/sql" },
{ name = "apache-airflow-providers-common-sql", marker = "extra == 'sql'",
editable = "providers/common/sql" },
@@ -4478,12 +4487,13 @@ requires-dist = [
{ name = "python-docx", marker = "extra == 'docx'", specifier = ">=1.0.0"
},
{ name = "sqlglot", marker = "extra == 'sql'", specifier = ">=30.0.0" },
]
-provides-extras = ["anthropic", "bedrock", "google", "openai", "mcp",
"code-mode", "shields", "skills", "avro", "parquet", "sql", "common-sql",
"langchain", "llamaindex", "pdf", "docx", "git"]
+provides-extras = ["anthropic", "bedrock", "google", "openai", "mcp",
"code-mode", "shields", "skills", "avro", "parquet", "sql", "aws",
"common-sql", "langchain", "llamaindex", "pdf", "docx", "git", "amazon"]
[package.metadata.requires-dev]
dev = [
{ name = "apache-airflow", editable = "." },
{ name = "apache-airflow-devel-common", editable = "devel-common" },
+ { name = "apache-airflow-providers-amazon", editable = "providers/amazon"
},
{ name = "apache-airflow-providers-common-compat", editable =
"providers/common/compat" },
{ name = "apache-airflow-providers-common-sql", editable =
"providers/common/sql" },
{ name = "apache-airflow-providers-common-sql", extras = ["datafusion"],
editable = "providers/common/sql" },
