Knowledge bases for Amazon Bedrock でメタデータフィルタリングを試してみる

Knowledge bases for Amazon Bedrock でメタデータを利用したコンテンツのフィルタリングができるようになりました。このブログではメタデータフィルタリングを利用して特定のドキュメントだけ引っ張り出してみました。
2024.05.03

こんにちは! AWS 事業本部コンサルティング部のたかくに(@takakuni_) です。

かなり前なのですが、 Knowledge bases for Amazon Bedrock でメタデータを利用したコンテンツのフィルタリングができるようになりました。

Amazon Bedrock のナレッジベースでメタデータのフィルタリングのサポートを開始

今回はこの機能を利用して、コンテンツのフィルタリングを行ってみようと思います。

おさらい

いきなり、メタデータフィルタリングを試すのもアレなので、おさらいします。

なぜメタデータフィルタリングが必要なのか

まずは、なぜメタデータフィルタリングが必要なのかをおさらいします。

RAG を組む上でよくあるのがプロンプトを投げた人の属性をもとにコンテンツをフィルタリングしたいケースです。主に以下の 2 点が目的になります。

  • 情報の秘匿性
    • 全員に見れる情報と各部署ごとに見れる情報など秘匿性のレベルが異なる
  • 情報の最適化
    • パーソナライズされた情報を返したい
    • 例:営業向けレベルの回答が欲しいのに、エンジニア向けレベルが返ってきてお腹いっぱいになる
    • 例:多国籍な企業の場合、国ごとに法律が違うため間違った回答になるケース

そんな課題に有効なのがメタデータフィルタリングです。

メタデータフィルタリング機能がなかった頃

メタデータフィルタリング機能がなかった頃にメタデータフィルタリングを行いたい場合、どうしてたのでしょうか? 申し訳ないのですが、私もわかりません....

わからないなりに、以下のようなアーキテクチャが妄想できます。

ナレッジベースごと分ける

一番、簡単ですね。ナレッジベースごとに分けて、ユーザーの属性ごとにクエリの投げ先を変えてあげる構成です。データベースをシングルテナントにするか、マルチテナントにするかでかなりコストが変わりそうですね。

ちなみに 1 リージョンごとのナレッジベースの最大数は 50 でした。

Quotas for Amazon Bedrock - Amazon Bedrock

プロンプトの活用

プロンプトを活用して情報を出し分ける方法です。LLM に依存するため、ハルシネーションが発生する可能性があります。

例えばですが、以下のように各ファイルの先頭行に、秘匿レベルを追記してあげて LLM に「秘匿レベルにマッチしない情報は無視してください」と追記してあげるといいかもしれません。

秘匿レベル:全社

# オフィスの利用方法マニュアル

## 1. 入退室について

- オフィスの利用可能時間は、平日の 9:00 から 18:00 までです。
- 入室の際は、受付で名前と所属をお伝えください。
- 退室時は、必ず施錠し、鍵を受付に返却してください。

または、 Retrieve API のレスポンスデータの中には s3Location キーが含まれるため、アプリ側でキーをもとに情報を出し分け、 LLM にデータを与えることができそうです。

{
    "nextToken": "string",
    "retrievalResults": [
        {
            "content": {
                "text": "string"
            },
            "location": {
                "s3Location": {
                    "uri": "string"
                },
                "type": "S3"
            },
            "metadata": {
                "string": {...}|[...]|123|123.4|"string"|True|None
            },
            "score": 123.0
        },
    ]
}

retrieve - Boto3 1.34.95 documentation

ただし、この方法だと 検索結果 (numberOfResults) に対してフィルタリングを行う順番 になるため、条件がマッチせず LLM に読み込ませるデータが不足するケースも発生します。近似値が高いデータ(ノイズデータを含む)が上位の結果に出たのち、フィルタリングが行われていくイメージです。 numberOfResults を 100 など大きい数字にした場合、データ不足は解決しますが、ノイズデータも合わせて回答文が生成されてしまうような状況が考えられます。

何が嬉しいのか

この辺りを解決してくれるのが、 Knowledge bases for Amazon Bedrock のメタデータフィルタリング機能です。この機能は、データソース内のファイルにメタデータを付与し、付与したメタデータを利用して検索結果の出力を行います。つまり、「同じベクトルデータベースを使える」かつ、「検索とフィルタリングの順番によって起因するデータ不足」を防ぐことができます。

メタデータの付与

各ファイルにメタデータを付与するには、 .metadata.json で終わるメタデータファイルを、コンテンツファイルと同じフォルダ階層においてあげる必要があります。

The file is in the same folder in the Amazon S3 bucket as its associated source document file.

Set up a data source for your knowledge base - Amazon Bedrock

takakuni% tree .
.
├── 就業規則.md
├── 就業規則.md.metadata.json
├── オフィスの利用方法.md
└── オフィスの利用方法.md.metadata.json

