gopidesupavan commented on code in PR #70087: URL: https://github.com/apache/airflow/pull/70087#discussion_r3612417058
########## providers/common/ai/src/airflow/providers/common/ai/toolsets/aws.py: ########## @@ -0,0 +1,507 @@ +# 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 + +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(e) + +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 + +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() + + +# Operations that put credentials or decrypted secrets into the agent's context. +# A wildcard in ``allowed_actions`` never matches these -- each must be listed +# verbatim to be callable. Defense in depth, not exhaustive: 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 _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: Review Comment: these are internal helpers i wouldn't expect users to uses these, anyway i have updated -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
