Claudeにばかり気を取られてた…エンベディングモデル側の最大トークン数


# 内容
Amazon Bedrockのナレッジベース機能を使って、RAG（Retrieval-Augmented Generation）システムを構築します。RAGシステムでは、テキスト情報をベクトルに変換するエンベディングモデル、ユーザーの入力に基づいて回答を生成するテキスト生成モデルなどの基盤モデルを利用するのが一般的です。最近のテキストモデルは多くのトークンを扱うことが出来ますが、エンベディングモデル側の最大入力トークン数に注意しておく必要があります。

# ベクトルDBに情報を格納する際の流れ

情報をベクトルデータベースに登録する際の処理イメージは以下の通りです。たとえば、「Aさんは静岡県出身」といったテキスト情報をあらかじめS3に格納しておきます。次に、ナレッジベースを同期すると、S3内の情報がエンベディングモデルによってベクトル（浮動小数点数のリスト）に変換され、ベクトルDBに保存されます。

![rag01.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2523652/c264a9e7-cc6d-4703-ba69-01e0902c5a9f.png)

# ベクトルDBに問い合わせるときの流れ

チャットシステムなどを通じて、ユーザーがナレッジベースに問い合わせを行う際の処理フローについて説明します。たとえば、ユーザーが「Aさんの出身地は？」と質問した場合、Claude 3のようなテキスト生成モデルがそのままベクトルDBにクエリを投げるわけではありません。

![rag02.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2523652/8a3590e3-33b3-4b36-a2c6-5e5e6c60558c.png)

まず、「Aさんの出身地は？」といった質問文をエンベディングモデルによりベクトル（浮動小数点のリスト）へと変換します。そのベクトルをもとに、ベクトルDB内の類似情報と照合して、関連するチャンクを取得します。その後、取得した情報をもとに、テキスト生成モデルが自然な回答を生成するといった流れになります。AWSの以下マニュアルに詳しい流れが書かれています。

https://docs.aws.amazon.com/ja_jp/bedrock/latest/userguide/kb-how-it-works.html


![rag03.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2523652/cb9b87a7-c798-45db-9e0e-fc82a5a82a24.png)


# RAGを構築していて発生したエラー
Bedrockを使ってRAGシステムを構築中、`RetrieveAndGenerateStream`メソッドを使用した際に、以下のようなエラーが発生しました。入力値または設定値に問題があるとのエラーですが、毎回発生したわけではないので、入力値に問題があるのではないかとあたりをつけました。

```
An error occurred (ValidationException) when calling the RetrieveAndGenerateStream operation: Invalid input or configuration provided. Check the input and Knowledge Base configuration and try your request again. (Service: BedrockAgentRuntime, Status Code: 400, Request ID: ****-****-****-****-****) (SDK Attempt Count: 1)
```

# テスト用プログラムの作成
原因切り分けのため、ナレッジベースに対して直接問い合わせを行い、回答を得るテストプログラムを作成しました。入力文字数、応答結果、エラー内容を返してくれます。

```python
import boto3
import json

REGION = "ap-northeast-1"
MODEL_ARN = "arn:aws:bedrock:ap-northeast-1::foundation-model/anthropic.claude-3-haiku-20240307-v1:0"
KNOWLEDGEBASE_ID = "*******"

# boto3 クライアント初期化
bedrock_agent = boto3.client("bedrock-agent-runtime", region_name=REGION)

# テストプロンプト
user_question = "これはテスト用の質問です。" * 80

context_prompt = f"""
以下はユーザーからの質問です：

<question>
{user_question}
</question>
"""

print(f" context_prompt文字数: {len(context_prompt)}")

# API 呼び出し
try:
    response = bedrock_agent.retrieve_and_generate_stream(
        input={"text": context_prompt},
        retrieveAndGenerateConfiguration={
            "type": "KNOWLEDGE_BASE",
            "knowledgeBaseConfiguration": {
                "knowledgeBaseId": KNOWLEDGEBASE_ID,
                "modelArn": MODEL_ARN,
            }
        }
    )

    # ストリーム受信
    print("応答:")
    for event in response.get("stream", []):
        if "output" in event:
            print(event["output"]["text"], end="")

except Exception as e:
    print("エラー発生:")
    print(e)
```

# 入力文字数による動作確認

下記のように、異なる文字数での確認を行いました。

✅ 問題無し

```python
# テストプロンプト
user_question = "これはテスト用の質問です。" * 80

context_prompt文字数: 1082
応答:

申し訳ありませんが、この質問に対する回答を検索結果から見つけることはできませんでした。この質問は、テスト用の質問であり、検索結果には関連する情報が含まれていないようです。
```


❌ 問題あり

```python
# テストプロンプト
user_question = "これはテスト用の質問です。" * 100

context_prompt文字数: 1342
エラー発生:
An error occurred (ValidationException) when calling the RetrieveAndGenerateStream operation: Invalid input or configuration provided. Check the input and Knowledge Base configuration and try your request again. (Service: BedrockAgentRuntime, Status Code: 400, Request ID: ****-****-****-****-****) (SDK Attempt Count: 1)
```

# 対応
使用していたエンベディングモデル「Cohere Embed Multilingual」の最大トークン数512に引っかかってエラーが発生していました。ナレッジベースで使用するエンベディングモデルをTitan Text Embeddings V2（最大トークン数：8k）にして再作成後、エラーが解消されました。