Knowledge bases for Amazon Bedrockのナレッジベースを複数構築する効率的な２つの方法

Amazon Bedrockで簡単にRAGを構築する仕組みとして、 **「Knowledge bases for Amazon Bedrock」** という機能が提供されています。

https://aws.amazon.com/jp/bedrock/knowledge-bases/

マネジメントコンソールのウィザードに従って簡単に構築でき、ベクトルデータベースとしてOpenSearch Serverlessのコレクションを自動で作成する **「クイック作成」** が選択できます。

とてもお手軽なのですが、複数のナレッジベースを構築すると、毎回OpenSearch Serverlessのコレクションが作成されます。

図：ナレッジベースごとにコレクションが作成されるイメージ

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/00dddb1c-f336-c7f1-99ab-806386d63d6f.png)

OpenSearch Serverlessの料金は、OCUという単位で計算され、1コレクションあたり最低1OCUが必要です（コレクションが２つの場合は2OCUが最低となります）。

コレクションをひとつに集約できれば最低料金が安く抑えられるので、その方法をふたつ紹介します。

:::note warn
投稿したあとに気づいたのですが、どうやら、複数のコレクションを作ってもOCUは共有するようです。そのため、コストがきになる場合でも、どんどんコレクションを作っても問題ないかもしれません。。

参考：https://docs.aws.amazon.com/opensearch-service/latest/developerguide/serverless-overview.html

> アカウントの最初のコレクションに対して最低 1 OCU [0.5 OCU x 2] が課金されます。後続のすべてのコレクションでこれらの OCU を共有できます。
>
> 固有の AWS KMS キーを持つコレクションは、他のコレクションと OCU を共有できません。

コレクションを2つ用意して半日放置してみましたが、インデックス作成と検索はどちらも0.5OCUでした。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/f63b86e8-345e-5378-668f-c8c20ece8c00.png)

:::


* 方法１：1コレクションに複数インデックスを作成する方法

    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/fa129940-250d-7045-5efc-a3ad2c8f4d37.png)

* 方法２：データソースを複数作成する方法

    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/199987cb-738f-837b-a780-12316008a956.png)

それぞれ紹介します。

:::note
OpenSearch Serverlessの用語は雑に以下のような感じです
コレクション： RDBで言うところのデータベース
インデックス： RDBで言うところのテーブル
:::

## 方法１：1コレクションに複数インデックスを作成する方法

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/fa129940-250d-7045-5efc-a3ad2c8f4d37.png)

Knowledge bases for Amazon Bedrockのクイック作成で、すでにひとつナレッジベースを作成した状態から開始します。

### インデックスを作成

OpenSearch Serverlessには1つのコレクションと1つのインデックスがある状態です。

* コレクション
    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/e16991a3-8dc2-88e4-95a8-0f13fddbcd07.png)

* インデックス
    ![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/75d9e072-cf76-f93c-6b0e-5a18515dc0cb.png)

作成済みのインデックスを確認すると、以下の内容で作成されていました。

![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/3c00fd82-16f0-7b34-45cd-e2afd2ccda57.png)

ベクトルフィールドやメタデータを真似して、 **「bedrock-knowledge-base-default-index-2」** を作成します。

![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/d3b26640-9505-8fae-7069-5a793e2345d1.png)

インデックスの追加作業は以上です。

### ナレッジベースを追加

ベクトルデータベースが作成できたので、Knowledge bases for Amazon Bedrockの画面でナレッジベースを追加します。

ナレッジベースの作成ウィザードで、ベクトルデータベースを指定する部分で、「作成したベクトルストアを選択」を選び、「Amazon OpenSearch Serverless用ベクトルエンジン」を選択します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/fbcb0454-aa24-a29d-0ff0-774d47ef894e.png)

先ほど作成したインデックスの情報を入力します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/d787156e-2632-b806-3e2c-24ff0879fb2f.png)

以上で設定完了です。

Pythonで呼び出す際の「knowledge_base_id」を切り替えることで、どのナレッジベースを使用するかを指定します。

