【コピペでRAG構築】Knowledge Base for Amazon Bedrock(Aurora Serverless v2 for PostgreSQL)

こんにちは。AWS CLIが好きな福島です。

はじめに

今回は、コピペでKnowledge Base for Amazon Bedrockを利用したRAGを構築する方法をご紹介します。ベクターストアとしては、Aurora Serverless v2 for PostgreSQLを利用するかつ夜間停止することでできる限りコストを抑えます。

構築には30分程度かかりますが、半分以上はAurora PostgreSQLの構築が完了するのを待っている時間です。

また、Knowledge Base for Amazon Bedrockの詳細について知りたい方は、以下のブログをご覧ください。

blog.serverworks.co.jp

補足

主なAurora Serverless v2の利用料

Auroraは、ストレージやData APIの利用料などコストがかかる点は複数ありますが、主に発生するACU時間あたりのコストを以下に記載します。

ACU 時間あたり、USD 0.12の利用料が発生しますが、最低0.5ACUで作成することができるため、月額$43.2(0.12 * 0.5(ACU) * 24(時間) * 30(日))、1ドル150円換算で月額6480円になります。また、Auroraの夜間停止によりさらにコストを抑えることができます。

aws.amazon.com

参考

今回は以下のドキュメントを参考に構築しました。

docs.aws.amazon.com

構成図

今回構築するリソースは以下の通りです。またリージョンはKnowledge Base for Amazon Bedrockがサポートしているオレゴンを利用します。

流れ

今回やることの流れは以下の通りです。

①CloudFormationのデプロイ
②Aurora Serverlessのセットアップ
③Knowledge Base for Amazon Bedrockの構築
④S3にドキュメントのアップロード
⑤S3内のドキュメントをAurora Serverlessに同期
⑥動作確認

補足ですが、CloudFormationで作成するリソースは以下の通りです。 Knowledge Base for Amazon Bedrockは、CloudFormationをサポートしていないため、AWS CLIで構築します。

VPC, Subnet, RouteTable
Aurora Serverless
SecretsManager * 2(Admin/Bedrockデータベースユーザー用)
IAM Role * 2(Bedrock用, EventBridge Scheduler用)
S3
EventBridge Scheduler * 2(Aurora停止用/起動用)

⓪事前準備

まずは必要な変数を設定します。用途は変数名の通りになりますが、特に希望がなければ以下をそのまま実行すればOKです。

STACK_NAME=test-kb-using-aurora
KNOWLEDGE_BASE_NAME=test-knowledge-base
DATA_SOURCE_NAME=test-knowledge-base-data-source
REGION=us-west-2
EMBEDDING_MODEL_ARN="arn:aws:bedrock:${REGION}::foundation-model/amazon.titan-embed-text-v1"
TEXT_MODEL_ARN="arn:aws:bedrock:${REGION}::foundation-model/anthropic.claude-v2:1"

また、Amazon Bedrockのモデルの有効化ができていない場合は、以下から実施しておきます。 (唯一、ここだけはコマンドラインから操作できないため、マネジメントコンソールから実施する必要があります。)

https://us-west-2.console.aws.amazon.com/bedrock/home?region=us-west-2#/modelaccess

①CloudFormationのデプロイ

Gitからテンプレートをクローンします。

git clone https://github.com/kazuya9831/blog-sample.git

CloudFormationをデプロイします。

aws cloudformation deploy \
--stack-name ${STACK_NAME} \
--template-file blog-sample/kb-for-amazon-bedrock-using-aurora/template.yml \
--capabilities CAPABILITY_NAMED_IAM \
--region ${REGION}

デプロイには15分程度かかるため、しばらく待ちます。

デプロイ待ち

デプロイを待つ間に構築するリソースのポイントを簡単に説明します。

データベース名

データベース名は、ragとしています。

https://github.com/kazuya9831/blog-sample/blob/main/kb-for-amazon-bedrock-using-aurora/template.yml#L153

Bedrockのデータベースユーザー用Secrets Manager

「②Aurora Serverlessのセットアップ」の中でBedrock用のデータベースユーザーを作成しますが、 Knowledge Base for Amazon Bedrockから作成したユーザーを利用するためには、認証情報をSecrets Managerに保存する必要があります。

そのため、事前にSecrets Managerに作成予定のユーザーの名前とパスワードを設定しています。

