Converse APIとTool useを使って動的に検索用クエリを生成する

はじめに

前回の記事｢Amazon Bedrock Converse APIとTool useで実現するAdvanced RAG: クエリ拡張によるBacklogヘルプセンターの検索｣では、以下のように生成したい検索用クエリの数だけqueryを並べていました。この場合､生成するクエリの数を変えるにはquery_4, query_5...など足したりあるいは削ったりする必要があるため､非効率でした｡

config/tools_definition.json

                "inputSchema": {
                    "json": {
                        "type": "object",
                        "properties": {
                            "query_1": {"type": "string", "description": "検索用クエリ。日本語と英語を混ぜた多様な単語を空白で区切って記述される。"},
                            "query_2": {"type": "string", "description": "検索用クエリ。日本語と英語を混ぜた多様な単語を空白で区切って記述される。"},
                            "query_3": {"type": "string", "description": "検索用クエリ。日本語と英語を混ぜた多様な単語を空白で区切って記述される。"}
                        },
                        "required": ["query_1", "query_2", "query_3"]
                    }
                }

@moritalous さんの記事に複数の内容を生成するにはToolの定義をArray型にする手法が書かれていたので、前回の記事の改善に使えないか試してみました｡

今回のアプリケーションの出力例

前回の実装では、出力の定義に制限がありました。具体的には、以下のような特徴がありました：

1. config/tools_definition.jsonファイル内で、出力の定義を固定で記述していました。
2. query_1, query_2, query_3といった形で、生成するクエリの数を事前に決めていました。

今回の実装では、この制限を取り除き、より柔軟な方法を採用しています。

今回はToolの定義にArray型を使い､queryをひとつだけ定義しました｡そして、インストラクションに与えられた質問文に基づいて、3個以上の検索用クエリを生成してください。と記述しました｡その結果､以下のようにjsonにquery_1~query_5が動的に生成されました｡また､インストラクションに検索用クエリを5個生成してください。と記述すればそのとおりに生成されました。つまり、tools_definition.json内の定義を変更せずに自然言語による指示を使って検索用クエリを動的に生成することができました。この方法により、コード変更なしに柔軟にクエリ生成の制御が可能になります。試した範囲では、インストラクションで指定したクエリ数に対する挙動は安定しているようです。

前回の記事で挙げた出力例のときと同じ質問に対して､以下のような回答が出力されました｡

2024-07-21 17:10:47,533 [INFO] query_1: ガントチャート Gantt chart プロジェクト管理 スケジュール 可視化
2024-07-21 17:10:48,013 [INFO] query_2: バーンダウンチャート burndown chart 進捗管理 比較 違い
2024-07-21 17:10:48,654 [INFO] query_3: プロジェクトマネジメント ガントチャート 活用法 メリット デメリット
2024-07-21 17:10:49,281 [INFO] query_4: ガントチャート 作成方法 ツール ソフトウェア 例
2024-07-21 17:10:49,901 [INFO] query_5: プロジェクト進捗管理 ガントチャート バーンダウンチャート 使い分け
2024-07-21 17:10:50,504 [INFO] query_0: ガントチャートとはどのようなものですか? バーンダウンチャートとは異なるものですか? プロジェクト管理での利用方法を教えてください。
2024-07-21 17:11:17,686 [INFO] Usage: {
  "inputTokens": 36963,
  "outputTokens": 1048,
  "totalTokens": 38011
}
********************************************************************************
Answer:
ガントチャートは作業の計画およびスケジュールを横型棒グラフで示した図です。各課題やタスクの開始時期・終了時期が把握しやすくなります。[1]

バーンダウンチャートとは異なるものです。バーンダウンチャートは計画と実績を表すグラフで、プロジェクトの進捗を可視化するために使用されます。[2][3]

プロジェクト管理でのガントチャートの利用方法は以下の通りです：

1. スケジュールの確認と調整：課題の開始日と期限日が帯状のグラフとして表示され、直感的に把握できます。[4]

2. 課題の操作：ドラッグ＆ドロップで期間を変更したり、課題詳細を開いて編集したりできます。担当者、状態、開始日、期限日を変更できます。[5]

3. 課題の絞り込みとグルーピング：表示する課題の絞り込みとグルーピングができます。[6]

