AWS Bedrock と Aurora でRAGを作る
本記事の概要
前置き
AWS の Bedrock には Knowledge Bases という機能がある。
同機能を使えば手軽に RAG の仕組みを構築できる。
Bedrock 上で管理されている LLM や Embedding モデルを用いつつ、データストア(ベクトルDB)は別で準備する形となる。
そこら辺は下記でまとめてある。
上記資料に記載の事情から、ベクトルDBは OpenSearch でも Pinecone でもなく Aurora Serverless V2 PostgreSQL を用いることとなった。
本資料の記載範囲
実際に、Bedrockその他の機能を用いてRAG機能のベースを構築するところまでのアレヤコレヤをまとめる。
構築に際してはTerraformを用いており、Terraformのコードもここに記載する。
構成
すごく簡単に、構成図
専用S3バケット(データソース)を準備。ここにナレッジの素材となるファイルをアップロードする
Knowledge Base は、同期処理を行うと同データソース内のデータを取得する
そのデータをチャンクに分解し、Embeddingモデルを用いてベクトルデータ化する
このベクトルデータをAurora Serverless V2 PostgreSQL内に作成した専用テーブルに格納する
アプリケーションからの問い合わせ文字列に対して、ベクトルDB検索を行い、結果を返す
といった流れになる。DB接続時にはSecrets Manager に保存されている username, password を用いる。
参考にした記事
こちらを参考にさせてもらった。
公式も参照
構築
Terraform ベースで構築内容を説明
※前提条件
Aurora Serverless V2 PostgreSQL と S3 バケットは作成済みとする。
開発作業の都合上、PostgreSQL はパブリック環境からの接続が可能(パブリックサブネットに配置)となっている。
プライベートサブネット配備の場合にBedrockからPostgreSQLに接続するために追加でなにか設定が必要か否かは未確認
main.tf
main.tfは何もしていなくて、各リソース作成を呼び出したり引数を渡したりしているだけ。
# ##################################
# 基本設定
terraform {
required_providers {
aws = {
source = "hashicorp/aws"
# バージョンが古いと Knowledge Bases に対応していないので注意
version = "5.76.0"
}
}
}
provider "aws" {
region = "ap-northeast-1"
}
# ##################################
# ##################################
# 各種リソース構築
# ##################################
# 作成済みのAurora内のDB構築
module "aurora_initialize_db" {
source = "./modules/aurora_initialize_db"
}
# Bedrock から ベクトルDB(PostgreSQL)に繋ぐためのシークレット
module "secret" {
source = "./modules/secret"
}
# Bedrockで必要となるIAMロール
module "bedrock_iam" {
source = "./modules/bedrock_iam"
}
# Bedrock周り
module "bedrock" {
source = "./modules/bedrock"
secret_arn = module.secret.db_conn_secret_arn
bedrock_iam_arn = module.bedrock_iam.bedrock_iam_role_arn
bedrock_attach_policies = module.bedrock_iam.bedrock_attach_policies
depends_on = [module.bedrock_iam]
}
aurora_initialize_db.tf
作成済みのAurora Serverless V2 PostgreSQL のDB, テーブル等の作成。
./modules/aurora_initialize_db.aurora_initialize_db.tf
※接続情報などは適宜外出しで変数管理すべし
ポイント
pgvector(PostgreSQLのベクトルストア機能)を有効化
Knowledge Base から利用されるベクトルストア用のテーブルを事前作成
# PostgreSQLに接続してDB作成
resource "null_resource" "create_db_sql" {
provisioner "local-exec" {
command = <<EOT
psql -h <ホスト名> \
-U <ユーザー名> \
-d "postgres" \
-c "
DO \$\$
BEGIN
IF NOT EXISTS (SELECT 1 FROM pg_database WHERE datname = '<DB名>') THEN
PERFORM dblink_exec('dbname=postgres', 'CREATE DATABASE <DB名>');
END IF;
END
\$\$;
"
EOT
environment = {
PGPASSWORD = <パスワード>
}
}
}
# 普通はアプリ専用のユーザーを作成して権限も付与すべきだけど省略
# PostgreSQLに接続してpgvector有効化
resource "null_resource" "create_extention_sql" {
provisioner "local-exec" {
command = <<EOT
psql -h <ホスト名> \
-U <ユーザー名> \
-d "<DB名>" \
-c "create extension if not exists vector;"
EOT
environment = {
PGPASSWORD = <パスワード>
}
}
}
# PostgreSQLに接続して、
# BedrockからKnowledge Basesのベクトルストアとして利用するためのテーブルを作成
# テーブル名や列名は Knowledge Base 構築時にも用いる(同じ文字列を使う)必要あり
resource "null_resource" "create_kb_table" {
provisioner "local-exec" {
command = <<EOT
psql -h <ホスト名> \
-U <ユーザー名> \
-d "<DB名>" \
-c "
CREATE TABLE IF NOT EXISTS bedrock_kb (
id uuid PRIMARY KEY,
embedding vector(1024),
chunks text,
metadata json
);
"
EOT
environment = {
PGPASSWORD = <パスワード>
}
}
}
secret.tf
Knowledge Base から PostgreSQL に接続するための認証情報をシークレットマネージャに登録。
./modules/secret/secret.tf
# 本来はベタ書きしない(別途変数として管理すべし
locals {
pg_conn_user_auth = {
username = <ユーザー名>
password = <パスワード>
}
}
resource "aws_secretsmanager_secret" "db_conn_secret" {
name = "db-conn-secret"
}
resource "aws_secretsmanager_secret_version" "latest" {
secret_id = aws_secretsmanager_secret.db_conn_secret.id
secret_string = jsonencode(local.pg_conn_user_auth)
}
# 後続の処理で利用する情報を output
output "db_conn_secret_arn" {
value = aws_secretsmanager_secret.db_conn_secret.arn
}
bedrock_iam.tf
Knowledge Base で処理を実行するための専用IAMロールを作成。
※前述の構成図ではロール周りの記載を割愛
./modules/bedrock_iam/bedrock_iam.tf
# ロール作成
resource "aws_iam_role" "bedrock_exec_role_for_knowledgeBase" {
name = "bedrock_exec_role_for_knowledgeBase"
assume_role_policy = jsonencode({
Version = "2012-10-17",
Statement = [
{
Sid = "AmazonBedrockKnowledgeBaseTrustPolicy",
Effect = "Allow",
Principal = {
Service = "bedrock.amazonaws.com"
},
Action = "sts:AssumeRole",
Condition = {
StringEquals = {
# AWSアカウントID。通常は変数化して外出しすべし
"aws:SourceAccount" = "XXXXXXXXXXXX"
},
ArnLike = {
"aws:SourceArn" = "arn:aws:bedrock:ap-northeast-1:XXXXXXXXXXXX:knowledge-base/*"
}
}
}
]
})
}
# ポリシー作成(以下4つ分)
resource "aws_iam_policy" "bedrock_invoke_policy" {
name = "bedrock_invoke_policy"
policy = jsonencode({
Version = "2012-10-17",
Statement = [
{
Sid = "BedrockInvokeModelStatement",
Effect = "Allow",
Action = [
"bedrock:InvokeModel"
],
Resource = [
"arn:aws:bedrock:ap-northeast-1::foundation-model/cohere.embed-multilingual-v3"
]
}
]
})
}
resource "aws_iam_policy" "bedrock_rds_policy" {
name = "bedrock_rds_policy"
policy = jsonencode({
Version = "2012-10-17",
Statement = [
{
Sid = "RdsDescribeStatementID",
Effect = "Allow",
Action = [
"rds:DescribeDBClusters"
],
Resource = [
"<AuroraクラスターのARN>"
]
},
{
Sid = "DataAPIStatementID",
Effect = "Allow",
Action = [
"rds-data:BatchExecuteStatement",
"rds-data:ExecuteStatement"
],
Resource = [
"<AuroraクラスターのARN>"
]
}
]
})
}
resource "aws_iam_policy" "bedrock_s3_policy" {
name = "bedrock_s3_policy"
policy = jsonencode({
Version = "2012-10-17",
Statement = [
{
Sid = "S3ListBucketStatement",
Effect = "Allow",
Action = [
"s3:ListBucket"
],
Resource = [
"<S3バケットARN>"
],
Condition = {
StringEquals = {
# AWSアカウントID
"aws:ResourceAccount" = ["XXXXXXXXXXXX"]
}
}
},
{
Sid = "S3GetObjectStatement",
Effect = "Allow",
Action = [
"s3:GetObject"
],
Resource = [
"<S3バケットARN>/*"
],
Condition = {
StringEquals = {
# AWSアカウントID
"aws:ResourceAccount" = ["XXXXXXXXXXXX"]
}
}
}
]
})
}
resource "aws_iam_policy" "bedrock_secrets_manager_policy" {
name = "bedrock_secrets_manager_policy"
policy = jsonencode({
Version = "2012-10-17",
Statement = [
{
Sid = "SecretsManagerGetStatement",
Effect = "Allow",
Action = [
"secretsmanager:GetSecretValue"
],
Resource = [
"<シークレットのARN>"
]
}
]
})
}
# ロールとポリシーを紐づけ
resource "aws_iam_role_policy_attachment" "bedrock_attach_policies" {
for_each = {
"bedrock_invoke_policy" = aws_iam_policy.bedrock_invoke_policy.arn
"bedrock_rds_policy" = aws_iam_policy.bedrock_rds_policy.arn
"bedrock_s3_policy" = aws_iam_policy.bedrock_s3_policy.arn
"bedrock_secrets_manager_policy" = aws_iam_policy.bedrock_secrets_manager_policy.arn
}
role = aws_iam_role.bedrock_exec_role_for_knowledgeBase.name
policy_arn = each.value
}
# 後続の処理で利用する情報を output
output "bedrock_iam_role_arn" {
value = aws_iam_role.bedrock_exec_role_for_knowledgeBase.arn
}
output "bedrock_attach_policies" {
value = aws_iam_role_policy_attachment.bedrock_attach_policies
}
bedrock.tf
いよいよKnowledge Base の作成。
./modules/bedrock/bedrock.tf
# ############################
# 前段処理で生成したリソースの値を利用
# ############################
variable "secret_arn" {}
variable "bedrock_iam_arn" {}
variable "bedrock_attach_policies" {}
########################################################
# Knowledge Base を Aurora(PostgreSQL)ベース、cohere の組み込みモデルベースで作成
########################################################
resource "aws_bedrockagent_knowledge_base" "this" {
name = "bedrock-kb"
role_arn = var.bedrock_iam_arn
knowledge_base_configuration {
type = "VECTOR"
vector_knowledge_base_configuration {
# cohere が提供するEmbedding Model を使用
embedding_model_arn = "arn:aws:bedrock:ap-northeast-1::foundation-model/cohere.embed-multilingual-v3"
}
}
storage_configuration {
type = "RDS"
rds_configuration {
credentials_secret_arn = var.secret_arn
resource_arn = "<AuroraクラスターのARN>"
database_name = "DB名" # 事前作成済みのDB
table_name = "bedrock_kb" # 事前作成済みのテーブル
field_mapping {
primary_key_field = "id" # テーブルの各カラム
vector_field = "embedding"
text_field = "chunks"
metadata_field = "metadata"
}
}
}
depends_on = [ var.bedrock_attach_policies ]
}
########################################################
# ナレッジベースが参照するデータソースのS3との紐づけ
########################################################
resource "awscc_bedrock_data_source" "this" {
name = "bedrock-kb-ds"
knowledge_base_id = aws_bedrockagent_knowledge_base.this.id
data_source_configuration = {
type = "S3"
s3_configuration = {
bucket_arn = "<S3バケットのARN>"
}
}
}
確認
これで、前述の構成図のような構成で環境が構築される。
実際にKnowledge Base を見てみる。
AWSコンソール
ナレッジベースが作られている
IAMロールやDBやシークレット、S3、Embeddingモデルなどの情報が紐づけられている。
補足
同期処理
AWSコンソール上に「同期」ボタンがあり、このボタンを押せばS3上のファイルがナレッジベースに同期されるが、スケジュール機能やリアルタイム同期機能はない。
→追加実装が必要
解説👇️
Bedrock が使えない
AWSコンソール上で実際に動作確認しようとすると、エラーが出た。
Your request rate is too high. Reduce the frequency of requests.
各モデルを有効化していても出る。
原因はQuota。デフォルトではQuotaが0に設定されていて、有効化しても使えない。
ここで各Quotaを変更すればいいのだが、AWS側のトラブルで24年11月現在、変更ができないケースがある。
その場合は、問い合わせが必要とのこと。
チャンキング戦略
今回の構築ではChunkの設定は行っていないが、カスタマイズ可能