```python:app.py
import json

import boto3

client = boto3.client("bedrock-agent-runtime")

question = "Bedrockとはなんですか？"
knowledge_base_id = "VB66V1AE0I"    ## ここを使用するナレッジベースで切り替え
model_arn = "anthropic.claude-3-haiku-20240307-v1:0"

response = client.retrieve_and_generate(
    input={"text": question},
    retrieveAndGenerateConfiguration={
        "type": "KNOWLEDGE_BASE",
        "knowledgeBaseConfiguration": {
            "knowledgeBaseId": knowledge_base_id,
            "modelArn": model_arn,
        },
    },
)

print(json.dumps(response, indent=2, ensure_ascii=False))
```

## 方法２：データソースを複数作成する方法

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/199987cb-738f-837b-a780-12316008a956.png)

インデックスを分けるのではなく、データソースを分けることも可能です。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/7cacb382-3797-bd9e-495a-47ab17992858.png)


インデックスを分ける方法と比較し、データソースを分ける方法は以下の特徴があります。

* ナレッジベースに含められるデータソースは最大5つまで
* Retrieve APIまたはRetrieveAndGenerate APIを呼び出す際に「どのデータソースを使用するか」を明示的に指定する必要がある
* 検索時に２つのデータソースにまたがって検索することができる

同期を行うと、キーが「x-amz-bedrock-kb-data-source-id」のメタデータにデータソースIDがセットされます。この値を検索時にフィルター条件として指定します。階層が深いですが以下のようにメタデータでフィルターを行います。

> retrieveAndGenerateConfiguration -> knowledgeBaseConfiguration -> retrievalConfiguration -> vectorSearchConfiguration -> filter

```python:フィルター条件でデータソースIDを指定
response = client.retrieve_and_generate(
    input={"text": question},
    retrieveAndGenerateConfiguration={
        "type": "KNOWLEDGE_BASE",
        "knowledgeBaseConfiguration": {
            "knowledgeBaseId": knowledge_base_id,
            "modelArn": model_arn,
            "retrievalConfiguration": {
                "vectorSearchConfiguration": {
                    "filter": {
                        "equals": {
                            "key": "x-amz-bedrock-kb-data-source-id",
                            "value": "5CSTAI5JWZ",
                        }
                    }
                }
            },
        },
    },
)
```

複数のフィルター条件を組み合わせることもできますので、２つのデータソースからまとめて検索することも可能です。（orAllを使用します）

```python: フィルター条件としてデータソースIDを複数指定
response = client.retrieve_and_generate(
    input={"text": question},
    retrieveAndGenerateConfiguration={
        "type": "KNOWLEDGE_BASE",
        "knowledgeBaseConfiguration": {
            "knowledgeBaseId": knowledge_base_id,
            "modelArn": model_arn,
            "retrievalConfiguration": {
                "vectorSearchConfiguration": {
                    "filter": {
                        "orAll": [
                            {
                                "equals": {
                                    "key": "x-amz-bedrock-kb-data-source-id",
                                    "value": "5CSTAI5JWZ",
                                }
                            },
                            {
                                "equals": {
                                    "key": "x-amz-bedrock-kb-data-source-id",
                                    "value": "BED4RXALWW",
                                }
                            },
                        ]
                    }
                }
            },
        },
    },
)
```

:::note
データソースが３つある場合に、そのうち２つとかの指定も可能です
:::

以下の２つのデータソースを用意して検証しました。

| データソースID | 格納したドキュメント |
| --- | --- |
| 5CSTAI5JWZ | AWS App Studioに関するドキュメント |
| BED4RXALWW | BedrockとBedrock Studioの関するドキュメント |

### AWS App Studioに関するドキュメントだけを検索対象にした場合

```json:vectorSearchConfigurationの内容
"vectorSearchConfiguration": {
    "filter": {
        "equals": {
            "key": "x-amz-bedrock-kb-data-source-id",
            "value": "5CSTAI5JWZ",
        }
    }
}
```

