Amazon Bedrock Knowledge Basesのベクトルストアとして低価格なS3 Vectorsが利用可能となり、AWS環境でとても手軽にRAGが構築できるようになりました。
私のチームが最近実施したあるPoCでも、プロジェクト管理ツールのチケットや業務マニュアル等をソースとしたRAGを多数構築し、Strands AgentsベースのAIエージェント(実行環境はBedrock AgentCore Runtime)に連結して、様々なユースケース(類似事例検索や作業工数見積もり等)への応用可能性を検証する機会があり、S3 Vectorsをフル活用しました。
本記事は、このPoCで得られた知見のなかから、Knowledge Base + S3 VectorsによるRAG作成を計画する時に押さえておきたい情報を整理して紹介します。紹介する内容は次の通りです。
- S3 Vectorsインデックスは事前作成し、その際に「フィルタリング不可能」なメタデータを定義しておく
- S3 VectorsのAPIで追加・更新したメタデータもKnowledge Baseから参照できる
- Strands AgentsのAIエージェントに2つ以上のKnowledge Baseを連結したい場合はカスタムツールを実装する
1. S3 Vectorsインデックスは事前作成し、その際に「フィルタリング不可能」なメタデータを定義しておく
S3 Vectorsは、ベクトルバケット内にインデックスがあり、インデックス内にベクトルがあるという構造になっていて、Knowledge Baseのうつわ(ベクトルストア)がインデックス、個々のチャンク(Knowledge Baseに保管するテキストデータ)を保持するのがベクトルです。
ベクトルはひとつの「ベクトルキー」、ひとつの「ディメンション」(チャンクをエンベディング処理して得られる数値群(行列))、複数の「メタデータ」で構成され、チャンクはメタデータになります。
メタデータには「フィルタリング不可能」と「フィルタリング可能」の2種類があります。この2種類の特徴は以下のとおりです。
| フィルタリング不可能なメタデータ | フィルタリング可能なメタデータ | |
|---|---|---|
| 作成時の考慮点 | インデックス作成時にフィルタリング不可能なメタデータのキー名を定義する(インデックス作成後の追加は不可) | 左記以外のメタデータはすべてフィルタリング可能として扱われる |
| 容量制限 | 40KB(ベクトルあたりのメタデータの総量で、フィルタリング不可能、フィルタリング可能を含む) | 2KB(ベクトルあたりのフィルタリング可能なメタデータの総量) |
| メタデータキーの数の上限 | 50個(フィルタリング不可能、フィルタリング可能を含む) | 10個 |
特に容量制限の観点から、メタデータはフィルタリングを使用しない限り「フィルタリング不可能」で作成しておいた方がよいことになります。
Knowledge Baseはベクトルの中にAMAZON_BEDROCK_TEXTというキー名のチャンク保持用のメタデータと、AMAZON_BEDROCK_METADATAというキー名の管理情報保持用のメタデータを使用しますが、以下のドキュメントではこれらをフィルタリング不可能なメタデータとして作成する手順を記述しています。手順を省略すると上限2KBのフィルタリング可能なメタデータとして扱われてしまいます。
Knowledge Base用に作成したベクトルストアを使用するための前提条件
ディメンションには反映したくないけど、検索結果の返り値には含めたい情報があるとします。この場合、情報をチャンクに含めるとエンベディング処理によってディメンションに反映してしまうので、チャンクとは別の場所に保持する必要があります。解決策としては任意のキー名のメタデータを作成し、検索結果の返り値にメタデータも含めるようにします。メタデータは容量等の制約回避のためフィルタリング不可能で作成したいことが多いと思いますが、フィルタリング不可能なメタデータはインデックス作成時にしかキー名を定義できないので、インデックス作成前に設計しておくことが肝要です。
Knowledge Base作成時に、S3 Vectorsインデックスを自動作成するオプションがあります。このオプションではKnowledge BaseはAMAZON_BEDROCK_TEXTとAMAZON_BEDROCK_METADATAをフィルタリング不可能なメタデータとして作成してくれます。しかしフィルタリング不可能なメタデータとして作成してくれるのはこの二つだけです。つまり自動作成に任せると、後述の.metadata.json等で追加する任意のメタデータはすべてフィルタリング可能なものになってしまい、容量や個数の制約を受けることになります。任意のメタデータをフィルタリング不可能として作成したい場合は、あらかじめS3 Vectorsインデックスを作成してメタデータを定義しておき、Knowledge Base作成時にそのインデックスを指定しましょう。
2. S3 VectorsのAPIで追加・更新したメタデータもKnowledge Baseから参照できる
Knowledge Baseでメタデータを付加したい場合、.metadata.jsonファイルを使用する方法が用意されています。
例えばS3バケットにあるdata01.txtというファイルをKnowledge Baseに取り込む場合、同じ場所にdata01.txt.metadata.jsonというファイルを配置してメタデータの定義(メタデータのキー名と値など)を記述しておくと、ファイル取り込み時にメタデータを作成してくれます。以下のAWSブログにこの方法の詳細な説明があります。
Knowledge Bases for Amazon Bedrock がメタデータフィルタリングをサポートし検索精度向上
作成されたメタデータは、.metadata.jsonに指定されたキー名と値のまま、S3 Vectorsインデックスの各ベクトルのメタデータになります。であれば、S3 Vectorsインデックス側で追加したメタデータも、.metadata.jsonファイルで定義したメタデータであるかのようにKnowledge Base側から参照できるということです。
動作確認をした時のpythonコードの一部(メタデータを追加する部分)を抜粋して以下に添付します。S3 Vectorsインデックスからベクトルを取得し、メタデータを追加して上書きして、Knowledge Baseのretrieveの応答にそのメタデータが含まれていることを確認しました。
response = s3vectors.get_vectors(
vectorBucketName=VECTOR_BUCKET_NAME,
indexName=INDEX_NAME,
keys=[key],
returnData=True,
returnMetadata=True
)
vector = response['vectors'][0]
(中略)
s3vectors.put_vectors(
vectorBucketName=VECTOR_BUCKET_NAME,
indexName=INDEX_NAME,
vectors=[{
'key': vector['key'],
'data': vector['data'],
'metadata': metadata
}]
)
3. Strands AgentsのAIエージェントに2つ以上のKnowledge Baseを連結したい場合はカスタムツールを実装する
最後はS3 Vectorsに限らないのですが、今回のPoCでStrands AgentsベースのAIエージェントに複数個のKnowledge Baseを連結する要件があり、そのために工夫した点を紹介します。
Strands Agentsには「ツール」という部品群があり、ツールを使ってAIエージェントを簡単に機能拡張できます。Knowledge Base連結用の"retrieve"というツールも標準で用意されています。(参照:https://github.com/strands-agents/tools)
使い方は次の通りです。
- 環境変数
STRANDS_KNOWLEDGE_BASE_IDにKnowledge Base IDをセット - Agentインスタンス作成時に、toolsの配列にretrieveを記述(例:
agent = Agent(tools=[retrieve])) - 環境変数
RETRIEVE_ENABLE_METADATA_DEFAULT=trueをセットすればメタデータも取得可能
ここで注意点ですが、Knowledge Base IDをSTRANDS_KNOWLEDGE_BASE_IDという環境変数で指定するため、この実行環境で使用できるKnowledge Baseはひとつだけになります。したがってAIエージェントに複数のKnowledge Baseを連結したい場合は、Strands Agents標準のretrieveツール以外の手段が必要です。
幸いStrands Agentsでは、ユーザー独自のカスタムツールも簡単に実装できます。以下に示すように、関数定義の前に@toolデコレータを配置し、定義の先頭のコメントでツールの説明を記述すればツールになります。ロジックのところにKnowledge Baseをretrieveする処理を記述すれば、Knowledge Base用のカスタムツールの出来上がりです。
@tool
def 関数名(引数, ・・・):
"""
ツールの説明を記述(AIエージェントがツール選択に利用)
"""
ロジック
return 返り値
まとめ
S3 VectorsとKnowledge Baseを組み合わせたRAG構築において、メタデータの事前設計、S3 Vectors APIによる柔軟なメタデータ管理、複数Knowledge Base連結のためのカスタムツール実装という3つの実践的な知見を紹介しました。これらはいずれも公式ドキュメントだけでは気づきにくい、実際のPoCを通じて得られた考慮点です。本記事がRAGを連結したAIエージェント構築の参考になれば幸いです。