こんにちは、つくぼし(tsukuboshi0755)です!
先日以下のアップデートにより、Knowledge Base for Amazon BedrockとAgents for Amazon BedrockがCloudFormationによるデプロイをサポートしました。
今回はこのアップデートを利用して、ベクターストアにAuroraを利用したKnowledge Base for Amazon Bedrockを、CloudFormation一撃で作成可能なテンプレートについて紹介します!
なおベクターストアをセットアップするために一部カスタムリソースを含む形になるため、その点にご留意頂けるとありがたいです。
構成
今回CloudFormationで作成する構成図は以下のとおりです。
Knowledge Base for Amazon Bedrockに必要なリソースとして、以下のリソースも合わせて作成します。
- ベクターストア:Aurora PostgreSQL
- 機密情報:Secrets Manager
- データソース:S3
またLambda-Backed Custom Resourceを利用して、Auroraに対するデータベースのセットアップを自動化します。
テンプレート
全体のコードは以下の通りです。
CloudFormationコード
template.yaml
Description: 'Knowledge Base for Amazon Bedrock with Aurora PostgreSQL (including db setup)'
Mappings:
DatabaseMap:
DatabaseName:
Name: bedrockkbdb
TableName:
Name: bedrock_integration.bedrock_kb
SchemaName:
Name: bedrock_integration
Username:
Name: bedrock_user
PrimaryKeyField:
Name: id
VectorField:
Name: embedding
TextField:
Name: chunks
MetadataField:
Name: metadata
Parameters:
DatabasePassword:
Type: String
Default: 'P@ssword123'
Description: 'The password for the database user.'
NoEcho: true
EmbeddingModelId:
Type: String
Default: amazon.titan-embed-text-v1
Description: 'The Id of the Bedrock model that is used to generate embeddings.'
Resources:
VPC:
Type: AWS::EC2::VPC
Properties:
CidrBlock: 192.168.0.0/24
EnableDnsSupport: true
EnableDnsHostnames: true
Tags:
- Key: Name
Value: !Sub ${AWS::StackName}-vpc
PrivateSubnet1:
Type: AWS::EC2::Subnet
Properties:
VpcId: !Ref VPC
AvailabilityZone: !Sub ${AWS::Region}a
CidrBlock: 192.168.0.0/28
MapPublicIpOnLaunch: false
Tags:
- Key: Name
Value: !Sub ${AWS::StackName}-private-subnet-1a
PrivateSubnet2:
Type: AWS::EC2::Subnet
Properties:
VpcId: !Ref VPC
AvailabilityZone: !Sub ${AWS::Region}c
CidrBlock: 192.168.0.16/28
MapPublicIpOnLaunch: false
Tags:
- Key: Name
Value: !Sub ${AWS::StackName}-private-subnet-1c
PrivateRouteTable:
Type: AWS::EC2::RouteTable
Properties:
VpcId: !Ref VPC
Tags:
- Key: Name
Value: !Sub ${AWS::StackName}-private-subnet-rtb
PrivateSubnet1RouteTableAssociation:
Type: AWS::EC2::SubnetRouteTableAssociation
Properties:
RouteTableId: !Ref PrivateRouteTable
SubnetId: !Ref PrivateSubnet1
PrivateSubnet2RouteTableAssociation:
Type: AWS::EC2::SubnetRouteTableAssociation
Properties:
RouteTableId: !Ref PrivateRouteTable
SubnetId: !Ref PrivateSubnet2
DBSubnetGroup:
Type: AWS::RDS::DBSubnetGroup
Properties:
DBSubnetGroupDescription: Bedrock RDS Subnet Group
SubnetIds:
- !Ref PrivateSubnet1
- !Ref PrivateSubnet2
AuroraCluster:
Type: AWS::RDS::DBCluster
DeletionPolicy: Delete
Properties:
DatabaseName: !FindInMap [DatabaseMap, DatabaseName, Name]
Engine: aurora-postgresql
EngineVersion: 15.5
DBSubnetGroupName: !Ref DBSubnetGroup
MasterUsername: postgresql
ManageMasterUserPassword: true
ServerlessV2ScalingConfiguration:
MinCapacity: 0.5
MaxCapacity: 1.0
StorageEncrypted: true
EnableHttpEndpoint: true
AuroraDBInstance:
Type: AWS::RDS::DBInstance
DeletionPolicy: Delete
Properties:
Engine: aurora-postgresql
DBInstanceClass: db.serverless
DBClusterIdentifier: !Ref AuroraCluster
SecretForAurora:
Type: AWS::SecretsManager::Secret
Properties:
Name: !Sub ${AWS::StackName}-db-secret-for-bedrock
Description: 'Secret for the database user for Bedrock'
SecretString: !Sub
- '{ "username":"${DatabaseUser}", "password":"${DatabasePassword}"}'
- DatabaseUser: !FindInMap [DatabaseMap, Username, Name]
DataSourceBucket:
Type: AWS::S3::Bucket
Properties:
BucketName: !Sub '${AWS::StackName}-ds-bucket-${AWS::AccountId}'
PublicAccessBlockConfiguration:
BlockPublicAcls: true
BlockPublicPolicy: true
IgnorePublicAcls: true
RestrictPublicBuckets: true
BucketEncryption:
ServerSideEncryptionConfiguration:
- ServerSideEncryptionByDefault:
SSEAlgorithm: AES256
BucketKeyEnabled: true
OwnershipControls:
Rules:
- ObjectOwnership: BucketOwnerEnforced
BedrockAccessPolicy:
Type: AWS::IAM::ManagedPolicy
Properties:
ManagedPolicyName: !Sub ${AWS::StackName}-bedrock-access-policy
PolicyDocument:
Version: '2012-10-17'
Statement:
- Effect: Allow
Action:
- 'bedrock:InvokeModel'
Resource: !Sub arn:aws:bedrock:${AWS::Region}::foundation-model/${EmbeddingModelId}
Roles:
- !Ref BedrockKnowledgeBaseRole
SecretsAccessPolicy:
Type: AWS::IAM::ManagedPolicy
Properties:
ManagedPolicyName: !Sub ${AWS::StackName}-secret-access-policy
PolicyDocument:
Version: '2012-10-17'
Statement:
- Effect: Allow
Action:
- 'secretsmanager:GetSecretValue'
Resource: !Ref SecretForAurora
Roles:
- !Ref BedrockKnowledgeBaseRole
AuroraAccessPolicy:
Type: AWS::IAM::ManagedPolicy
Properties:
ManagedPolicyName: !Sub ${AWS::StackName}-aurora-access-policy
PolicyDocument:
Version: '2012-10-17'
Statement:
- Effect: Allow
Action:
- 'rds:DescribeDBClusters'
- 'rds-data:BatchExecuteStatement'
- 'rds-data:ExecuteStatement'
Resource: !GetAtt AuroraCluster.DBClusterArn
- Effect: Allow
Action:
- 'secretsmanager:GetSecretValue'
Resource: !GetAtt AuroraCluster.MasterUserSecret.SecretArn
Roles:
- !Ref BedrockKnowledgeBaseRole
- !Ref ExecSQLFunctionRole
S3AccessPolicy:
Type: AWS::IAM::ManagedPolicy
Properties:
ManagedPolicyName: !Sub ${AWS::StackName}-s3-access-policy
PolicyDocument:
Version: '2012-10-17'
Statement:
- Sid: S3ListBucketStatement
Effect: Allow
Action:
- 's3:ListBucket'
Resource: !GetAtt DataSourceBucket.Arn
Condition:
StringEquals:
aws:ResourceAccount: !Ref 'AWS::AccountId'
- Sid: S3GetObjectStatement
Effect: Allow
Action:
- 's3:GetObject'
Resource: !Sub
- '${DataSourceBucketArn}/*'
- DataSourceBucketArn: !GetAtt DataSourceBucket.Arn
Condition:
StringEquals:
aws:ResourceAccount: !Ref 'AWS::AccountId'
Roles:
- !Ref BedrockKnowledgeBaseRole
SetupAuroraData:
Type: 'Custom::SetupAuroraData'
Properties:
ServiceToken: !GetAtt ExecSQLFunction.Arn
ResourceArn: !GetAtt AuroraCluster.DBClusterArn
SecretArn: !GetAtt AuroraCluster.MasterUserSecret.SecretArn
DatabaseName: !FindInMap [DatabaseMap, DatabaseName, Name]
DatabasePassword: !Ref DatabasePassword
TableName: !FindInMap [DatabaseMap, TableName, Name]
SchemaName: !FindInMap [DatabaseMap, SchemaName, Name]
UserName: !FindInMap [DatabaseMap, Username, Name]
MetadataField: !FindInMap [DatabaseMap, MetadataField, Name]
PrimaryKeyField: !FindInMap [DatabaseMap, PrimaryKeyField, Name]
TextField: !FindInMap [DatabaseMap, TextField, Name]
VectorField: !FindInMap [DatabaseMap, VectorField, Name]
DependsOn: AuroraDBInstance
ExecSQLFunctionRole:
Type: 'AWS::IAM::Role'
Properties:
RoleName: !Sub ${AWS::StackName}-execsql-function-role
AssumeRolePolicyDocument:
Version: '2012-10-17'
Statement:
- Effect: Allow
Principal:
Service: lambda.amazonaws.com
Action: 'sts:AssumeRole'
ManagedPolicyArns:
- arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
ExecSQLFunction:
Type: AWS::Lambda::Function
Properties:
FunctionName: !Sub ${AWS::StackName}-execsql-function
Handler: index.lambda_handler
Role: !GetAtt ExecSQLFunctionRole.Arn
Runtime: python3.12
Timeout: 600
LoggingConfig:
LogFormat: JSON
ApplicationLogLevel: INFO
SystemLogLevel: INFO
Code:
ZipFile: |
import boto3
import cfnresponse
import logging
from typing import Any, Dict
logger = logging.getLogger()
rds_data = boto3.client('rds-data')
def execute_statement(resource_arn: str, database_name: str, secret_arn: str, sql: str) -> Any:
response = rds_data.execute_statement(
resourceArn=resource_arn,
database=database_name,
secretArn=secret_arn,
sql=sql
)
return response
def lambda_handler(event: Dict[str, Any], context: Any) -> None:
try:
resource_arn = event['ResourceProperties']['ResourceArn']
secret_arn = event['ResourceProperties']['SecretArn']
database_name = event['ResourceProperties']['DatabaseName']
database_password = event['ResourceProperties']['DatabasePassword']
table_name = event['ResourceProperties']['TableName']
schema_name = event['ResourceProperties']['SchemaName']
user_name = event['ResourceProperties']['UserName']
metadata_field = event['ResourceProperties']['MetadataField']
primary_key_field = event['ResourceProperties']['PrimaryKeyField']
text_field = event['ResourceProperties']['TextField']
vector_field = event['ResourceProperties']['VectorField']
if event['RequestType'] == 'Create':
create_extension = f"""
CREATE EXTENSION IF NOT EXISTS vector;
"""
create_extension_res = execute_statement(resource_arn, database_name, secret_arn, create_extension)
logger.info(f"Create Extension Response: {create_extension_res}")
create_role = f"""
CREATE ROLE {user_name} WITH PASSWORD '{database_password}' LOGIN;
"""
create_role_res = execute_statement(resource_arn, database_name, secret_arn, create_role)
logger.info(f"Create Role Response: {create_role_res}")
create_shema = f"""
CREATE SCHEMA {schema_name};
"""
create_shema_res = execute_statement(resource_arn, database_name, secret_arn, create_shema)
logger.info(f"Create Schema Response: {create_shema_res}")
grant_schema = f"""
GRANT ALL ON SCHEMA {schema_name} to {user_name};
"""
grant_schema_res = execute_statement(resource_arn, database_name, secret_arn, grant_schema)
logger.info(f"Grant Schema Response: {grant_schema_res}")
create_table = f"""
CREATE TABLE {table_name} ({primary_key_field} uuid PRIMARY KEY, {vector_field} vector(1536), {text_field} text, {metadata_field} json)
"""
create_table_res = execute_statement(resource_arn, database_name, secret_arn, create_table)
logger.info(f"Create Table Response: {create_table_res}")
grant_table = f"""
GRANT ALL ON TABLE {table_name} TO {user_name};
"""
grant_table_res = execute_statement(resource_arn, database_name, secret_arn, grant_table)
logger.info(f"Grant Table Response: {grant_table_res}")
create_index = f"""
CREATE INDEX on {table_name} USING hnsw ({vector_field} vector_cosine_ops);
"""
create_index_res = execute_statement(resource_arn, database_name, secret_arn, create_index)
logger.info(f"Create Index Response: {create_index_res}")
cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
if event['RequestType'] == 'Update':
cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
if event['RequestType'] == 'Delete':
cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
except Exception as e:
cfnresponse.send(event, context, cfnresponse.FAILED, {'Message': str(e)})
BedrockKnowledgeBaseRole:
Type: AWS::IAM::Role
Properties:
RoleName: !Sub ${AWS::StackName}-bedrock-kb-role
AssumeRolePolicyDocument:
Version: '2012-10-17'
Statement:
- Effect: Allow
Principal:
Service: [bedrock.amazonaws.com]
Action: ['sts:AssumeRole']
BedrockKnowledgeBase:
Type: AWS::Bedrock::KnowledgeBase
Properties:
Name: !Sub ${AWS::StackName}-knowledge-base
KnowledgeBaseConfiguration:
Type: VECTOR
VectorKnowledgeBaseConfiguration:
EmbeddingModelArn: !Sub arn:aws:bedrock:${AWS::Region}::foundation-model/${EmbeddingModelId}
RoleArn: !GetAtt BedrockKnowledgeBaseRole.Arn
StorageConfiguration:
Type: RDS
RdsConfiguration:
CredentialsSecretArn: !Ref SecretForAurora
DatabaseName: !FindInMap [DatabaseMap, DatabaseName, Name]
FieldMapping:
MetadataField: !FindInMap [DatabaseMap, MetadataField, Name]
PrimaryKeyField: !FindInMap [DatabaseMap, PrimaryKeyField, Name]
TextField: !FindInMap [DatabaseMap, TextField, Name]
VectorField: !FindInMap [DatabaseMap, VectorField, Name]
ResourceArn: !GetAtt AuroraCluster.DBClusterArn
TableName: !FindInMap [DatabaseMap, TableName, Name]
DependsOn: SetupAuroraData
BedrockKnowledgeBaseDS:
Type: AWS::Bedrock::DataSource
Properties:
KnowledgeBaseId: !Ref BedrockKnowledgeBase
Name: !Sub ${AWS::StackName}-data-source
DataSourceConfiguration:
Type: S3
S3Configuration:
BucketArn: !GetAtt DataSourceBucket.Arn
Outputs:
BedrockKnowledgeBaseId:
Value: !Ref BedrockKnowledgeBase
BedrockDataSourceId:
Value: !Ref BedrockKnowledgeBaseDS
AuroraClusterId:
Value: !Ref AuroraCluster
DSBucketName:
Value: !Ref DataSourceBucket
以下より、本テンプレートで作成される各リソースについて説明します。
VPC
VPC:
Type: AWS::EC2::VPC
Properties:
CidrBlock: 192.168.0.0/24
EnableDnsSupport: true
EnableDnsHostnames: true
Tags:
- Key: Name
Value: !Sub ${AWS::StackName}-vpc
PrivateSubnet1:
Type: AWS::EC2::Subnet
Properties:
VpcId: !Ref VPC
AvailabilityZone: !Sub ${AWS::Region}a
CidrBlock: 192.168.0.0/28
MapPublicIpOnLaunch: false
Tags:
- Key: Name
Value: !Sub ${AWS::StackName}-private-subnet-1a
PrivateSubnet2:
Type: AWS::EC2::Subnet
Properties:
VpcId: !Ref VPC
AvailabilityZone: !Sub ${AWS::Region}c
CidrBlock: 192.168.0.16/28
MapPublicIpOnLaunch: false
Tags:
- Key: Name
Value: !Sub ${AWS::StackName}-private-subnet-1c
PrivateRouteTable:
Type: AWS::EC2::RouteTable
Properties:
VpcId: !Ref VPC
Tags:
- Key: Name
Value: !Sub ${AWS::StackName}-private-subnet-rtb
PrivateSubnet1RouteTableAssociation:
Type: AWS::EC2::SubnetRouteTableAssociation
Properties:
RouteTableId: !Ref PrivateRouteTable
SubnetId: !Ref PrivateSubnet1
PrivateSubnet2RouteTableAssociation:
Type: AWS::EC2::SubnetRouteTableAssociation
Properties:
RouteTableId: !Ref PrivateRouteTable
SubnetId: !Ref PrivateSubnet2
DBSubnetGroup:
Type: AWS::RDS::DBSubnetGroup
Properties:
DBSubnetGroupDescription: Bedrock RDS Subnet Group
SubnetIds:
- !Ref PrivateSubnet1
- !Ref PrivateSubnet2
Auroraに必要なVPC、サブネット、ルートテーブルを作成します。
なお今回の主題ではないため、詳細な説明は割愛します。
Aurora
DBSubnetGroup:
Type: AWS::RDS::DBSubnetGroup
Properties:
DBSubnetGroupDescription: Bedrock RDS Subnet Group
SubnetIds:
- !Ref PrivateSubnet1
- !Ref PrivateSubnet2
AuroraCluster:
Type: AWS::RDS::DBCluster
DeletionPolicy: Delete
Properties:
DatabaseName: !FindInMap [DatabaseMap, DatabaseName, Name]
Engine: aurora-postgresql
EngineVersion: 15.5
DBSubnetGroupName: !Ref DBSubnetGroup
MasterUsername: postgresql
ManageMasterUserPassword: true
ServerlessV2ScalingConfiguration:
MinCapacity: 0.5
MaxCapacity: 1.0
StorageEncrypted: true
EnableHttpEndpoint: true
AuroraDBInstance:
Type: AWS::RDS::DBInstance
DeletionPolicy: Delete
Properties:
Engine: aurora-postgresql
DBInstanceClass: db.serverless
DBClusterIdentifier: !Ref AuroraCluster
Knowledge Base for Amazon Bedrockのベクターストアとして必要なAurora PostgreSQLを作成します。
MasterUsername
及びManageMasterUserPassword
の両方を指定する事で、以下の通りAdminユーザー用Secrets ManagerがAuroraと同時に作成されます。
またMinCapacity
及びMaxCapacity
には、ACUのスケーリング設定を指定します。
さらにEnableHttpEndpoint
を指定する事で、後述のLambda-Backed Custom Resourceで使用するRDS Data APIを有効化します。
なお今回は以下を参考に、検証用途のためバックアップ不要としてDeletionPolicy
をDelete
で指定していますが、必要に応じてSnapshot
やRetain
に変更してください。
Secrets Manager
SecretForAurora:
Type: AWS::SecretsManager::Secret
Properties:
Name: !Sub ${AWS::StackName}-db-secret-for-bedrock
Description: 'Secret for the database user for Bedrock'
SecretString: !Sub
- '{ "username":"${DatabaseUser}", "password":"${DatabasePassword}"}'
- DatabaseUser: !FindInMap [DatabaseMap, Username, Name]
Knowledge Base for Amazon BedrockからAuroraに接続するためのユーザー/パスワード情報をSecrets Managerから取得する必要があるため、Secrets Managerを作成します。
なお本来認証情報をハードコーディングする事はセキュリティ上好ましくない事にご留意ください。
S3 Bucket
DataSourceBucket:
Type: AWS::S3::Bucket
Properties:
BucketName: !Sub '${AWS::StackName}-ds-bucket-${AWS::AccountId}'
PublicAccessBlockConfiguration:
BlockPublicAcls: true
BlockPublicPolicy: true
IgnorePublicAcls: true
RestrictPublicBuckets: true
BucketEncryption:
ServerSideEncryptionConfiguration:
- ServerSideEncryptionByDefault:
SSEAlgorithm: AES256
BucketKeyEnabled: true
OwnershipControls:
Rules:
- ObjectOwnership: BucketOwnerEnforced
Knowledge Base for Amazon Bedrockのデータソースとして必要なS3バケットを作成します。
このバケットに対してユーザーがドキュメントをアップロードする事で、Knowledge Base for Amazon Bedrockがデータソースとして利用できるようになります。
IAM Policy
BedrockAccessPolicy:
Type: AWS::IAM::ManagedPolicy
Properties:
ManagedPolicyName: !Sub ${AWS::StackName}-bedrock-access-policy
PolicyDocument:
Version: '2012-10-17'
Statement:
- Effect: Allow
Action:
- 'bedrock:InvokeModel'
Resource: !Sub arn:aws:bedrock:${AWS::Region}::foundation-model/${EmbeddingModelId}
Roles:
- !Ref BedrockKnowledgeBaseRole
SecretsAccessPolicy:
Type: AWS::IAM::ManagedPolicy
Properties:
ManagedPolicyName: !Sub ${AWS::StackName}-secret-access-policy
PolicyDocument:
Version: '2012-10-17'
Statement:
- Effect: Allow
Action:
- 'secretsmanager:GetSecretValue'
Resource: !Ref SecretForAurora
Roles:
- !Ref BedrockKnowledgeBaseRole
AuroraAccessPolicy:
Type: AWS::IAM::ManagedPolicy
Properties:
ManagedPolicyName: !Sub ${AWS::StackName}-aurora-access-policy
PolicyDocument:
Version: '2012-10-17'
Statement:
- Effect: Allow
Action:
- rds:DescribeDBClusters
- rds-data:BatchExecuteStatement
- rds-data:ExecuteStatement
Resource: !GetAtt AuroraCluster.DBClusterArn
- Effect: Allow
Action:
- 'secretsmanager:GetSecretValue'
Resource: !GetAtt AuroraCluster.MasterUserSecret.SecretArn
Roles:
- !Ref BedrockKnowledgeBaseRole
- !Ref ExecSQLFunctionRole
S3AccessPolicy:
Type: AWS::IAM::ManagedPolicy
Properties:
ManagedPolicyName: !Sub ${AWS::StackName}-s3-access-policy
PolicyDocument:
Version: '2012-10-17'
Statement:
- Sid: S3ListBucketStatement
Effect: Allow
Action:
- 's3:ListBucket'
Resource: !GetAtt DataSourceBucket.Arn
Condition:
StringEquals:
aws:ResourceAccount: !Ref 'AWS::AccountId'
- Sid: S3GetObjectStatement
Effect: Allow
Action:
- 's3:GetObject'
Resource: !Sub
- '${DataSourceBucketArn}/*'
- DataSourceBucketArn: !GetAtt DataSourceBucket.Arn
Condition:
StringEquals:
aws:ResourceAccount: !Ref 'AWS::AccountId'
Roles:
- !Ref BedrockKnowledgeBaseRole
Lambda-Backed Custom Resource及びKnowledge Base for Amazon Bedrockに必要なIAMポリシーを作成します。
Bedrock、Secrets Manager、Aurora、S3に対するアクセス権限を設定し、IAMロールにアタッチします。
Lambda-Backed Custom Resource
SetupAuroraData:
Type: 'Custom::SetupAuroraData'
Properties:
ServiceToken: !GetAtt ExecSQLFunction.Arn
ResourceArn: !GetAtt AuroraCluster.DBClusterArn
SecretArn: !GetAtt AuroraCluster.MasterUserSecret.SecretArn
DatabaseName: !FindInMap [DatabaseMap, DatabaseName, Name]
DatabasePassword: !Ref DatabasePassword
TableName: !FindInMap [DatabaseMap, TableName, Name]
SchemaName: !FindInMap [DatabaseMap, SchemaName, Name]
UserName: !FindInMap [DatabaseMap, Username, Name]
MetadataField: !FindInMap [DatabaseMap, MetadataField, Name]
PrimaryKeyField: !FindInMap [DatabaseMap, PrimaryKeyField, Name]
TextField: !FindInMap [DatabaseMap, TextField, Name]
VectorField: !FindInMap [DatabaseMap, VectorField, Name]
DependsOn: AuroraDBInstance
ExecSQLFunctionRole:
Type: 'AWS::IAM::Role'
Properties:
RoleName: !Sub ${AWS::StackName}-execsql-function-role
AssumeRolePolicyDocument:
Version: '2012-10-17'
Statement:
- Effect: Allow
Principal:
Service: lambda.amazonaws.com
Action: 'sts:AssumeRole'
ManagedPolicyArns:
- arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
ExecSQLFunction:
Type: AWS::Lambda::Function
Properties:
FunctionName: !Sub ${AWS::StackName}-execsql-function
Handler: index.lambda_handler
Role: !GetAtt ExecSQLFunctionRole.Arn
Runtime: python3.12
Timeout: 600
LoggingConfig:
LogFormat: JSON
ApplicationLogLevel: INFO
SystemLogLevel: INFO
Code:
ZipFile: |
import boto3
import cfnresponse
import logging
from typing import Any, Dict
logger = logging.getLogger()
rds_data = boto3.client('rds-data')
def execute_statement(resource_arn: str, database_name: str, secret_arn: str, sql: str) -> Any:
response = rds_data.execute_statement(
resourceArn=resource_arn,
database=database_name,
secretArn=secret_arn,
sql=sql
)
return response
def lambda_handler(event: Dict[str, Any], context: Any) -> None:
try:
resource_arn = event['ResourceProperties']['ResourceArn']
secret_arn = event['ResourceProperties']['SecretArn']
database_name = event['ResourceProperties']['DatabaseName']
database_password = event['ResourceProperties']['DatabasePassword']
table_name = event['ResourceProperties']['TableName']
schema_name = event['ResourceProperties']['SchemaName']
user_name = event['ResourceProperties']['UserName']
metadata_field = event['ResourceProperties']['MetadataField']
primary_key_field = event['ResourceProperties']['PrimaryKeyField']
text_field = event['ResourceProperties']['TextField']
vector_field = event['ResourceProperties']['VectorField']
if event['RequestType'] == 'Create':
create_extension = f"""
CREATE EXTENSION IF NOT EXISTS vector;
"""
create_extension_res = execute_statement(resource_arn, database_name, secret_arn, create_extension)
logger.info(f"Create Extension Response: {create_extension_res}")
create_role = f"""
CREATE ROLE {user_name} WITH PASSWORD '{database_password}' LOGIN;
"""
create_role_res = execute_statement(resource_arn, database_name, secret_arn, create_role)
logger.info(f"Create Role Response: {create_role_res}")
create_shema = f"""
CREATE SCHEMA {schema_name};
"""
create_shema_res = execute_statement(resource_arn, database_name, secret_arn, create_shema)
logger.info(f"Create Schema Response: {create_shema_res}")
grant_schema = f"""
GRANT ALL ON SCHEMA {schema_name} to {user_name};
"""
grant_schema_res = execute_statement(resource_arn, database_name, secret_arn, grant_schema)
logger.info(f"Grant Schema Response: {grant_schema_res}")
create_table = f"""
CREATE TABLE {table_name} ({primary_key_field} uuid PRIMARY KEY, {vector_field} vector(1536), {text_field} text, {metadata_field} json)
"""
create_table_res = execute_statement(resource_arn, database_name, secret_arn, create_table)
logger.info(f"Create Table Response: {create_table_res}")
grant_table = f"""
GRANT ALL ON TABLE {table_name} TO {user_name};
"""
grant_table_res = execute_statement(resource_arn, database_name, secret_arn, grant_table)
logger.info(f"Grant Table Response: {grant_table_res}")
create_index = f"""
CREATE INDEX on {table_name} USING hnsw ({vector_field} vector_cosine_ops);
"""
create_index_res = execute_statement(resource_arn, database_name, secret_arn, create_index)
logger.info(f"Create Index Response: {create_index_res}")
cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
if event['RequestType'] == 'Update':
cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
if event['RequestType'] == 'Delete':
cfnresponse.send(event, context, cfnresponse.SUCCESS, {})
except Exception as e:
cfnresponse.send(event, context, cfnresponse.FAILED, {'Message': str(e)})
今回のテンプレートの肝となる箇所です。
Knowledge Base for Amazon BedrockのベクターストアとしてAuroraを利用するためには、以下の通りSQLによるベクターストアのセットアップを実施する必要があります。
こちらはCloudFormationでは対応していないため、カスタムリソースを使用してリクエストタイプがCreateとなる時のみ、セットアップ用のLambda関数を実行する事で、ベクターストアのセットアップを自動化します。
カスタムリソース及びリクエストタイプについては、以下の記事をご参照ください。
またセットアップ用のLambda関数では、Auroraに対してHTTP経由でアクセスを行える機能であるRDS Data APIを使用します。
RDS Data APIについては、以下の記事をご参照ください。
カスタムリソースとして定義したLambda関数では、以下のセットアップ処理に関するSQLをRDS Data APIを介して実行します。
- ベクターストア用の拡張機能の作成
- ユーザーの作成
- スキーマの作成
- スキーマに対する権限をユーザーに付与
- テーブルの作成
- テーブルに対する権限をユーザーに付与
- インデックスの作成
なお公式ドキュメントでは「テーブルに対する権限をユーザーに付与」処理はありませんが、ない場合Knowledge Base for Amazon Bedrock作成時にエラーが生じるため追加しています。
Knowledge Base for Amazon Bedrock
BedrockKnowledgeBaseRole:
Type: AWS::IAM::Role
Properties:
RoleName: !Sub ${AWS::StackName}-bedrock-kb-role
AssumeRolePolicyDocument:
Version: '2012-10-17'
Statement:
- Effect: Allow
Principal:
Service: [bedrock.amazonaws.com]
Action: ['sts:AssumeRole']
BedrockKnowledgeBase:
Type: AWS::Bedrock::KnowledgeBase
Properties:
Name: !Sub ${AWS::StackName}-knowledge-base
KnowledgeBaseConfiguration:
Type: VECTOR
VectorKnowledgeBaseConfiguration:
EmbeddingModelArn: !Sub arn:aws:bedrock:${AWS::Region}::foundation-model/
RoleArn: !GetAtt BedrockKnowledgeBaseRole.Arn
StorageConfiguration:
Type: RDS
RdsConfiguration:
CredentialsSecretArn: !Ref SecretForAurora
DatabaseName: !FindInMap [DatabaseMap, DatabaseName, Name]
FieldMapping:
MetadataField: !FindInMap [DatabaseMap, MetadataField, Name]
PrimaryKeyField: !FindInMap [DatabaseMap, PrimaryKeyField, Name]
TextField: !FindInMap [DatabaseMap, TextField, Name]
VectorField: !FindInMap [DatabaseMap, VectorField, Name]
ResourceArn: !GetAtt AuroraCluster.DBClusterArn
TableName: !FindInMap [DatabaseMap, TableName, Name]
DependsOn: SetupAuroraData
BedrockKnowledgeBaseDS:
Type: AWS::Bedrock::DataSource
Properties:
KnowledgeBaseId: !Ref BedrockKnowledgeBase
Name: !Sub ${AWS::StackName}-data-source
DataSourceConfiguration:
Type: S3
S3Configuration:
BucketArn: !GetAtt DataSourceBucket.Arn
最後にKnowledge Base for Amazon Bedrockを作成し、事前に作成したAurora、Secrets Manager、S3を紐づけます。
EmbeddingModelArn
では、Knowledge Base for Amazon Bedrockに対して利用する埋め込みモデルのARNを指定します。
またRdsConfiguration
では、Auroraに対する接続情報及びデータベースのスキーマ情報を指定します。
動作確認
上記のテンプレートをデプロイする事で、Knowledge Base for Amazon Bedrockを利用して正常に文書データを検索できるか確認します。
なお今回は、2024/4時点でKnowledge Base for Amazon Bedrockが利用可能なバージニアリージョンus-east-1
でデプロイを行います。
CloudFormation テンプレートのデプロイ
今回は以下のパラメータで、CloudFormationスタックをデプロイします。
DatabasePassword
では、Auroraに接続するためのパスワードを指定します。
機密情報となるため、適切なパスワードを使用してください。
EmbeddingModelId
では、Knowledge Base for Amazon Bedrockに対して利用する埋め込みモデルのIDを指定します。
今回はamazon.titan-embed-text-v1
を指定します。
なおAuroraのデプロイ時間が長く、およそ20-30分程度かかります。
S3バケットへのデータ投入
次にCloudShell上で以下のコマンドを実施し、PDFファイルを作成したS3バケットにアップロードします。
# S3 バケット名を設定
BUCKET_NAME=<バケット名>
# AWS の公式ドキュメントの PDF ファイルをダウンロード
mkdir -p pdf
cd pdf
wget https://docs.aws.amazon.com/ja_jp/amazondynamodb/latest/developerguide/dynamodb-dg.pdf -O DynamoDB.pdf
wget https://docs.aws.amazon.com/ja_jp/lambda/latest/dg/lambda-dg.pdf -O Lambda.pdf
wget https://docs.aws.amazon.com/ja_jp/vpc/latest/userguide/vpc-ug.pdf -O VPC.pdf
wget https://docs.aws.amazon.com/ja_jp/bedrock/latest/userguide/bedrock-ug.pdf -O Bedrock.pdf
cd ..
# S3 バケットに PDF ファイルをアップロード
aws s3 cp pdf s3://${BUCKET_NAME} --recursive
ファイルのアップロードが成功すると、以下の通り作成したS3バケットに4つのPDFファイルが作成されます。
Knowledge Baseデータソースの同期
続いてBedrockのコンソールに移動し、左ペインの"ナレッジベース"に切り替え、作成したナレッジベースをクリックします。
ナレッジベース画面中部あたりのデータソースに移動し、同期をクリックします。
同期中は、ステータスがSyncingになります。
データソースの同期にも、およそ20-30分程度かかります。
同期が完了すると、ステータスがReadyに戻り、最終同期時刻が表示されます。
Knowledge Baseでの検索実行
最後に検索が可能か確認するために、ナレッジベース画面のテストをクリックし、ナレッジベーステストを実施します。
例えば"回答を生成"をOFFにした状態で、「IPアドレス」と入力し送信すると、該当の用語に類似するデータソースのチャンクが表示されます。
また"回答を生成"をONにした状態で、「Lambdaに割り当てられる最大メモリを教えてください」と入力し送信すると、データソースのチャンクを元に正しく「10,240MB」という回答を返します。
これでナレッジベースが正しく動作している事を確認できました!
最後に
今回はベクターストアにAuroraを利用したKnowledge Base for Amazon Bedrockを、CloudFormation一撃で作成可能なテンプレートについて紹介しました。
今までKnowledge Baseはコンソールから作成する必要がありましたが、CloudFormationによるデプロイが可能となった事で、より簡単にKnowledge Baseを作成できるようになっています。
Knowledge Base for Amazon Bedrockをまだ触った事がない方は、ぜひ一度お試しください!
以上、つくぼし(tsukuboshi0755)でした!
参考文献
テンプレートの作成にあたって、以下のブログを参考にさせて頂きました。