Amazon BedrockでRAGを構築し､Cohere Command R+でストリーミングテキストを生成する

# はじめに

Amazon BedrockでCohere Command RとCommand R+が利用可能となったので､[Amazon Bedrock+PineconeでRAGを構築し､IPAの"安全なウェブサイトの作り方"に詳しいSlackチャットボットを作成する](https://qiita.com/revsystem/items/a07b1cc159b2d3483113)の記事で構築したRAGを使って質問に回答するコードを書いてみました｡
入力テキストをもとにした検索結果から回答を生成する場合､期待した精度の回答を得られないことがあります｡Advanced RAG用に最適化されているCommando R+モデルを使うことで精度の向上やハルシネーションの軽減が期待できるようです｡

# 参考情報

https://aws.amazon.com/jp/about-aws/whats-new/2024/04/cohere-command-r-r-plus-amazon-bedrock/

https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere-command-r-plus.html

https://docs.aws.amazon.com/bedrock/latest/userguide/inference-invoke.html

https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime/client/invoke_model_with_response_stream.html

https://aws.amazon.com/jp/blogs/news/verifying-the-accuracy-contribution-of-advanced-rag-methods-on-rag-systems-built-with-amazon-kendra-and-amazon-bedrock/

https://community.aws/content/2f38mlpgBYSMBBGeUJNMjyUmRyw/build-retrieval-augmented-generation-rag-based-generative-ai-genai-application-with-amazon-bedrock

https://docs.cohere.com/docs/retrieval-augmented-generation-rag

https://cohere.com/blog/command-r-plus-microsoft-azure

https://dev.classmethod.jp/articles/amazon-bedrock-cohere-command-r-rag


# デモ動画

![デモ動画](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/61809/5166b7b0-a1e7-0f79-7971-318e1c8253e6.gif)

ここでは､｢SQLインジェクションはどのようにして発生するのか｡その原因と対策を解説してください｡Web開発者がSQLインジェクションを防ぐために取るべき対策は何ですか｡｣という質問に対して｢SQLインジェクション 原因 対策｣という検索クエリが生成されました｡この検索クエリを元に検索結果を取得し､応答テキストが生成されました｡
言語モデルは､cohere.command-r-plus-v1:0を使用しました｡

> 2024-06-02 03:03:58,966 [INFO] Found credentials in shared credentials file: ~/.aws/credentials
> 2024-06-02 03:04:00,369 [INFO] input_text:
>     SQLインジェクションはどのようにして発生するのか｡その原因と対策を解説してください｡Web開発者がSQLインジェクションを防ぐために取るべき対策は何ですか｡
> 
> 2024-06-02 03:04:00,370 [INFO] Search queries: [{'text': 'SQLインジェクション 原因 対策', 'generation_id':
> 'f060c8cd-d9b5-4535-8ef6-89c7e6171d90'}]
> 
> [Stream output]====================================
> SQLインジェクションは、データベースへの命令文を構成する入力値を送信することで発生 する脆弱性です。悪意のある
> リクエストにより、データベースの不正利用や情報漏えい、改ざん、消去などの被害を受ける可能性があります。
> 
> SQLインジェクションを防ぐための対策として、以下のような方法が挙げられます。
> 
> - プレースホルダを用いてSQL文を組み立てる：SQL文の雛形の中に変数の場所を示す記号（プレースホルダ）を置き、後に実際の値を機械的な処理で割り当てる方法です。この方法では、機械的な処理でSQL文が組み立てられるため、SQLインジェクションの脆弱性を解消できます。
> - 静的プレースホルダと動的プレースホルダ：プレースホルダに実際の値を割り当てる処理をバインドと呼び、静的プレースホルダと動的プレースホルダの2つの方式があります。静 的プレースホルダは、SQLのISO/JIS規格では、準備された文（Prepared Statement）と呼ばれ、アプリケーション側のデータベース接続ライブラリ内で値をエスケープ処理してプレースホルダにはめ込む方式です。動的プレースホルダは、原理的にSQLインジェクション脆弱 性の可能性がなくなるという点で優れています。
> - SQL文のリテラルを正しく構成する：SQL文の組み立てを文字列連結により行う場合は、エスケープ処理等を行うデータベースエンジンのAPIを用いて、SQL文のリテラルを正しく構成する必要があります。
> - 権限の管理：ウェブアプリケーションがデータベースに接続する際に使用するアカウントの権限を必要最低限のものに制限します。
> - エラーメッセージの管理：データベースに関連するエラーメッセージを、利用者のブラウザ上に表示させないようにします。エラーメッセージには、データベースの種類やエラーの原因、実行エラーを起こしたSQL文などの情報が含まれており、SQLインジェクション攻撃につながる有用な情報となる可能性があるためです。


cohere.command-r-v1:0を使った場合の出力例も以下に示します｡cohere.command-r-v1:0の場合は､検索クエリーが英語で生成されることが多いようです｡一方､`日本語で回答を提供すること｡`という指示の有無にかかわらず日本語での質問に対する応答テキストは日本語で生成されることが多いようです｡


<details><summary>cohere.command-r-v1:0を使った場合の出力例</summary>

> 2024-06-02 02:00:38,454 [INFO] Found credentials in shared credentials file: ~/.aws/credentials
> 2024-06-02 02:00:39,651 [INFO] input_text:
>     SQLインジェクションはどのようにして発生するのか｡その原因と対策を解説してください｡Web開発者がSQLインジェ
> クションを防ぐために取るべき対策は何ですか｡
> 
> 2024-06-02 02:00:39,651 [INFO] Search queries: [{'text': 'how does SQL injection happen?', 'generation_id':
> '58e46ae8-bb00-4992-8d44-bf715d50425f'}, {'text': 'SQL injection prevention', 'generation_id':
> '58e46ae8-bb00-4992-8d44-bf715d50425f'}]
> 
> [Stream output]====================================
> SQLインジェクション（SQL Injection）は、ウェブアプリケーションに渡されるパラメータに直接SQL文が指定されることで発生するセキュリティ問題です。パラメータ値が改変され 、データベースが不正利用される可能性があります。
> 
> ## SQLインジェクションの原因
> SQL文のリテラル部分をパラメータ化when普通、アプリケーションからSQLを利用する場合に、パラメータ化された部分を実際の値に展開するときに、リテラルとして文法的に正しく生成されないことがあります。そうすると、パラメータに与えられた値がリテラルの外にはみ出し、元のSQL文 Meaningを変更できるようになってしまいます。これがSQLインジェクションの原因です。
> 
> 例えば、
> ```SQL
> SELECT * FROM atable WHERE id=$id
> ```
> というSQL文がある場合、$idに以下の値が与えられると、文 meaningが変わってしまいます。
> ```
> 0; DELETE FROM atable
> ```
> このようにしてデータベースの内容がすべて削除されることにもなりかねません。
> また、エラーメッセージの内容にデータベースの種類やエラーの原因が含まれてしまうと、それらがSQLインジェクション攻撃に悪用される可能性があります。そのため、エラーメッ セージは利用者のブラウザに表示しないようにしましょう。
> ## 対策
> 根本的な解決策として、バインド機構を利用するのが有効です。それ桁 projekができない 場合は、以下の対策が挙げられます。
> - ウェブアプリケーションに渡されるパラメータに、SQL文を直接指定しない
> - エラーメッセージをそのままブラウザに表示しない
> - データベースアカウントに適切な権限を与える
> 
> さらに、実行環境次第ではありますが、エスケープ処理を適切に行わないと生じるAPIの脆 弱性もあります。そのような場合は、修正パッチを適用させるか、別の手法を検討しましょう。
</details>


## 開発環境構築

### 作業環境のOSバージョン

Windows 11上のWSLでUbuntu 23.04を動かしています。

```console
$ cat /etc/os-release | grep PRETTY_NAME
PRETTY_NAME="Ubuntu 23.10"
```

### Python環境

```console
$ python3 --version
Python 3.12.0
```

```console
$ python3 -m venv .venv
source .venv/bin/activate
```

```
$ pip3 install --upgrade pip
$ pip3 --version
pip 24.0 from /home/xxx/.pyenv/versions/3.12.0/lib/python3.12/site-packages/pip (python 3.12)
```

```console
$ pip install boto3
```

## サンプルコード

コードは､[[アップデート] Amazon Bedrockで新モデル「Cohere Command R/R+」が利用可能になったので、RAGで使ってみた](https://dev.classmethod.jp/articles/amazon-bedrock-cohere-command-r-rag/) に紹介されているものを利用しました｡ 応答テキストを生成する箇所で[invoke_model_with_response_stream](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime/client/invoke_model_with_response_stream.html)を使ってストリーム出力しています｡


```python:sample.py
import json
import logging

import boto3

logging.basicConfig(
    format="%(asctime)s [%(levelname)s] %(message)s",
    level=logging.INFO
)

logger = logging.getLogger(__name__)


def generate_message(bedrock, agent, model_id, knowledgebase_id, input_text):

    # Step 1: ユーザー入力テキストを基に検索クエリーを生成する
    request_body = json.dumps({
        "message": input_text,
        "search_queries_only": True, # 検索クエリーのみを生成する
    })

    response = bedrock.invoke_model(
        body=request_body,
        contentType="application/json",
        accept="application/json",
        modelId=model_id,
    )

    response_body = json.loads(response.get("body").read())
    search_queries = response_body["search_queries"]
    logging.info("input_text: %s", input_text)
    logging.info("Search queries: %s", search_queries, )

    # Step 2: Retrieverに対してクエリーを実行して検索結果を得る

    documents = []

    for query in search_queries:
        query_text = query["text"]
        response = agent.retrieve(
            knowledgeBaseId=knowledgebase_id,
            retrievalQuery={
                "text": query_text
            },
            retrievalConfiguration={
                "vectorSearchConfiguration": {
                    "numberOfResults": 3
                }
            },
        )

        retrieval_results = response['retrievalResults']
        logging.debug("response: %s", response)
        logging.debug("retrieval_results: %s", retrieval_results)

        for result in retrieval_results:
            snippet = result['content']['text']

            documents.append(
                {
                    "query": query_text,
                    "snippet": snippet,
                }
            )
        logging.debug("documents: %s", documents)

    # Step 3: ユーザー入力テキストと検索結果を基に応答テキストを生成する
    system_instruction = """
    あなたは親切で知識豊富なチャットアシスタントです。
    あなたにはユーザーが知りたい情報に関連する複数のドキュメントの抜粋が提供されます｡
    これらの情報をもとに、ユーザーの質問に対する回答を提供してください。
    質問に答えるための情報がない場合は、「情報が不十分で回答できません」と答えてください。
    また、質問への回答は以下の点に留意してください｡


    - 日本語で回答を提供すること｡
    - 質問者は､Webサイトの開発や運用に携わるエンジニアです｡
    - 質問者に対して詳細な説明をエンジニア向けの内容､用語で提供すること｡
    - 質問に対する回答に複数の可能性が含まれる場合､それぞれの可能性について検討し詳細な回答を提供すること｡
    - 回答には､解決策や対応策の例を含ること｡例は質問者の理解を助けます｡
    """

    request_body = json.dumps({
        "preamble": system_instruction,
        "message": input_text,
        "documents": documents,
        "max_tokens": 5000,
        "temperature": 1,
    })

    response = bedrock.invoke_model_with_response_stream(
        body=request_body,
        contentType="application/json",
        accept="application/json",
        modelId=model_id,
    )


    print("\n[Stream output]====================================")
    # ストリーム出力
    for event in response["body"]:
        chunk = json.loads(event["chunk"]["bytes"])
        if chunk['event_type'] == 'text-generation':
            yield chunk['text']

    logging.debug(json.dumps(chunk["response"], indent=2, ensure_ascii=False))


def main():

    bedrock = boto3.client(
        service_name="bedrock-runtime",
        region_name="us-east-1"
    )
    agent = boto3.client(
        service_name="bedrock-agent-runtime",
        region_name="us-east-1"
    )

    model_id = 'cohere.command-r-plus-v1:0'
    knowledgebase_id = "xxxxxxxxxx"

    input_text = """
    SQLインジェクションはどのようにして発生するのか｡その原因と対策を解説してください｡Web開発者がSQLインジェクションを防ぐために取るべき対策は何ですか｡
     """

    for sentence in generate_message(bedrock, agent, model_id, knowledgebase_id, input_text):
        print(sentence, end="", flush=True)


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

## 実行

```
$ python3 ./sample.py
```


## コードについて

`generate_message`関数は､ユーザー入力テキストを基に検索クエリーを生成し､Retrieverに対してクエリーを実行して検索結果を得ます｡
Step1において以下のようなプロンプトを用意し検索クエリーの生成を調整できないか試しましたが､検索クエリーの生成数や内容は意図したとおりにはなりませんでした｡


```python:うまくいかなかった例
    # Step 1: ユーザー入力テキストを基に検索クエリーを生成する
    request_body = json.dumps({
        "temperature": 0.0,
        "p": 0.99,
        "k": 250,
        "max_tokens": 1000,
        "preamble": "質問文に基づいて、3個の検索クエリを生成してください。\
                     各クエリは30トークン以内とし、日本語と英語を適切に混ぜてください｡",
        # "chat_history" is not used for "search_queries_only" with empty []
        "message": input_text,
        "search_queries_only": True,
    })

    response = bedrock.invoke_model(
        body=request_body,
        contentType="application/json",
        accept="application/json",
        modelId="cohere.command-r-v1:0",
    )
```

Step3では以下のようなプロンプトを用意し､`input_text`と検索結果をもとに応答テキストを生成するよう指示を出しています｡これらの指示によって応答テキストの文体や内容を調整することができました｡

```python:回答テキストに関するプロンプト
    # Step 3: ユーザー入力テキストと検索結果を基に応答テキストを生成する
    system_instruction = """
    あなたは親切で知識豊富なチャットアシスタントです。
    あなたにはユーザーが知りたい情報に関連する複数のドキュメントの抜粋が提供されます｡
    これらの情報をもとに、ユーザーの質問に対する回答を提供してください。
    質問に答えるための情報がない場合は、「情報が不十分で回答できません」と答えてください。
    また、質問への回答は以下の点に留意してください｡


    - 日本語で回答を提供すること｡
    - 質問者は､Webサイトの開発や運用に携わるエンジニアです｡
    - 質問者に対して詳細な説明をエンジニア向けの内容､用語で提供すること｡
    - 質問に対する回答に複数の可能性が含まれる場合､それぞれの可能性について検討し詳細な回答を提供すること｡
    - 回答には､解決策や対応策の例を含ること｡例は質問者の理解を助けます｡
    """

    request_body = json.dumps({
        "preamble": system_instruction,
        "message": input_text,
        "documents": documents,
        "max_tokens": 5000,
        "temperature": 1,
    })

    response = bedrock.invoke_model_with_response_stream(
        body=request_body,
        contentType="application/json",
        accept="application/json",
        modelId=model_id,
    )

```

# まとめ

Cohere Command R+モデルを使って質問に回答するコードを書いてみました｡モデルがAdvanced RAG用に最適化されているだけに､良い検索結果を得られるような検索クエリーの生成や検索結果をもとにした回答の生成が容易に実現できました｡
個人的な感覚として､従来のRAGでは入力テキストやプロンプトを工夫しても想定した回答を得られないことがありました｡Cohere Command R+モデルを用いたAdvanced RAGを構築することで精度の向上やハルシネーションの軽減が期待できそうです｡

前述のデモ画面の項目でも書きましたが､Command RとCommand R+では検索クエリの生成に言語の違いが現れました｡記事内では省略しましたが､コード内のStep3のプロンプトに､`回答は大阪弁で提供すること｡`という指示を加えた場合､cohere.command-r-v1:0ではわずかに大阪弁が混じりましたが全体的に方言のない文章となりました｡しかし､cohere.command-r-plus-v1:0の場合､応答テキスト全体が大阪弁となりました｡Command RとCommando R+では知識ベースが異なるのかもしれません｡

このあたりの違いと費用を踏まえて､検索クエリ生成処理と応答テキスト生成処理のそれぞれでモデルを使い分けても良さそうです｡