ferruzzi commented on code in PR #39245: URL: https://github.com/apache/airflow/pull/39245#discussion_r1585192344
########## tests/system/providers/amazon/aws/example_bedrock_knowledge_base.py: ########## @@ -0,0 +1,519 @@ +# 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 json +import logging +import os.path +import tempfile +from datetime import datetime +from time import sleep +from urllib.request import urlretrieve + +import boto3 +from botocore.exceptions import ClientError +from opensearchpy import ( + AuthorizationException, + AWSV4SignerAuth, + OpenSearch, + RequestsHttpConnection, +) + +from airflow import DAG +from airflow.decorators import task +from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentHook +from airflow.providers.amazon.aws.hooks.opensearch_serverless import OpenSearchServerlessHook +from airflow.providers.amazon.aws.hooks.s3 import S3Hook +from airflow.providers.amazon.aws.hooks.sts import StsHook +from airflow.providers.amazon.aws.operators.bedrock import ( + BedrockCreateDataSourceOperator, + BedrockCreateKnowledgeBaseOperator, + BedrockIngestDataOperator, +) +from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator +from airflow.providers.amazon.aws.sensors.bedrock import ( + BedrockIngestionJobSensor, + BedrockKnowledgeBaseActiveSensor, +) +from airflow.providers.amazon.aws.sensors.opensearch_serverless import ( + OpenSearchServerlessCollectionActiveSensor, +) +from airflow.utils.helpers import chain +from airflow.utils.trigger_rule import TriggerRule +from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder + +############################################################################### +# NOTE: The account running this test must first manually request access to +# the `Titan Embeddings G1 - Text` foundation model via the Bedrock console. +# Gaining access to the model can take 24 hours from the time of request. +############################################################################### + +# Externally fetched variables: +ROLE_ARN_KEY = "ROLE_ARN" +sys_test_context_task = SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).build() + +DAG_ID = "example_bedrock_knowledge_base" + +log = logging.getLogger(__name__) + + +@task +def create_opensearch_policies(bedrock_role_arn: str, collection_name: str): + """ + Create security, network and data access policies within Amazon OpenSearch Serverless. + + :param bedrock_role_arn: Arn of the Bedrock Knowledge Base Execution Role. + :param collection_name: Name of the OpenSearch collection to apply the policies to. + """ + + encryption_policy_name = f"{naming_prefix}rag-sp-{env_id}" + network_policy_name = f"{naming_prefix}rag-np-{env_id}" + access_policy_name = f"{naming_prefix}rag-ap-{env_id}" + + def _create_security_policy(name, policy_type, policy): + try: + aoss_client.create_security_policy(name=name, policy=json.dumps(policy), type=policy_type) + except ClientError as e: + if e.response["Error"]["Code"] == "ConflictException": + log.info("OpenSearch security policy %s already exists.", name) + raise + + def _create_access_policy(name, policy_type, policy): + try: + aoss_client.create_access_policy(name=name, policy=json.dumps(policy), type=policy_type) + except ClientError as e: + if e.response["Error"]["Code"] == "ConflictException": + log.info("OpenSearch data access policy %s already exists.", name) + raise + + _create_security_policy( + name=encryption_policy_name, + policy_type="encryption", + policy={ + "Rules": [{"Resource": [f"collection/{collection_name}"], "ResourceType": "collection"}], + "AWSOwnedKey": True, + }, + ) + + _create_security_policy( + name=network_policy_name, + policy_type="network", + policy=[ + { + "Rules": [{"Resource": [f"collection/{collection_name}"], "ResourceType": "collection"}], + "AllowFromPublic": True, + } + ], + ) + + _create_access_policy( + name=access_policy_name, + policy_type="data", + policy=[ + { + "Rules": [ + { + "Resource": [f"collection/{collection_name}"], + "Permission": [ + "aoss:CreateCollectionItems", + "aoss:DeleteCollectionItems", + "aoss:UpdateCollectionItems", + "aoss:DescribeCollectionItems", + ], + "ResourceType": "collection", + }, + { + "Resource": [f"index/{collection_name}/*"], + "Permission": [ + "aoss:CreateIndex", + "aoss:DeleteIndex", + "aoss:UpdateIndex", + "aoss:DescribeIndex", + "aoss:ReadDocument", + "aoss:WriteDocument", + ], + "ResourceType": "index", + }, + ], + "Principal": [(StsHook().conn.get_caller_identity()["Arn"]), bedrock_role_arn], + } + ], + ) + + +@task +def create_collection(collection_name: str): + """ + Call the Amazon OpenSearch Serverless API and create a collection with the provided name. + + :param collection_name: The name of the Collection to create. + """ + log.info("\nCreating collection: %s.", collection_name) + return aoss_client.create_collection(name=collection_name, type="VECTORSEARCH")["createCollectionDetail"][ + "id" + ] + + +@task +def create_vector_index(collection_id: str, region: str): + """ + Use the OpenSearchPy client to create the vector index for the Amazon Open Search Serverless Collection. + + :param collection_id: ID of the collection to be indexed. + :param region: Name of the AWS region the collection resides in. + """ + # Build the OpenSearch client + oss_client = OpenSearch( + hosts=[{"host": f"{collection_id}.{region}.aoss.amazonaws.com", "port": 443}], + http_auth=AWSV4SignerAuth(boto3.Session().get_credentials(), region, "aoss"), + use_ssl=True, + verify_certs=True, + connection_class=RequestsHttpConnection, + timeout=300, + ) + index_config = { + "settings": { + "index.knn": "true", + "number_of_shards": 1, + "knn.algo_param.ef_search": 512, + "number_of_replicas": 0, + }, + "mappings": { + "properties": { + "vector": { + "type": "knn_vector", + "dimension": 1536, + "method": {"name": "hnsw", "engine": "faiss", "space_type": "l2"}, + }, + "text": {"type": "text"}, + "text-metadata": {"type": "text"}, + } + }, + } + + retries = 35 + while retries > 0: + try: + response = oss_client.indices.create(index=index_name, body=json.dumps(index_config)) + log.info("Creating index: %s.", response) + break + except AuthorizationException as e: + # Index creation can take up to a minute and there is no (apparent?) way to check the current state. + log.info( + "Access denied; policy permissions have likely not yet propagated, %s tries remaining.", + retries, + ) + log.debug(e) + retries -= 1 + sleep(2) + + +@task +def copy_data_to_s3(bucket: str): + """ + Download some sample data and upload it to 3S. + + :param bucket: Name of the Amazon S3 bucket to send the data to. + """ + + # Monkey patch the list of names available for NamedTempFile so we can pick the names of the downloaded files. + backup_get_candidate_names = tempfile._get_candidate_names # type: ignore[attr-defined] + destinations = iter( + [ + "AMZN-2022-Shareholder-Letter.pdf", + "AMZN-2021-Shareholder-Letter.pdf", + "AMZN-2020-Shareholder-Letter.pdf", + "AMZN-2019-Shareholder-Letter.pdf", + ] + ) + tempfile._get_candidate_names = lambda: destinations # type: ignore[attr-defined] + + # Download the sample data files, save them as named temp files using the names above, and upload to S3. + sources = [ + "https://s2.q4cdn.com/299287126/files/doc_financials/2023/ar/2022-Shareholder-Letter.pdf";, + "https://s2.q4cdn.com/299287126/files/doc_financials/2022/ar/2021-Shareholder-Letter.pdf";, + "https://s2.q4cdn.com/299287126/files/doc_financials/2021/ar/Amazon-2020-Shareholder-Letter-and-1997-Shareholder-Letter.pdf";, + "https://s2.q4cdn.com/299287126/files/doc_financials/2020/ar/2019-Shareholder-Letter.pdf";, + ] + + for source in sources: + with tempfile.NamedTemporaryFile(mode="w", prefix="") as data_file: + urlretrieve(source, data_file.name) + S3Hook().conn.upload_file( + Filename=data_file.name, Bucket=bucket, Key=os.path.basename(data_file.name) + ) + + # Revert the monkey patch. + tempfile._get_candidate_names = backup_get_candidate_names # type: ignore[attr-defined] + # Verify the path reversion worked. + with tempfile.NamedTemporaryFile(mode="w", prefix=""): + # If the reversion above did not apply correctly, this will fail with + # a StopIteration error because the iterator will run out of names. + ... + + +@task +def get_collection_arn(collection_id: str): + """ + Return a collection ARN for a given collection ID. + + :param collection_id: ID of the collection to be indexed. + """ + return next( + colxn["arn"] + for colxn in aoss_client.list_collections()["collectionSummaries"] + if colxn["id"] == collection_id + ) + + +# [START howto_operator_bedrock_delete_data_source] +@task(trigger_rule=TriggerRule.ALL_DONE) +def delete_data_source(knowledge_base_id: str, data_source_id: str): + """ + Delete the Amazon Bedrock data source created earlier. + + .. seealso:: + For more information on how to use this sensor, take a look at the guide: + :ref:`howto_operator:BedrockDeleteDataSource` + + :param knowledge_base_id: The unique identifier of the knowledge base which the data source is attached to. + :param data_source_id: The unique identifier of the data source to delete. + """ + log.info("Deleting data source %s from Knowledge Base %s.", data_source_id, knowledge_base_id) + bedrock_agent_client.delete_data_source(dataSourceId=data_source_id, knowledgeBaseId=knowledge_base_id) + + +# [END howto_operator_bedrock_delete_data_source] + + +# [START howto_operator_bedrock_delete_knowledge_base] +@task(trigger_rule=TriggerRule.ALL_DONE) +def delete_knowledge_base(knowledge_base_id: str): + """ + Delete the Amazon Bedrock knowledge base created earlier. + + .. seealso:: + For more information on how to use this sensor, take a look at the guide: + :ref:`howto/operator:BedrockDeleteKnowledgeBase` + + :param knowledge_base_id: The unique identifier of the knowledge base to delete. + """ + log.info("Deleting Knowledge Base %s.", knowledge_base_id) + bedrock_agent_client.delete_knowledge_base(knowledgeBaseId=knowledge_base_id) + + +# [END howto_operator_bedrock_delete_knowledge_base] + + +@task(trigger_rule=TriggerRule.ALL_DONE) +def delete_vector_index(collection_id: str): + """ + Delete the vector index created earlier. + + :param collection_id: ID of the collection to be indexed. + """ + host = f"{collection_id}.{region_name}.aoss.amazonaws.com" + credentials = boto3.Session().get_credentials() + awsauth = AWSV4SignerAuth(credentials, region_name, "aoss") + + # Build the OpenSearch client + oss_client = OpenSearch( + hosts=[{"host": host, "port": 443}], + http_auth=awsauth, + use_ssl=True, + verify_certs=True, + connection_class=RequestsHttpConnection, + timeout=300, + ) + oss_client.indices.delete(index=index_name) + + +@task(trigger_rule=TriggerRule.ALL_DONE) +def delete_collection(collection_id: str): + """ + Delete the OpenSearch collection created earlier. + + :param collection_id: ID of the collection to be indexed. + """ + log.info("Deleting collection %s.", collection_id) + aoss_client.delete_collection(id=collection_id) + + +@task(trigger_rule=TriggerRule.ALL_DONE) +def delete_opensearch_policies(): + """Delete the security, network and data access policies created earlier.""" + + access_policies = aoss_client.list_access_policies( + type="data", resource=[f"collection/{vector_store_name}"] + )["accessPolicySummaries"] + log.info("Found access policies for %s: %s", vector_store_name, access_policies) + if not access_policies: + raise Exception("No access policies found?") + for policy in access_policies: + log.info("Deleting access policy for %s: %s", vector_store_name, policy["name"]) + aoss_client.delete_access_policy(name=policy["name"], type="data") + + for policy_type in ["encryption", "network"]: + policies = aoss_client.list_security_policies( + type=policy_type, resource=[f"collection/{vector_store_name}"] + )["securityPolicySummaries"] + if not policies: + raise Exception("No security policies found?") + log.info("Found %s security policies for %s: %s", policy_type, vector_store_name, policies) + for policy in policies: + log.info("Deleting %s security policy for %s: %s", policy_type, vector_store_name, policy["name"]) + aoss_client.delete_security_policy(name=policy["name"], type=policy_type) + + +with DAG( + dag_id=DAG_ID, + schedule="@once", + start_date=datetime(2021, 1, 1), + tags=["example"], + catchup=False, +) as dag: + test_context = sys_test_context_task() + env_id = test_context["ENV_ID"] + + aoss_client = OpenSearchServerlessHook().conn + bedrock_agent_client = BedrockAgentHook().conn Review Comment: They both get used a few times so I didn't want to recreate them multiple times. Will the `aws_con_id=None` skip the db call and fix the static check? For a system test, I think that's a fine solution. I can try it out and see. Thanks. -- 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]
