Knowledge Base for Amazon Bedrock でメタデータフィルタリングが可能に

## はじめに
Knowledge Base for Amazon Bedrock にメタデータフィルタリング機能が追加されました。これにより、クエリに関連性の高いドキュメントのみを検索対象とすることができるため、検索精度の向上が期待できます。

https://aws.amazon.com/jp/about-aws/whats-new/2024/03/knowledge-bases-amazon-bedrock-metadata-filtering/

たとえばニュース記事をデータソースとする場合、事前に西暦 (year) やジャンルといった情報をメタデータとして設定しておくことで、これらを使用したフィルタリングが可能になります。



## やってみる (Amazon Aurora 編)
ベクトルデータベースに Aurora を使用したナレッジベースの環境が手元にあったので、こちらで試してみます。

Aurora を使用したナレッジベースの構築方法は以下の記事で紹介しています。

https://qiita.com/hayao_k/items/45e59c1c2a183c27b20d


### メタデータファイルの準備
データソースとなるファイル単位で以下のようなスキーマの JSON ファイルを準備します。

```scheme.json
{
   "metadataAttributes": {
      "${attribute1}": "${value1}",
      "${attribute2}": "${value2}",
      ...
   }
}
```

属性値は以下のデータ型がサポートされています。

* String
* Number
* Boolean

メタデータのファイル名はデータソースとなるファイル名に `.metadata.json` を追加した形式とする必要があります。例えば `A.txt` に対してメタデータファイルを用意する場合、`A.txt.metadata.json` を作成し、データソースが配置されている同じ S3 バケットに配置します。

以下にご注意ください。

* メタデータのファイルサイズは 10 KB 以下にする必要があります
* ソースドキュメントと同じフォルダーに配置する必要があります

今回は Bedrock のユーザーガイドと API Reference (どちらも英語版) を S3 に配置し、それぞれ以下のようなメタデータファイルを準備しました。

```json:bedrock-ug.pdf.metadata.json
{
    "metadataAttributes": {
        "category": "User Guide",
        "service": "Bedrock",
        "year": 2024
    }
}
```

```json:bedrock-api.pdf.metadata.json
{
    "metadataAttributes": {
        "category": "API Reference",
        "service": "Bedrock",
        "year": 2024
    }
}
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/125105/8a8f03c4-f27a-fd64-f6ef-67de64ebcff6.png)


### ベクトルデータベースへの列追加
使用しているベクトルデータベースをメタデータに対応した形へ変更する必要があります。

Amazon Aurora データベースクラスター内の既存のベクトルインデックスにメタデータを追加する場合は、メタデータの属性ごとにテーブルへ列を追加します。

今回のケースでは以下のような SQL を実行して既存の Aurora クラスター上のテーブルに列を追加します。

```sql
ALTER TABLE bedrock_integration.bedrock_kb
  ADD COLUMN category VARCHAR(128),
  ADD COLUMN service VARCHAR(128),
  ADD COLUMN year SMALLINT;