なお、メタデータファイルの最大サイズは 10 KB です。

The file size doesn't exceed the quota of 10 KB.

Set up a data source for your knowledge base - Amazon Bedrock

メタデータファイルの形式

メタデータファイルはキーバリュー形式で定義し、文字列、数字、 Bool の 3 つの型をバリューに定義できます。

{
	"metadataAttributes": {
		"target": "company-wide",
		"year": 2024,
		"for_managers": true
	}
}

Set up a data source for your knowledge base - Amazon Bedrock

Amazon Aurora をベクトルデータベースとして利用している場合は、次のように 事前に メタデータ用のカラムを追加する必要があります。

ALTER TABLE bedrock_integration.bedrock_kb
ADD COLUMN IF NOT EXISTS target VARCHAR(100),
ADD COLUMN IF NOT EXISTS year SMALLINT,
ADD COLUMN IF NOT EXISTS for_managers BOOLEAN;

If the vector index for your knowledge base is in an Amazon Aurora database cluster, check that the table for your index contains a column for each metadata property in your metadata files before starting ingestion.

Sync to ingest your data sources into the knowledge base - Amazon Bedrock

ビルトインメタデータ

x-amz-bedrock-kb-source-uri メタデータはメタデータファイルを用意せずに利用可能なビルトインメタデータとして提供されています。

In addition to custom metadata, you can also filter using S3 prefixes (which is a built-in metadata, so you don’t need to provide any metadata files).

Knowledge Bases for Amazon Bedrock now supports metadata filtering to improve retrieval accuracy | AWS Machine Learning Blog

※ ドキュメントには記載がなく AWS Blog からの引用になります。

例えば次のようなフォルダ構造でコンテンツが配置されていると想定します。

takakuni % tree .
.
├── company-wide
│   ├── 就業規則.md
│   ├── 会議室マニュアル.md
│   └── オフィスの利用方法.md
└── department
    ├── engineer
    │   ├── 個人検証用AWSアカウントのセットアップ方法.md
    │   └── 目標設定について.md
    └── sales
        ├── 営業資料.pdf
        └── 目標設定について.md

この状態で以下のフィルタリングを行うと、部署別のドキュメントに対して検索できます。

import os
import boto3

bucket_name = os.environ.get("BUCKET_NAME")
region = os.environ.get("AWS_REGION")
kb_id = os.environ.get("KNOWLEDGE_BASE_ID")
attribute = os.environ.get("USER_ATTRIBUTE")
client = boto3.client("bedrock-agent-runtime", region_name=region)

response = client.retrieve(
    knowledgeBaseId=kb_id,
    retrievalConfiguration={
        "vectorSearchConfiguration": {
            "filter": {
                "startsWith": {
                    "key": "x-amz-bedrock-kb-source-uri",
                    "value": f"s3://{bucket_name}/department/{attribute}/",
                }
            },
            "numberOfResults": 1,
        },
    },
    retrievalQuery={"text": "個人目標の設定方法を教えてください"},
)

※ ただし、 startsWith は、ベクトルデータベースが Open Search Serverless の場合にのみ利用可能です。

This filter is currently only supported for Amazon OpenSearch Serverless vector stores.

RetrievalFilter - Amazon Bedrock

やってみる

それでは、実際にメタデータを書くオブジェクトに設定し、メタデータフィルタリングを試してみます。以下のように Amazon Aurora をベクトルデータベースにした場合のメタデータフィルタリングを試してみます。

以下の Terraform コードを利用してリソースをデプロイします。

フィルタリングを行うデータ

データ構造は以下を想定します。メタデータフィルタリングによって、2024 年度の事業計画がフィルタリングできたらクリアとします。

company-wide
├── 2023年度事業計画.md
├── 2023年度事業計画.md.metadata.json
├── 2024年度事業計画.md
├── 2024年度事業計画.md.metadata.json
├── 就業規則.md
├── 就業規則.md.metadata.json
├── 会議室マニュアル.md
├── 会議室マニュアル.md.metadata.json
├── オフィスの利用方法.md
└── オフィスの利用方法.md.metadata.json

今回は次の形式で、.metadata.json ファイルを作成しました。2024年度事業計画.md.metadata.json のみ、for_managers が true かつ、 year が 2024 です。

{
	"metadataAttributes": {
		"target": "company-wide",
		"year": 2024,
		"for_managers": true
	}
}

カラムの追加

まずはカラムの追加です。Aurora データベースに対して、カラムを追加してあげます。 Data API が有効なため、 Data API 経由で実行します。

export REGION="us-west-2"
export CLUSTER_ARN="Aurora クラスター ARN"
export SECRET_ARN="シークレット ARN"
export DATABASE_NAME="postgresql"

