Knowledge Base for Amazon Bedrock のベクターストアとして Aurora PostgreSQL が利用可能に！ #AWS

この記事は セゾン情報システムズ Advent Calendar 2023 シリーズ 2, 23 日目の記事でず。

はじめに

2023/12/21 のアップデートで Knowledge Base for Amazon Bedrock のベクターストアとして Amazon Aurora PostgreSQL が選択可能になりました🎉

上記のアップデートには以下の 2 つが含まれていますが、本記事では後者について扱います。

Amazon Aurora ML から SQL で Amazon Bedrock 上の基盤モデルにアクセス可能に
Knowledge Base for Amazon Bedrock のベクターストアとして Aurora PostgreSQL を選択可能に

Knowledge Base for Amazon Bedrock は任意の S3 バケットに検索対象とさせたいオブジェクトを配置するだけで、Amazon Bedrcok 上でマネージドな RAG の構成を簡単に実装できる機能です。そのほかの詳細や設定手順については以下の記事で触れていますので、もしよろしければご参照ください。

Aurora PostgreSQL の価格面のメリット

これまでは以下の 3 つのデータベースがサポートされていました。

Vector Engine for Amazon OpenSearch Servelerless
Pincone
Redis Enterprise Cloud

Pincone および Redis Enterprise Cloud は AWS のサービスではありません。AWS のサービスを利用する必要がある場合は、Vector Engine for Amazon OpenSearch Serverless が唯一の選択肢でした。OpenSearch Serverless は Knowledge Base の Quick Setup から簡単に作成できる一方で、小規模利用におけるコストに課題があります。

OpenSearch Serverless のコンピューティング料金は OpenSearch Compute Unit (OCU) の使用量によって請求されます。2023/12/24 時点でバージニア北部リージョンの料金は $0.24 per OCU per hour です。

開発/テストモードを有効化 (スタンバイレプリカを無効化) した場合でも、最低 2 OCU が請求されるため　0.24 × 2 × 730 = $350.4/月程度の料金が発生します。

Aurora PostgreSQL では Serverless v2 の最小容量である 0.5 ACU のケースで 0.12 × 0.5 × 730 = $43.8/月程度の料金になります。

プロビジョンドインスタンスの場合、db.r6g.large で $189.80/月の料金ですが、使用していない間は最大 7 日間インスタンスの停止が行えるため、更に料金を節約することも可能です。

前提条件

Aurora PostgreSQL クラスターを Knowledge Base for Amazon Bedrock で使用するには以下の前提条件を満たす必要があります。

エンジンバージョン
- 15.4 以降
- 14.9 以降
- 13.12 以降
- 12.16 以降
pgvector 拡張機能のバージョン
- 0.5.0 以降
Data API の有効化
Secrets Manager で管理されるユーザー情報による接続

Knowledge base から Aurora への接続は Data API を介しておこなわれます。え、Data API って Aurora Serverless v1 だけでは！？と思われた皆さん、ご安心ください。同日のアップデートで Data API が Aurora PostgreSQL のプロビジョンドインスタンスおよび Serverless v2 でもサポートされました。

プロビジョンドインスタンスの場合、バースト可能クラス (t 系) では Data API はサポートされていませんのでご注意ください。

Limitations with RDS Data API

Data API isn't supported on T DB instance classes.

やってみる

Aurora PostgreSQL クラスターの作成

今回は以下の設定で新規クラスターを作成しました。

エンジンバージョン: 15.5
認証情報の設定: AWS Secrets Manager でマスター認証情報を管理
DB インスタンスクラス: Serverless v2
容量の範囲: 最小 0.5 ACU、最大 1 ACU
マルチ AZ 配置: Aurora レプリカを作成しない
パブリックアクセス: あり
モニタリング: Perfomance Insights オフ、拡張モニタリング無効化
RDS Data API: 有効化

Knowledge base からアクセスにはパブリックアクセスが有効化されている必要があります。VPC セキュリティグループでは HTTPS (TCP 443) のインバウンド接続を許可します。

Data API の有効化をお忘れなく。

Aurora PostgreSQL クラスターの設定

コンソールのクエリエディタを使用してマスターユーザーでデフォルトの postgres データベースに接続します。

pgvector 拡張機能をインストールします。

-- pgvegtor 拡張機能のインストール
CREATE EXTENSION IF NOT EXISTS vector;

-- pgvector のバージョン確認
SELECT extversion FROM pg_extension WHERE extname='vector';

pgvector のバージョンは 0.5.0 以降である必要があります。

スキーマおよびロールを作成し、権限を付与します。

-- スキーマの作成
CREATE SCHEMA bedrock_integration;

-- ロールの作成
CREATE ROLE bedrock_user WITH PASSWORD 'your_password' LOGIN;

-- スキーマに対する権限を付与
GRANT ALL ON SCHEMA bedrock_integration to bedrock_user;

AWS Secrets Manager で bedrock_user の RDS データベースの認証情報を作成します。

クエリエディタで bedrock_user の認証情報を使用して postgres データベースに接続し、テーブルおよびインデックスを作成します。

-- テーブルの作成
CREATE TABLE bedrock_integration.bedrock_kb (id uuid PRIMARY KEY, embedding vector(1536), chunks text, metadata json);

-- インデックスの作成
CREATE INDEX on bedrock_integration.bedrock_kb USING hnsw (embedding vector_cosine_ops);

Knowledge base の作成

Bedrock コンソールの Orchestration から Knowledge base を選択し、Create knowledge base へ進みます。

Step 1 で Knowledge base の名前および詳細を入力します。ここではすべてデフォルト値のまま次へ進みます。サービスロールも新規作成を選択します。

Step 2 でデータソースの配置先となる S3 バケットを指定します。オブジェクトがカスタマー管理の KMS キーで暗号化されている場合は Add customer-managed KMS key for S3 data にチェックを入れます。

Step 3 でベクターストアの設定を行います。Choose a vector store you have created から　Amazon Aurora を選択します。

Aurora Cluster の接続情報を入力します。

Amazon Aurora DB Cluster ARN: arn:aws:rds:<region>:<account_id>:cluster:<your_cluster>
Database name: postgres
Table Name: bedrock_integration.bedrock_kb
Secret ARN: arn:aws:secretsmanager:<region>:<account_id>:secert:<your_secert>
Customer managed KMS key for your secret: secret をカスタマー管理キーで暗号化している場合に KMS キーの ARN を指定

Index field mapping で次の情報を入力します。(bedrock_kb テーブルの定義とマッピングします。)

Vector field: embedding
Text field: chunks
Bedrock-managed metadata field: metadata
Primary key: id

Step 4 で入力内容を確認し、Knowledge base の作成を完了します。Knowledge base の作成が完了したら最初に S3 バケット上に配置したデータを同期する必要があるため、Sync をクリックします。

今回はあらかじめ PDF 版のAmazon Bedrock のユーザーガイド (英語) を S3 バケットに格納しています。

動作確認

コンソールの Test Knowledge base から Bedrock のユーザーガイドを元に回答が生成されることが確認できました。

クエリエディタに接続し、bedrock_kb テーブルを参照してみると、ベクトルデータが格納されていることがわかります。

select * from bedrock_integration.bedrock_kb limit 10;

参考

以上です。
参考になれば幸いです。