https://github.com/kazuya9831/blog-sample/blob/main/kb-for-amazon-bedrock-using-aurora/template.yml#L138-L139

本来、認証情報はハードコーティングするべきではないため、その点ご留意ください。

Adminのデータベースユーザー用Secrets Manager

Adminのデータベースユーザー用Secrets Managerは、Auroraの作成と同時に作成されます。そのため、テンプレートにはAdminのデータベースユーザー用のSecrets Managerを構築するコードが含まれていません。

https://github.com/kazuya9831/blog-sample/blob/main/kb-for-amazon-bedrock-using-aurora/template.yml#L158

Auroraのキャパシティユニット

Auroraのキャパシティユニットは、最低0.5、最大1.0で設定しています。 https://github.com/kazuya9831/blog-sample/blob/main/kb-for-amazon-bedrock-using-aurora/template.yml#L160-L161

AuroraのData API有効化

Knowledge Base for Amazon BedrockでAuroraを利用するためには、Data APIが必要なため、有効化しています。 Data APIにより、API経由でSQLを実行できるようになります。

https://github.com/kazuya9831/blog-sample/blob/main/kb-for-amazon-bedrock-using-aurora/template.yml#L162

Bedrock用のIAMロール名

Bedrock用のIAMロール名には命名規則があり、「AmazonBedrockExecutionRoleForKnowledgeBase_」をつける必要があります。

https://github.com/kazuya9831/blog-sample/blob/main/kb-for-amazon-bedrock-using-aurora/template.yml#L182

Auroraの自動起動/停止時間

Auroraは、8:00に自動起動、20:00に自動停止するようにしています。

②Aurora Serverlessのセットアップ

必要な情報を設定

以下を実行します。 DATABASE_NAMEとTABLE_NAME以外は、CloudFormationのOutputsから値を取得しています。

DATABASE_NAME="rag"
TABLE_NAME="bedrock_integration.bedrock_kb"

RESOURCE_ARN=$(aws cloudformation describe-stacks \
--stack-name ${STACK_NAME} \
--query "Stacks[].Outputs[?OutputKey=='TestAuroraClusterArn'].OutputValue" \
--output text --region ${REGION})

ADMIN_SECRET_ARN=$(aws cloudformation describe-stacks \
--stack-name ${STACK_NAME} \
--query "Stacks[].Outputs[?OutputKey=='TestAuroraAdminUserSecretArn'].OutputValue" \
--output text --region ${REGION})

BEDROCK_SECRET_ARN=$(aws cloudformation describe-stacks \
--stack-name ${STACK_NAME} \
--query "Stacks[].Outputs[?OutputKey=='TestAuroraBedrockUserSecretArn'].OutputValue" \
--output text --region ${REGION})

ROLE_ARN=$(aws cloudformation describe-stacks \
--stack-name ${STACK_NAME} \
--query "Stacks[].Outputs[?OutputKey=='TestIAMRoleForKnowledgeBase'].OutputValue" \
--output text --region ${REGION})

S3_FOR_KNOWLEDGE_BASE_DATA_SOURCE=$(aws cloudformation describe-stacks \
--stack-name ${STACK_NAME} \
--query "Stacks[].Outputs[?OutputKey=='TestS3Bucket'].OutputValue" \
--output text --region ${REGION})

pgvectorのセットアップ

aws rds-data execute-statement \
--resource-arn ${RESOURCE_ARN} \
--database ${DATABASE_NAME} \
--secret-arn ${ADMIN_SECRET_ARN} \
--region ${REGION} \
--sql "CREATE EXTENSION IF NOT EXISTS vector;"

補足

以下のコマンドでバージョンを確認することができますが、pgvector 0.5.0 以降のバージョンである必要があります。

aws rds-data execute-statement \
--resource-arn ${RESOURCE_ARN} \
--database ${DATABASE_NAME} \
--secret-arn ${ADMIN_SECRET_ARN} \
--region ${REGION} \
--sql "SELECT extversion FROM pg_extension WHERE extname='vector';"

スキーマの作成

aws rds-data execute-statement \
--resource-arn ${RESOURCE_ARN} \
--database ${DATABASE_NAME} \
--secret-arn ${ADMIN_SECRET_ARN} \
--region ${REGION} \
--sql "CREATE SCHEMA bedrock_integration;"

新しいロールの作成