4. エクスポート：ガントチャートをExcel形式でダウンロードし、Backlogに参加していないメンバーにもスケジュールを共有できます。[7]

これらの機能を活用することで、プロジェクトのスケジュールが可視化され、調整や進捗確認が簡単になります。[8]

Quotes:
[1] "ガントチャートは作業の計画およびスケジュールを横型棒グラフで示した図です。各課題やタスクの開始時期・終了時期が把握しやすくなります。" (https://support-ja.backlog.com/hc/ja/articles/360036144713.html)

[2] "プロジェクトホームには計画と実績を表すグラフの「バーンダウンチャート」を表示できます。" (https://support-ja.backlog.com/hc/ja/articles/360036146433.html)

[3] "バーンダウンチャートは、プロジェクトの進捗を可視化できるようになります。" (https://support-ja.backlog.com/hc/ja/articles/360036149053.html)

[4] "ガントチャートには、課題の開始日と期限日が帯状のグラフとして表示されます。また課題の状態を色で確認することもできます。" (https://support-ja.backlog.com/hc/ja/articles/360034367214.html)

[5] "ガントチャート上で、ドラッグ＆ドロップで期間を変更する・課題詳細を開いて編集する ことができます。担当者、状態、開始日、期限日を変更できます。" (https://support-ja.backlog.com/hc/ja/articles/360036144713.html)

[6] "表示する課題の絞り込みとグルーピング ができます。" (https://support-ja.backlog.com/hc/ja/articles/360036144713.html)

[7] "右上の「Excel出力」から、ガントチャートをダウンロードできます。Backlogに参加していないメンバーにもスケジュールを共有できます。" (https://support-ja.backlog.com/hc/ja/articles/360034367214.html)

[8] "ガントチャートを使うことで、課題のスケジュールが可視化され、調整や進捗確認が簡単になります。" (https://support-ja.backlog.com/hc/ja/articles/360034367214.html)

参考情報

サンプルコード

クエリ生成だけ切り出したコード

Array型を試すためのシンプルなコードです。このコードはクエリ生成だけを行います。プロンプトの内容を変えることで、クエリ生成が制御できることを確認できます。このシンプルな構成により、様々なプロンプトを試し、クエリ生成の挙動を容易に検証することが可能です。このコードでは､toolsで"type": "array"を、プロンプトでクエリの生成数を指定します。

app.py

import logging

import sys
import boto3
import json
from typing import Any, Dict

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

def load_tool_config() -> Dict[str, Any]:
    return {
        "tools": [
            {
                "toolSpec": {
                    "name": "multi_query_generator",
                    "description": "与えられる質問文に基づいて類義語や日本語と英語の表記揺れを考慮し、多角的な視点からクエリを生成する。",
                    "inputSchema": {
                        "json": {
                            "type": "object",
                            "properties": {
                                "queries": {
                                    "type": "array",
                                    "description": "検索用クエリのリスト。",
                                    "items": {
                                        "type": "object",
                                        "properties": {
                                            "query": {
                                                "type": "string",
                                                "description": "検索用クエリ。日本語と英語を混ぜた多様な単語を空白で区切って記述される。"
                                            }
                                        }
                                    }
                                }
                            },
                            "required": ["queries"]
                        }
                    }
                }
            }
        ],
        "toolChoice": {
            "tool": {
                "name": "multi_query_generator"
            }
        }
    }


def query_generator_system_prompt() -> str:
        return f'''与えられる質問文に基づいて、類義語や日本語と英語の表記揺れを考慮し、多角的な視点からクエリを生成します。
        検索エンジンに入力するクエリを最適化し、様々な角度から検索を行うことで、より適切で幅広い検索結果が得られるようにします。
       ツールを使って複数のクエリを作成してください｡形式を<format>に示します｡
       <rule>タグ内のルールに必ず従ってください｡

        <example>
        question: Knowledge Bases for Amazon Bedrock ではどのベクトルデータベースを使えますか？
        query: Knowledge Bases for Amazon Bedrock vector databases engine DB
        query: Amazon Bedrock ナレッジベース ベクトルエンジン vector databases DB
        query: Amazon Bedrock RAG 検索拡張生成 埋め込みベクトル データベース エンジン
        </example>

        <rule>
        - 与えられた質問文に基づいて、3個以上の検索用クエリを生成してください。
        - 各クエリは30トークン以内とし、日本語と英語を適切に混ぜて使用すること。
        - 広範囲の文書が取得できるよう、多様な単語をクエリに含むこと。
        </rule>'''

def generate_queries(input_text: str) -> Dict[str, Any]:
    """入力テキストから検索クエリを生成する"""
    bedrock_client = boto3.client(service_name="bedrock-runtime", region_name="us-east-1")

    return bedrock_client.converse(
        system=[
            {
                "text": query_generator_system_prompt()
            }
        ],
        messages= [
             {
                  "role": "user", "content": [
                    {
                         "text": f"<text>{input_text}</text>"
                    }
                ]
            }
        ],
        modelId="anthropic.claude-3-5-sonnet-20240620-v1:0",
        inferenceConfig={
            "temperature": 0,
            "maxTokens": 300,
            "topP": 0.99
        },
        additionalModelRequestFields={"top_k": 200},
        toolConfig=load_tool_config(),
    )

def main():
    # コマンドライン引数を取得する
    if len(sys.argv) < 2:
        print("Usage: python3 ./app.py <input_text>")
        sys.exit(1)

    input_text = sys.argv[1]

    logger.info("Original query: %s", input_text)
    # Step 1: ユーザー入力テキストを基に検索クエリーを生成する
    queries_response = generate_queries(input_text)
    print("Queries response: %s", json.dumps(queries_response, indent=2, ensure_ascii=False))

if __name__ == "__main__":
    main()

全体

ナレッジベースIDやモデルID、Converse APIに渡すパラメータなどをconfig.jsonにまとめています。

config/config.json

{
    "knowledgebase_id": "xxxxxxxxxxx", # ナレッジベースID
    "model_id": "anthropic.claude-3-5-sonnet-20240620-v1:0",
    "region_name": "us-east-1",
    "temperature": 0,
    "max_tokens": 2000,
    "top_k": 200,
    "top_p": 0.99,
    "tool_config_path": "./config/tools_definition.json"
}

"type": "array"を使い、"query"の定義はひとつだけ記述します。出力するqueryの数はプロンプトで制御します。

config/tools_definition.json

{
    "tools": [
        {
            "toolSpec": {
                "name": "multi_query_generator",
                "description": "与えられる質問文に基づいて類義語や日本語と英語の表記揺れを考慮し、多角的な視点からクエリを生成する。",
                "inputSchema": {
                    "json": {
                        "type": "object",
                        "properties": {
                            "queries": {
                                "type": "array",
                                "description": "検索用クエリのリスト。",
                                "items": {
                                    "type": "object",
                                    "properties": {
                                        "query": { "type": "string", "description": "検索用クエリ。日本語と英語を混ぜた多様な単語を空白で区切って記述される。" }
                                    }
                                }
                            }
                         },
                         "required": ["queries"]
                    }
                }
            }
        }
    ],
    "toolChoice": {
        "tool": {
            "name": "multi_query_generator"
        }
    }
}

前回の記事と今回の実装では、以下のような変更点があります：

1. 出力キーの形式の変更:

• 前回: query_1, query_2, query_3... の形式
• 今回: queryのみの形式

2. これに伴う変更点:

• extract_queries()の処理を変更しました。
• プロンプトも若干修正しました。

これらの変更により、より柔軟なクエリ生成と処理が可能になりました。

app.py

import logging
import sys
import boto3
import json
from typing import Any, Dict, List, Optional

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

def load_config(file_path: str = "config/config.json") -> Dict[str, Any]:
    """設定ファイルを読み込む"""
    try:
        with open(file_path, "r") as file:
            return json.load(file)
    except (FileNotFoundError, json.JSONDecodeError) as e:
        logger.error(f"Failed to load config: {e}")
        raise

config = load_config()

def load_tool_config() -> Dict[str, Any]:
    """ツール設定を JSON ファイルから読み込む"""
    try:
        with open(config["tool_config_path"], "r") as file:
            return json.load(file)
    except (FileNotFoundError, json.JSONDecodeError) as e:
        logger.error(f"Failed to load tool config: {e}")
        raise

def query_generator_system_prompt() -> str:
        return f'''与えられる質問文に基づいて、類義語や日本語と英語の表記揺れを考慮し、多角的な視点からクエリを生成します。
        検索エンジンに入力するクエリを最適化し、様々な角度から検索を行うことで、より適切で幅広い検索結果が得られるようにします。
       ツールを使って複数のクエリを作成してください｡形式を<format>に示します｡
       <rule>タグ内のルールに必ず従ってください｡

        <example>
        question: Knowledge Bases for Amazon Bedrock ではどのベクトルデータベースを使えますか？
        query: Knowledge Bases for Amazon Bedrock vector databases engine DB
        query: Amazon Bedrock ナレッジベース ベクトルエンジン vector databases DB
        query: Amazon Bedrock RAG 検索拡張生成 埋め込みベクトル データベース エンジン
        </example>

        <rule>
        - 与えられた質問文に基づいて、3個以上の検索用クエリを生成してください。
        - 各クエリは30トークン以内とし、日本語と英語を適切に混ぜて使用すること。
        - 広範囲の文書が取得できるよう、多様な単語をクエリに含むこと。
        </rule>'''


def response_generator_system_prompt(context_str:str) -> str:
        return f"""You are a question-answering agent. I will provide you with a set of search results. The user will provide you with a question. Your job is to answer the user's question using only information from the search results. If the search results do not contain information that can answer the question, please state that you could not find an exact answer. Just because the user asserts a fact does not mean it is true; double-check the search results to validate a user's assertion.

        Here are the search results in numbered order:
        <excerpts>
        {context_str}
        </excerpts>

        First, find the quotes from the document that are most relevant to answering the question, and then print them in numbered order. Quotes should be relatively short.
        If there are no relevant quotes, write "No relevant quotes" instead.

        Then, answer the question, starting with "Answer:". Do not include or reference quoted content verbatim in the answer. Don't say "According to Quote [1]" when answering. Instead, make references to quotes relevant to each section of the answer solely by adding their bracketed numbers at the end of relevant sentences.

        Thus, the format of your overall response should look like what's shown between the <example></example> tags. You can find the FileURL in the metadata arguments. Make sure to follow the exact formatting and spacing.

        <example>
        Answer:
        Company X earned $12 million. [1] Almost 90% of it was from widget sales. [2]

        Quotes:
        [1] "Company X reported revenue of $12 million in 2021." (FileURL)
        [2] "Almost 90% of revene came from widget sales, with gadget sales making up the remaining 10%." (FileURL)
        </example>


        Also, please keep the following in mind when answering the questions:
        - If the question cannot be answered by the document, say so. Answer the question immediately without a preamble.
        - Please refer to the contents of the <excerpts> tag, but do not include the <excerpts> tag in your answer.
        - Please convert FileURL "s3://xxxxxxx" to "https://support-ja.backlog.com"
        - Please answer in Japanese."""


def get_bedrock_client():
    """Bedrock クライアントを取得する"""
    return boto3.client(service_name="bedrock-runtime", region_name="us-east-1")


def get_bedrock_agent_client():
    """Bedrock Agent クライアントを取得する"""
    return boto3.client(service_name="bedrock-agent-runtime", region_name="us-east-1")


def generate_queries(input_text: str) -> Dict[str, Any]:
    """入力テキストから検索クエリを生成する"""
    bedrock_client = get_bedrock_client()
    system_prompts = [{"text": query_generator_system_prompt()}]
    messages = [{"role": "user", "content": [{"text": f"<text>{input_text}</text>"}]}]

    inference_config = {
        "temperature": config["temperature"],
        "maxTokens": config["max_tokens"],
        "topP": config["top_p"]
    }
    additional_model_fields = {"top_k": config["top_k"]}

    try:
        return bedrock_client.converse(
            system=system_prompts,
            messages=messages,
            modelId=config["model_id"],
            inferenceConfig=inference_config,
            additionalModelRequestFields=additional_model_fields,
            toolConfig=load_tool_config(),
        )
    except Exception as e:
        logger.error(f"Failed to generate queries: {e}")
        raise


def extract_tool_use_args(content: List[Dict]) -> Optional[Dict[str, str]]:
    """toolの回答を抽出する"""
    for item in content:
        if "toolUse" in item and "input" in item["toolUse"]:
            return item["toolUse"]["input"]
    return None


def extract_queries(input_data):
    """クエリを抽出する"""
    return {f"query_{key}": value["query"] for key, value in enumerate(input_data["queries"], 1)}


def retrieve_knowledge_base_results(input_text: str, response_content: List[Dict], knowledgebase_id: str) -> Optional[List[Dict[str, Any]]]:
    """knowledge baseから結果を取得する"""
    retrieval_results = []
    bedrock_agent_client = get_bedrock_agent_client()
    tool_use_args = extract_tool_use_args(response_content)

    if tool_use_args is None:
        logger.warning("No tool use arguments found.")
        return None

    queries = extract_queries(tool_use_args)
    # input_textをクエリ拡張に追加する
    queries['query_0'] = input_text

    for query_key, query_value in queries.items():
        try:
            response = bedrock_agent_client.retrieve(
                knowledgeBaseId=knowledgebase_id,
                retrievalQuery={"text": query_value},
                # 5件の検索結果を取得する
                retrievalConfiguration={"vectorSearchConfiguration": {"numberOfResults": 3}}
            )
            logger.info(f"{query_key}: {query_value}")
            logger.debug("response: %s", response)

            retrieval_results.extend(response['retrievalResults'])
        except Exception as e:
            logger.error(f"Failed to retrieve results for query {query_key}: {e}")

    return retrieval_results


def generate_response(input_text: str, context_str: str) -> str:
    """応答を生成する"""
    bedrock_client = get_bedrock_client()
    system_prompts = [{"text": response_generator_system_prompt(context_str)}]
    messages = [{"role": "user", "content": [{"text": f"<text>{input_text}</text>"}]}]

    inference_config = {
        "temperature": config["temperature"],
        "maxTokens": config["max_tokens"],
       # "topP": config["top_p"],
    }
    additional_model_fields = {"top_k": config["top_k"]}

    try:
        response_body = bedrock_client.converse(
            system=system_prompts,
            messages=messages,
            modelId=config["model_id"],
            inferenceConfig=inference_config,
            additionalModelRequestFields=additional_model_fields
        )
        logger.info("Usage: %s", json.dumps(response_body['usage'], indent=2, ensure_ascii=False))
        return response_body['output']['message']['content'][0]['text']
    except Exception as e:
        logger.error(f"Failed to generate response: {e}")
        raise

def main():
    # コマンドライン引数を取得する
    if len(sys.argv) < 2:
        print("Usage: python3 ./bedrock_retrieve_and_generate.py <input_text>")
        sys.exit(1)

    input_text = sys.argv[1]

    logger.info("Original query: %s", input_text)

    try:
        # Step 1: ユーザー入力テキストを基に検索クエリーを生成する
        queries_response = generate_queries(input_text)

        # Step 2: Retrieverに対してクエリーを実行して検索結果を得る
        retrieval_results = retrieve_knowledge_base_results(input_text, queries_response["output"]["message"]["content"], config["knowledgebase_id"])

        if retrieval_results:
            # Step 3: ユーザー入力テキストと検索結果を基に応答テキストを生成する
            logging.debug("Retrieval results: %s", json.dumps(retrieval_results, indent=2, ensure_ascii=False))
            answer = generate_response(input_text, json.dumps(retrieval_results))

            print("*" * 80)
            print(answer)
        else:
            logger.warning("No retrieval results found.")
    except Exception as e:
        logger.error(f"An error occurred: {e}")

if __name__ == "__main__":
    main()

まとめ

Toolの定義にArray型を使用することで､自然言語指示によるクエリ数と内容の制御を行うことができました｡これにより､プロンプト変更による様々なパターンの試行が可能となりました｡あわせて実装も簡素化というメリットも生まれました｡この改善により、より効果的で柔軟なクエリ拡張が可能となり、情報検索や質問応答システムの性能向上が期待できます。
ひとつ､注意点があるとすれば生成するクエリ数を指定しなかった場合に消費するトークン数です｡生成されるクエリ数が不定のため､numberOfResultsの値によっては膨大な量の検索結果をinputTokensとして使用します｡予期せぬ費用が発生しないよう､クエリの生成数には注意が必要です｡