[プレビュー]Amazon Bedrock KnowledgeBaseとAmazon S3 Vectorsをクロスアカウント接続してみた

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "CrossAccountAccess",
            "Effect": "Allow",
            "Principal": {
                "AWS": "arn:aws:iam::123456789012:root"
            },
            "Action": "s3vectors:*",
            "Resource": "arn:aws:s3vectors:us-west-2:210987654321:bucket/s3-vector-bucket"
        },
        {
            "Sid": "BedrockServiceAccess",
            "Effect": "Allow",
            "Principal": {
                "AWS": "arn:aws:iam::123456789012:role/AmazonBedrockExecutionRoleForKnowledgeBase"
            },
            "Action": [
                "s3vectors:GetIndex",
                "s3vectors:QueryVectors",
                "s3vectors:PutVectors",
                "s3vectors:GetVectors",
                "s3vectors:DeleteVectors"
            ],
            "Resource": [
                "arn:aws:s3vectors:us-west-2:210987654321:bucket/s3-vector-bucket",
                "arn:aws:s3vectors:us-west-2:210987654321:bucket/s3-vector-bucket/*"
            ]
        }
    ]
}

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "BedrockInvokeModelStatement",
            "Effect": "Allow",
            "Action": [
                "bedrock:InvokeModel"
            ],
            "Resource": [
                "arn:aws:bedrock:us-west-2::foundation-model/amazon.titan-embed-text-v2:0"
            ]
        },
        {
            "Sid": "S3ListBucketStatement",
            "Effect": "Allow",
            "Action": [
                "s3:ListBucket"
            ],
            "Resource": [
                "arn:aws:s3:::kb-documents-bucket-123456789012"
            ],
            "Condition": {
                "StringEquals": {
                    "aws:ResourceAccount": [
                        "123456789012"
                    ]
                }
            }
        },
        {
            "Sid": "S3GetObjectStatement",
            "Effect": "Allow",
            "Action": [
                "s3:GetObject"
            ],
            "Resource": [
                "arn:aws:s3:::kb-documents-bucket-123456789012/*"
            ],
            "Condition": {
                "StringEquals": {
                    "aws:ResourceAccount": [
                        "123456789012"
                    ]
                }
            }
        },
        {
            "Sid": "S3VectorsPermissions",
            "Effect": "Allow",
            "Action": [
                "s3vectors:GetIndex",
                "s3vectors:QueryVectors",
                "s3vectors:PutVectors",
                "s3vectors:GetVectors",
                "s3vectors:DeleteVectors"
            ],
            "Resource": "arn:aws:s3vectors:us-west-2:210987654321:bucket/s3-vector-bucket/index/*"
        }
    ]
}

import logging
import boto3
import json
from botocore.exceptions import ClientError

"""
アカウントA： 123456789012
アカウントB： 210987654321
"""

def create_knowledge_base(bedrock_agent_client):
    """
    Create a knowledge base with S3 vector storage

    :param bedrock_agent_client: boto3 Bedrock Agent client
    :return: Knowledge base ID if created successfully, else None
    """

    try:
        # Knowledge base configuration
        knowledge_base_config = {
            "type": "VECTOR",
            "vectorKnowledgeBaseConfiguration": {
                "embeddingModelArn": "arn:aws:bedrock:us-west-2::foundation-model/amazon.titan-embed-text-v2:0"
            },
        }

        # Storage configuration for S3 Vector
        storage_config = {
            "type": "S3_VECTORS",
            "s3VectorsConfiguration": {
                "indexArn": "arn:aws:s3vectors:us-west-2:210987654321:bucket/s3-vector-bucket/index/default-index",
                "vectorBucketArn": "arn:aws:s3vectors:us-west-2:210987654321:bucket/s3-vector-bucket",
            },
        }

        # Create the knowledge base
        response = bedrock_agent_client.create_knowledge_base(
            name="test-kb-s3-vectors",
            description="Knowledge base using S3 Vector storage for cross-account setup",
            roleArn="arn:aws:iam::123456789012:role/AmazonBedrockExecutionRoleForKnowledgeBase",
            knowledgeBaseConfiguration=knowledge_base_config,
            storageConfiguration=storage_config,
        )

        knowledge_base_id = response["knowledgeBase"]["knowledgeBaseId"]
        logging.info(
            f"Knowledge base created successfully with ID: {knowledge_base_id}"
        )
        return knowledge_base_id

    except ClientError as e:
        logging.error(f"Failed to create knowledge base: {e}")
        return None

def create_data_source(bedrock_agent_client, knowledge_base_id):
    """
    Create a data source for the knowledge base

    :param bedrock_agent_client: boto3 Bedrock Agent client
    :param knowledge_base_id: ID of the knowledge base
    :return: Data source ID if created successfully, else None
    """

    try:
        # Data source configuration
        data_source_config = {
            "type": "S3",
            "s3Configuration": {
                "bucketArn": "arn:aws:s3:::kb-documents-bucket-123456789012",
                "inclusionPrefixes": ["documents/"],
            },
        }

        # Create the data source
        response = bedrock_agent_client.create_data_source(
            knowledgeBaseId=knowledge_base_id,
            name="data-source",
            dataSourceConfiguration=data_source_config,
        )

        data_source_id = response["dataSource"]["dataSourceId"]
        logging.info(f"Data source created successfully with ID: {data_source_id}")
        return data_source_id

    except ClientError as e:
        logging.error(f"Failed to create data source: {e}")
        return None

def start_ingestion_job(bedrock_agent_client, knowledge_base_id, data_source_id):
    """
    Start an ingestion job to process documents in the data source

    :param bedrock_agent_client: boto3 Bedrock Agent client
    :param knowledge_base_id: ID of the knowledge base
    :param data_source_id: ID of the data source
    :return: Ingestion job ID if started successfully, else None
    """

    try:
        response = bedrock_agent_client.start_ingestion_job(
            knowledgeBaseId=knowledge_base_id, dataSourceId=data_source_id
        )

        ingestion_job_id = response["ingestionJob"]["ingestionJobId"]
        logging.info(f"Ingestion job started successfully with ID: {ingestion_job_id}")
        return ingestion_job_id

    except ClientError as e:
        logging.error(f"Failed to start ingestion job: {e}")
        return None

if __name__ == "__main__":
    # Configure logging
    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")

    # Create Bedrock Agent client
    bedrock_agent_client = boto3.client("bedrock-agent", region_name="us-west-2")

    # Create knowledge base
    knowledge_base_id = create_knowledge_base(bedrock_agent_client)

    if knowledge_base_id:
        print(f"Knowledge base created successfully with ID: {knowledge_base_id}")

        # Create data source
        data_source_id = create_data_source(
            bedrock_agent_client=bedrock_agent_client,
            knowledge_base_id=knowledge_base_id,
        )

        if data_source_id:
            print(f"Data source created successfully with ID: {data_source_id}")

            # Start ingestion job
            ingestion_job_id = start_ingestion_job(
                bedrock_agent_client=bedrock_agent_client,
                knowledge_base_id=knowledge_base_id,
                data_source_id=data_source_id,
            )

            if ingestion_job_id:
                print(f"Ingestion job started successfully with ID: {ingestion_job_id}")
                print(
                    "Documents will be processed and vectorized into the S3 Vector bucket"
                )
            else:
                print("Failed to start ingestion job")
        else:
            print("Failed to create data source")
    else:
        print("Failed to create knowledge base")