airflow.providers.common.ai.toolsets.aws¶

Curated AWS toolset exposing allow-listed boto3 operations as pydantic-ai tools.

Classes¶

AWSToolset

Curated toolset that gives an LLM agent allow-listed access to AWS APIs.

Module Contents¶

class airflow.providers.common.ai.toolsets.aws.AWSToolset(aws_conn_id='aws_default', *, allowed_actions, region_name=None, max_items=1000, max_output_bytes=65536)[source]¶

Bases: pydantic_ai.toolsets.abstract.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 (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.

Parameters:

aws_conn_id (str) – Airflow connection ID for AWS credentials.
allowed_actions (list[str]) –
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.
region_name (str | None) – AWS region for API calls. None (default) uses the region configured on the connection. Deliberately not exposed as a tool argument.
max_items (int) – Upper bound on items aggregated from paginated operations (applied via the paginator’s MaxItems). Default 1000.
max_output_bytes (int) – Upper bound on the serialized response returned to the agent. Larger payloads are clipped and flagged with "truncated": true. Default 65536.

property id: str[source]¶

An ID for the toolset that is unique among all toolsets registered with the same agent.

If you’re implementing a concrete implementation that users can instantiate more than once, you should let them optionally pass a custom ID to the constructor and return that here.

A toolset needs to have an ID in order to be used in a durable execution environment like Temporal, in which case the ID will be used to identify the toolset’s activities within the workflow.

async get_tools(ctx)[source]¶

The tools that are available in this toolset.

async call_tool(name, tool_args, ctx, tool)[source]¶

Call a tool with the given arguments.

Args:: name: The name of the tool to call. tool_args: The arguments to pass to the tool. ctx: The run context. tool: The tool definition returned by [get_tools][pydantic_ai.toolsets.AbstractToolset.get_tools] that was called.