質問：
BedrockとBedrock StudioとApp Studioについて教えて下さい。日本語でお願いします。

回答：
Bedrock は、Minecraft のゲームエンジンです。Bedrock Edition は、Windows 10、Xbox、PlayStation、Nintendo Switch、iOS、Android などの様々なプラットフォームで動作するバージョンです。Bedrock は、クロスプラットフォーム対応や高度な機能を備えています。 Bedrock Studio は、Bedrock Edition の開発を行う Mojang Studios の一部門です。Bedrock Studio は、Bedrock Edition の開発、更新、サポートを担当しています。 AWS App Studio は、AWS が提供するサービスで、自然言語を使ってエンタープライズレベルのアプリケーションを作成できるツールです。IT プロジェクトマネージャーやデータエンジニア、エンタープライズアーキテクトなどの技術的スキルのない人でも、内部業務アプリケーションを迅速に構築できるようサポートします。

:::note
App Studioについては正確に回答できていますが、Bedrockについてはハルシネーションが発生しています。
:::

### AWS App Studio、BedrockとBedrock Studioに関するドキュメントを検索対象にした場合

```json:vectorSearchConfigurationの内容
"vectorSearchConfiguration": {
    "filter": {
        "orAll": [
            {
                "equals": {
                    "key": "x-amz-bedrock-kb-data-source-id",
                    "value": "5CSTAI5JWZ",
                }
            },
            {
                "equals": {
                    "key": "x-amz-bedrock-kb-data-source-id",
                    "value": "BED4RXALWW",
                }
            },
        ]
    }
}
```

質問：
BedrockとBedrock StudioとApp Studioについて教えて下さい。日本語でお願いします。

回答：
Amazon Bedrock は、Amazon が提供する完全管理型のサービスで、さまざまなAIスタートアップ企業が開発したファウンデーションモデルを統一されたAPIを通して利用できるようにしたものです。ユースケースに合わせて適切なファウンデーションモデルを選択でき、セキュリティ、プライバシー、責任あるAIなどの機能も提供しています。ユーザーは自社のデータを使ってモデルをカスタマイズしたり、AWS上のツールを使ってアプリケーションに統合・デプロイすることができます。 AWS App Studio は、自然言語を使ってエンタープライズレベルのアプリケーションを作成できるサービスです。IT管理者やデータエンジニア、エンタープライズアーキテクトなどの技術者以外の人でも、セキュリティとフル管理機能を備えたアプリケーションを迅速に構築できるようになります。在庫管理、請求処理、複雑な承認フローなどの内部業務プロセスの近代化に活用できます。 Amazon Bedrock Studio は、開発環境を設定することなく、Amazon Bedrock のモデルや機能を使ったプロトタイプアプリを簡単に作成できるWebアプリケーションです。ナレッジベースの活用や、ガードレールの実装、関数呼び出しなどの機能を使ったアプリケーションを構築できます。Bedrock Studioを使うには、所属する組織からログイン情報を取得する必要があります。

:::note
BedrockとBedrock StudioとApp Studioの全てに正確に回答できました
:::

## まとめ

以上、一つのコレクションで２つのナレッジベースを構築する方法を紹介しました。（２つ目の方法は擬似的にナレッジベースが２つですが）

| 観点 | 方法１<br/>1コレクションに複数インデックス | 方法２<br/>データソースを複数 |
| --- | --- | --- |
| OpenSearch コレクションの数 | 1 | 1|
| OpenSearch インデックスの数 | 2 | 1 |
| データソースの数 | 2 | 2 |
| 指定方法 | ナレッジベースID | データソースID |
| 作成上限数 | 100個（アカウントあたりのナレッジベース数の上限） | 5個（ナレッジベースあたりのデータソースの上限） |
| ポイント | 構築難易度は高くない。上限数を気にせず作成可能。複数のナレッジにまたがって検索することはできない | Knowledge bases for Amazon Bedrockの画面のみで構築可能。上限の制約が厳しい。検索時の条件でナレッジをまたがって検索することも可能 |