aws rds-data execute-statement \
--resource-arn ${RESOURCE_ARN} \
--database ${DATABASE_NAME} \
--secret-arn ${ADMIN_SECRET_ARN} \
--region ${REGION} \
--sql "CREATE ROLE bedrock_user WITH PASSWORD 'P@ssword123' LOGIN;"

テーブル作成

aws rds-data execute-statement \
--resource-arn ${RESOURCE_ARN} \
--database ${DATABASE_NAME} \
--secret-arn ${ADMIN_SECRET_ARN} \
--region ${REGION} \
--sql "CREATE TABLE bedrock_integration.bedrock_kb (id uuid PRIMARY KEY, embedding vector(1536), chunks text, metadata json);"

権限設定

aws rds-data execute-statement \
--resource-arn ${RESOURCE_ARN} \
--database ${DATABASE_NAME} \
--secret-arn ${ADMIN_SECRET_ARN} \
--region ${REGION} \
--sql "GRANT ALL ON SCHEMA bedrock_integration to bedrock_user;"

aws rds-data execute-statement \
--resource-arn ${RESOURCE_ARN} \
--database ${DATABASE_NAME} \
--secret-arn ${ADMIN_SECRET_ARN} \
--region ${REGION} \
--sql "GRANT ALL ON TABLE bedrock_integration.bedrock_kb TO bedrock_user;"

補足

テーブルへの権限付与はドキュメントには記載されていなかったのですが、以下のエラーが出たため、権限を付与しました。

An error occurred (ValidationException) when calling the CreateKnowledgeBase operation: 
The knowledge base storage configuration provided is invalid... 
The vector database encountered an error while processing the request: 
ERROR: permission denied for table bedrock_kb; 
SQLState: 42501 (Service: RdsData, Status Code: 400, Request ID: 9e529ca4-e284-47e6-976e-57c175f31721)

インデックスの作成

aws rds-data execute-statement \
--resource-arn ${RESOURCE_ARN} \
--database ${DATABASE_NAME} \
--secret-arn ${ADMIN_SECRET_ARN} \
--region ${REGION} \
--sql "CREATE INDEX on bedrock_integration.bedrock_kb USING hnsw (embedding vector_cosine_ops);"

③Knowledge Base for Amazon Bedrockの構築

Knowledge Base for Amazon Bedrockの構築

aws bedrock-agent create-knowledge-base \
--name "${KNOWLEDGE_BASE_NAME}" \
--knowledge-base-configuration 'type=VECTOR,vectorKnowledgeBaseConfiguration={embeddingModelArn='${EMBEDDING_MODEL_ARN}'}' \
--storage-configuration 'type=RDS,rdsConfiguration={resourceArn='${RESOURCE_ARN}',credentialsSecretArn='${BEDROCK_SECRET_ARN}',databaseName='${DATABASE_NAME}',tableName='${TABLE_NAME}',fieldMapping={primaryKeyField=id,vectorField=embedding,textField=chunks,metadataField=metadata}}' \
--role-arn "${ROLE_ARN}" --region ${REGION}

データソースへの追加

Knowledge BaseのIDが必要なため、取得します。

KNOWLEDGE_BASE_ID=$(aws bedrock-agent list-knowledge-bases \
--query "knowledgeBaseSummaries[?name=='${KNOWLEDGE_BASE_NAME}'].knowledgeBaseId" \
--output text --region ${REGION})

データソースを追加します。

aws bedrock-agent create-data-source \
--knowledge-base-id ${KNOWLEDGE_BASE_ID} \
--name ${KNOWLEDGE_BASE_NAME}-data-source \
--data-source-configuration 'type=S3,s3Configuration={bucketArn=arn:aws:s3:::'${S3_FOR_KNOWLEDGE_BASE_DATA_SOURCE}'}' \
--region ${REGION}

④S3にドキュメントのアップロード

架空の人物のプローフィールを作成します。

cat << EOF > profile.txt
名前: 山田太郎
年齢: 28歳
職業: ソフトウェアエンジニア
趣味: プログラミング、登山、旅行

名前: 佐藤花子
年齢: 32歳
職業: グラフィックデザイナー
趣味: 写真撮影、読書、映画鑑賞

名前: 伊藤一郎
年齢: 40歳
職業: 獣医
趣味: サイクリング、ブログ執筆、料理
EOF

作成したドキュメントをS3にアップロードします。