aws rds-data --region $REGION execute-statement \
  --resource-arn $CLUSTER_ARN \
  --secret-arn $SECRET_ARN \
  --database $DATABASE_NAME \
  --sql "ALTER TABLE bedrock_integration.bedrock_kb \
  ADD COLUMN IF NOT EXISTS target VARCHAR(100), \
  ADD COLUMN IF NOT EXISTS year SMALLINT, \
  ADD COLUMN IF NOT EXISTS for_managers BOOLEAN;"

実行例

[cloudshell-user@ip-10-130-63-143 ~]$ export REGION="us-west-2"
[cloudshell-user@ip-10-130-63-143 ~]$ export CLUSTER_ARN=arn:aws:rds:us-west-2:XXXXXXXXXXXX:cluster:kb-metadata-vctrdb-cluster
[cloudshell-user@ip-10-130-63-143 ~]$ export SECRET_ARN=arn:aws:secretsmanager:us-west-2:XXXXXXXXXXXX:secret:kb-metadata-kb-vctrdb-secret-XXXXXX
[cloudshell-user@ip-10-130-63-143 ~]$ export DATABASE_NAME="postgresql"
[cloudshell-user@ip-10-130-63-143 ~]$ aws rds-data --region $REGION execute-statement \
>   --resource-arn $CLUSTER_ARN \
>   --secret-arn $SECRET_ARN \
>   --database $DATABASE_NAME \
>   --sql "ALTER TABLE bedrock_integration.bedrock_kb \
>   ADD COLUMN IF NOT EXISTS target VARCHAR(100), \
>   ADD COLUMN IF NOT EXISTS year SMALLINT, \
>   ADD COLUMN IF NOT EXISTS for_managers BOOLEAN;"
{
    "numberOfRecordsUpdated": 0,
    "generatedFields": []
}

クエリエディタでクエリを叩いてみると列が追加されていることがわかります。

データの追加

サンプルドキュメントcompany-wide フォルダだけ、 S3 にアップロードします。

アップロードが完了したら、ナレッジベースコンソールからデータソースを選択し、 Snyc をクリックします。

無事、同期が完了して Source files と Metadata files が 5 つずつ認識されていますね。

Aurora データベースからも、各列にデータが格納されていることがわかります。

マネジメントコンソールから確認

それでは、 2024 年度の経営計画をフィルタリングしていきます。まずはマネジメントコンソールから確認します。

2024 年と for_managers を true にした状態で期待した結果が返ってきました。

今度は 2024 年かつ for_managers を false にした状態です。事業計画についてはデータがないと回答が返ってきました。

AWS SDK から確認

最後に AWS SDK (boto3) から フィルタリングを行ってみます。環境変数はすでに設定済みとします。

import os
import boto3

bucket_name = os.environ.get("BUCKET_NAME")
region = os.environ.get("AWS_REGION")
kb_id = os.environ.get("KNOWLEDGE_BASE_ID")
client = boto3.client("bedrock-agent-runtime", region_name=region)

filter = {
    "andAll": [
        {"equals": {"key": "year", "value": 2024}},
        {"equals": {"key": "for_managers", "value": True}},
    ]
}

response = client.retrieve_and_generate(
    input={"text": "来年はどんな事業に注力する予定ですか?"},
    retrieveAndGenerateConfiguration={
        "knowledgeBaseConfiguration": {
            "knowledgeBaseId": kb_id,
            "modelArn": "arn:aws:bedrock:us-west-2::foundation-model/anthropic.claude-3-sonnet-20240229-v1:0",
            "retrievalConfiguration": {
                "vectorSearchConfiguration": {
                    "filter": filter,
                    "numberOfResults": 5,
                }
            },
        },
        "type": "KNOWLEDGE_BASE",
    },
)

print(response["output"]["text"])

無事、結果が返ってきました。

(app-py3.12) takakuni@app % python medatada_filter.py
株式会社サンプルは2024年度に、宇宙旅行事業とAI搭載ロボットペット事業の2つの新規事業に注力する予定です。 宇宙旅行事業では、月面でのリゾート「ルナ・パラダイス」の建設、宇宙飛行士養成プログラムの提供、宇宙食の開発と販売を行います。AI搭載ロボットペット事業では、感情認識AIを搭載した「エモ・ペット」の開発、セレブリティとのコラボレーションモデルの販売、ロボットペット用アクセサリーの開発と販売を行う予定です。

まとめ

以上、「Knowledge bases for Amazon Bedrock でメタデータフィルタリングを試してみる」でした。

今後、 Amazon Kendra のように、 CDE やフォルダレベルでメタデータの付与とかできるようなアップデートに期待ですね。(Lambda とか使えば S3 のイベント通知で自動生成できそうではありますが)

このブログがどなたかの参考になれば幸いです。 AWS 事業本部コンサルティング部のたかくに(@takakuni_) でした!