```

:::note info
本記事では詳細に触れませんが、ベクトルデータベースに Amazon OpenSearch Serverless を使用している場合は以下のいずれか対応が必要です。詳細はドキュメントを参照してください。

* コンソールで新規のナレッジベースの作成時に新しいベクトルストアをクイック作成する
* 既存のベクトルストアを使用したい場合は Engine を `faiss` とした新しいベクトルインデックスを作成し、ナレッジベース作成時に新しいインデックスを指定

※ つまりいずれの場合もナレッジベース自体の再作成は必要になります
:::

### 再同期
コンソールまたは [StartIngestionJob API](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent_StartIngestionJob.html) を使用してデータソースを再同期します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/125105/c3840d29-7337-9fbc-ac7d-fc852e52470e.png)


:::note warn
事前にベクトルデータベースへの列追加が完了していない場合は同期エラーが発生します。
:::

### 動作確認 (コンソール)
以下はフィルターを設定せず、「ナレッジベースから回答を得るにはどうすればいいですか」という質問をした結果です。回答結果からユーザーガイドをソースとして回答していることがわかります。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/125105/77cdfe75-88d1-c064-11a2-6d353ef55f9c.png)

次に API Reference をソースとして回答させるために `category = API Reference` を Filters に設定して実行します。

:::note warn
コンソールでフィルターを設定する際は値にシングルクォーテーションやダブルクォーテーションを付与しないでください。付与すると適切にフィルタリングされません。
:::

API Referenceの情報をソースとした回答を得ることができました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/125105/0b0a9357-8862-f378-98ba-17fc36593fa3.png)

回答結果の詳細からソースチャンクを表示すると、チャンクに関連付けられたメタデータを確認できます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/125105/6b66d345-624f-1fda-8505-596fc5d0a971.png)

フィルターに設定可能な演算子はドキュメントに記載がありますので、以下を参照ください。

https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-config.html

AND 条件で `year = 2023` 追加すると、該当するメタデータを持つソースデータがないため、ナレッジベースから回答は得られません。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/125105/e9f3fad0-27c2-6aa8-6c5d-a3400837d430.png)

Aurora クラスター上のデータを SQL で参照すると、追加した列にメタデータの値がそのまま挿入されていることがわかります。

```sql
select * from bedrock_integration.bedrock_kb where year = 2024 limit 10;
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/125105/b3f8ac2a-cedd-f5ef-f8e8-d43a6bd6a4d2.png)

### 動作確認 (boto3)
コンソールと同様の確認を boto3 でも試してみます。boto3 のバージョンは `v1.34.73` で動作確認しています。

以下はコード例です。

```python:kb.py
import boto3

def main():
    client = boto3.client("bedrock-agent-runtime")
    
    response = client.retrieve_and_generate(
        input={"text": "ナレッジベースから回答を得るにはどうしたらいいですか？"},
        retrieveAndGenerateConfiguration={
            "type": 'KNOWLEDGE_BASE',
            "knowledgeBaseConfiguration": {
                "knowledgeBaseId": "KB12345678", 
                "modelArn": "arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-3-sonnet-20240229-v1:0",
                "retrievalConfiguration": {
                    "vectorSearchConfiguration": {
                        "numberOfResults": 5,
                        "filter": {
                            "equals": {
                                "key": "category",
                                "value": "API Reference"
                            }
                        }
                    }
                }
            }
        }
    )
    
    print(f'Answer: {response["output"]["text"]}')
    print(f'\nMetadata: {response["citations"][0]["retrievedReferences"][0]["metadata"]}')

if __name__ == "__main__":
    main()
```

実行すると以下の結果を得られました。

```shell-session
$ python kb.y
Answer: ナレッジベースから回答を得るには、以下の手順が必要です:

1. データソースを設定し、サポートされているベクトル格納庫を構成する必要があります。詳細は「データの取り込み準備」を参照してください。

2. Amazon Bedrock コンソールを使用して、Amazon OpenSearch Service 内にベクトル格納庫を作成・管理させることもできます。詳細は「ナレッジベースの作成」を参照してください。 3. ナレッジベースを作成する際には、以下の情報を指定する必要があります:
- ナレッジベース名と説明
- ナレッジベース作成のための IAM ロール ARN
- 使用する埋め込みモデル
- ベクトル格納庫の構成 (OpenSearch Service、Aurora、Pinecone、Redis Enterprise Cloud のいずれか)

Metadata: {'category': 'API Reference', 'year': 2024.0, 'service': 'Bedrock'}
```

複合条件としたい場合の Filter の設定例は以下のようになります。

```json
                        "filter": {
                            "andAll": [
                                {
                                    "equals": {
                                        "key": "category",
                                        "value": "API Reference"
                                    }
                                },
                                {
                                    "equals": {
                                        "key": "year",
                                        "value": 2023
                                    }
                                }
                            ]
                        }
```
## 参考
https://docs.aws.amazon.com/bedrock/latest/userguide/knowledge-base-ds.html#kb-ds-metadata

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