aws s3 cp profile.txt s3://${S3_FOR_KNOWLEDGE_BASE_DATA_SOURCE}

⑤S3内のドキュメントをAurora Serverlessに同期

同期するためには、データソースのIDが必要なため取得します。

DATA_SOURCE_ID=$(aws bedrock-agent list-data-sources \
--knowledge-base-id ${KNOWLEDGE_BASE_ID} \
--query "dataSourceSummaries[?name=='${DATA_SOURCE_NAME}'].dataSourceId" \
--output text --region ${REGION})

同期コマンドを実行します。次のコマンドでステータスを確認するため、INGESTION_IDを取得します。

INGESTION_ID=$(aws bedrock-agent start-ingestion-job  \
--knowledge-base-id ${KNOWLEDGE_BASE_ID} \
--data-source-id ${DATA_SOURCE_ID} \
--query ingestionJob.ingestionJobId \
--output text --region ${REGION} )

同期が完了したか確認します。

aws bedrock-agent get-ingestion-job  \
--knowledge-base-id ${KNOWLEDGE_BASE_ID} \
--data-source-id ${DATA_SOURCE_ID} \
--ingestion-job-id ${INGESTION_ID} \
--query ingestionJob.status \
--output text --region ${REGION}

⑥動作確認

質問

aws bedrock-agent-runtime retrieve-and-generate \
--region ${REGION} \
--retrieve-and-generate-configuration type="KNOWLEDGE_BASE",knowledgeBaseConfiguration="{knowledgeBaseId=${KNOWLEDGE_BASE_ID},modelArn=${TEXT_MODEL_ARN}}" \
--input text=山田太郎って何歳？

実行結果例

$ aws bedrock-agent-runtime retrieve-and-generate \
--region ${REGION} \
--retrieve-and-generate-configuration type="KNOWLEDGE_BASE",knowledgeBaseConfiguration="{knowledgeBaseId=${KNOWLEDGE_BASE_ID},modelArn=${TEXT_MODEL_ARN}}" \
--input text=山田太郎って何歳？
{
    "sessionId": "025178da-ee46-45d2-a063-068c44b28af7",
    "output": {
        "text": "山田太郎の年齢は28歳です。"
    },
    "citations": [
        {
            "generatedResponsePart": {
                "textResponsePart": {
                    "text": "山田太郎の年齢は28歳です。",
                    "span": {
                        "start": 0,
                        "end": 13
                    }
                }
            },
            "retrievedReferences": [
                {
                    "content": {
                        "text": "名前: 山田太郎 年齢: 28歳 職業: ソフトウェアエンジニア 趣味: プログラミング、登山、旅行  名前: 佐藤花子 年齢: 32歳 職業: グラフィックデザイナー 趣味: 写真撮影、読書、映画鑑賞  名前: 伊藤一郎 年齢: 40歳 職業: 獣医 趣味: サイクリング、ブログ執筆、料理"
                    },
                    "location": {
                        "type": "S3",
                        "s3Location": {
                            "uri": "s3://test-knowledge-base-data-sources-bucket-271409176351/profile.txt"
                        }
                    }
                }
            ]
        }
    ]
}

検索

aws bedrock-agent-runtime retrieve \
--region ${REGION} \
--knowledge-base-id ${KNOWLEDGE_BASE_ID} \
--retrieval-query text=山田太郎って何歳？

実行結果例

$ aws bedrock-agent-runtime retrieve \
--region ${REGION} \
--knowledge-base-id ${KNOWLEDGE_BASE_ID} \
--retrieval-query text=山田太郎って何歳？

{
    "retrievalResults": [
        {
            "content": {
                "text": "名前: 山田太郎 年齢: 28歳 職業: ソフトウェアエンジニア 趣味: プログラミング、登山、旅行  名前: 佐藤花子 年齢: 32歳 職
業: グラフィックデザイナー 趣味: 写真撮影、読書、映画鑑賞  名前: 伊藤一郎 年齢: 40歳 職業: 獣医 趣味: サイクリング、ブログ執筆、料理"
            },
            "location": {
                "type": "S3",
                "s3Location": {
                    "uri": "s3://test-knowledge-base-data-sources-bucket-271409176351/profile.txt"
                }
            },
            "score": 0.32750651774536255
        }
    ]
}

終わりに

今回は、コピペでRAGを構築する方法をご紹介しました。

どなたかのお役に立てれば幸いです。