[アップデート] Amazon Bedrock Knowledge Bases がクロスリージョン推論をサポートしました

[アップデート] Amazon Bedrock Knowledge Bases がクロスリージョン推論をサポートしました

Amazon Bedrock Knowledge Bases がクロスリージョン推論をサポートしました。モデルを使った検索精度を向上させる機能がいくつかあるのですが、各機能にどのような動きになるのか確認してみました。
Clock Icon2024.09.16

こんにちは!AWS 事業本部コンサルティング部のたかくに(@takakuni_)です。

Amazon Bedrock Knowledge bases がクロスリージョン推論をサポートしました。

https://aws.amazon.com/about-aws/whats-new/2024/09/amazon-bedrock-knowledge-bases-cross-region-inference/

クロスリージョン推論 is 何?何が嬉しいの?という方はこちらも合わせてご覧ください。

https://dev.classmethod.jp/articles/amazon-bedrock-cross-region-inference-support/

おさらい

Amazon Bedrock Knowledge bases では、モデルの推論を行う機能が多く搭載されています。

  • エンべディングモデル
    • データソースのデータをベクトル化、ベクトルデータベースへのデータ保管
    • セマンティックチャンキングを利用した文どうしのコサイン類似度の比較
    • プロンプトのベクトル化
  • テキスト生成モデル
    • RetrieveAndGenerate API でベクトル検索後の回答文の生成
    • RetrieveAndGenerate API で sessionId を指定した場合のクエリ書き換え
    • Advanced Parsing オプションによる非テキスト情報(グラフや表)の解析
    • クエリ分解機能によるプロンプトのクエリ分解
    • Chat with your document 機能を使った RAG の提供

GA してからアップデートを追うようになりましたが、凄まじい成長具合ですね。

余談は置いといて、今回のクロスリージョン推論は「テキスト生成モデル」のシーンでサポートしました。おそらく、推論プロファイル自体が現時点で Claude 3、3.5 ファミリーのみサポートしているためだと思われます。

https://docs.aws.amazon.com/bedrock/latest/userguide/cross-region-inference-support.html

やってみる

それではクロスリージョン推論モリモリなナレッジベースを作ってみましょう。今回利用した Terraform はこちらです。

https://github.com/takakuni-classmethod/genai-blog/tree/main/knowledge_bases_cross_region_inference

今回は作成したナレッジベースを利用し、高度な解析オプション、sessionId を利用したクエリ書き換え、クエリ分解を試してみます。

マネジメントコンソールから確認

作成されたリソースを、マネジメントコンソールから見てみます。

高度な解析オプションのモデルがどちらも選択されていないような表示です。

2024-09-15 at 21.31.10-Amazon Bedrock@2x.png

ちなみに、推論プロファイルを指定せず、リージョナルに用意されたモデル ID を指定した場合はモデルが選ばれて表示されます。

2024-09-15 at 21.36.26-Amazon Bedrock@2x.png

それ以外は、とくに違いが見当たりませんでした。

データソースの同期

おなじみだとは思いますが、いつものようにデータソースの同期を行ってみます。

2024-09-15 at 21.37.59-Amazon Bedrock@2x.png

高度な解析オプション

高度な解析オプションで推論プロファイルを利用する場合は、ナレッジベースのサービスロールに追加の権限が必要です。

従来は次のポリシーに従い IAM 設計を行ってきましたが、

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

Prerequisites for cross-region inference に記載されている、以下の IAM ポリシーを追加する必要があります。

{
	"Version": "2012-10-17",
	"Statement": [
		{
			"Effect": "Allow",
			"Action": [
				"bedrock:InvokeModel*",
				"bedrock:GetInferenceProfile",
				"bedrock:ListInferenceProfiles"
			],
			"Resource": [
				"arn:aws:bedrock:*::foundation-model/*",
				"arn:aws:bedrock:*:*:inference-profile/*"
			]
		}
	]
}

https://docs.aws.amazon.com/bedrock/latest/userguide/cross-region-inference-prereq.html

Actions に GetInferenceProfileListInferenceProfiles の追加と InvokeModel のリソースを絞っている場合は、推論プロファイルの ARN を追加してあげましょう。

ログを見てみる

まずは CloudWatch Logs でログを見てみます。とくに目新しいことはないのですが、推論プロファイル特有の inferenceRegion キーがログに記録されています。

私がデプロイしたのは us-west-2 リージョンのため、推論プロファイル側で us-east-1 にルーティングしていることがわかります。

ModelInvocationLog.json
{
	"schemaType": "ModelInvocationLog",
	"schemaVersion": "1.0",
	"timestamp": "2024-09-16T01:07:18Z",
	"accountId": "123456789012",
	"identity": {
		"arn": "arn:aws:sts::123456789012:assumed-role/infrnc-prfl-kb-role/XQ0MMWVNRU-DocumentLoaderTask-602fadb1-326a-4764-8d53-da8aa1487"
	},
	"region": "us-west-2",
	"requestId": "17829a89-3d2d-4466-b3c5-9b7f816f368c",
	"operation": "InvokeModel",
	"modelId": "arn:aws:bedrock:us-west-2:123456789012:inference-profile/us.anthropic.claude-3-sonnet-20240229-v1:0",
	"input": {
		"inputContentType": "application/json",
		"inputBodyJson": {
			"anthropic_version": "bedrock-2023-05-31",
			"messages": [
				{
					"role": "user",
					"content": [
						{
							"type": "text",
							"text": "Transcribe the text content from an image page and output in Markdown syntax (not code blocks). Follow these steps:\n\n1. Examine the provided page carefully.\n\n2. Identify all elements present in the page, including headers, body text, footnotes, tables, visulizations, captions, and page numbers, etc.\n\n3. Use markdown syntax to format your output:\n    - Headings: # for main, ## for sections, ### for subsections, etc.\n    - Lists: * or - for bulleted, 1. 2. 3. for numbered\n    - Do not repeat yourself\n\n4. If the element is a visulization\n    - Provide a detailed description in natural language\n    - Do not transcribe text in the visualization after providing the description\n\n5. If the element is a table\n    - Create a markdown table, ensuring every row has the same number of columns\n    - Maintain cell alignment as closely as possible\n    - Do not split a table into multiple tables\n    - If a merged cell spans multiple rows or columns, place the text in the top-left cell and output ' ' for other\n    - Use | for column separators, |-|-| for header row separators\n    - If a cell has multiple items, list them in separate rows\n    - If the table contains sub-headers, separate the sub-headers from the headers in another row\n\n6. If the element is a paragraph\n    - Transcribe each text element precisely as it appears\n\n7. If the element is a header, footer, footnote, page number\n    - Transcribe each text element precisely as it appears\n\nOutput Example:\n\nA bar chart showing annual sales figures, with the y-axis labeled \"Sales ($Million)\" and the x-axis labeled \"Year\". The chart has bars for 2018 ($12M), 2019 ($18M), 2020 ($8M), and 2021 ($22M).\nFigure 3: This chart shows annual sales in millions. The year 2020 was significantly down due to the COVID-19 pandemic.\n\n# Annual Report\n\n## Financial Highlights\n\n* Revenue: $40M\n* Profit: $12M\n* EPS: $1.25\n\n\n| | Year Ended December 31, | |\n| | 2021 | 2022 |\n|-|-|-|\n| Cash provided by (used in): | | |\n| Operating activities | $ 46,327 | $ 46,752 |\n| Investing activities | (58,154) | (37,601) |\n| Financing activities | 6,291 | 9,718 |\n\nHere is the image.\n"
						},
						{
							"type": "image",
							"source": {
								"type": "base64",
								"media_type": "image/png",
								"data": "省略"
							}
						},
						{
							"type": "text",
							"text": "\nBased on the image provided, try to format the context into markdown format.\n<context>\n会社概要\n\n\n</div>\n</context>\n"
						}
					]
				}
			],
			"max_tokens": 2048,
			"temperature": 0.0,
			"top_p": 1.0,
			"stop_sequences": ["\nObservation"],
			"top_k": 50
		},
		"inputTokenCount": 1393
	},
	"output": {
		"outputContentType": "application/json",
		"outputBodyJson": {
			"id": "msg_bdrk_01MGL4ovqy6yViVtbqkedNdH",
			"type": "message",
			"role": "assistant",
			"model": "claude-3-sonnet-20240229",
			"content": [{ "type": "text", "text": "# 会社概要" }],
			"stop_reason": "end_turn",
			"stop_sequence": null,
			"usage": { "input_tokens": 1393, "output_tokens": 9 }
		},
		"outputTokenCount": 9
	},
+	"inferenceRegion": "us-east-1"
}

ルーティングはファイル単位かスライド単位か

結論、スライド単位で利用されるリージョンのルーティングが行われます。次は 1 つの PDF ファイルを解析した際の inferenceRegion の結果です。

2024-09-16 at 10.53.30-ログ  CloudWatch  us-west-2@2x.png

次のような構造で解析が行われると予想されるため、 inferenceRegion がスライド単位で異なるのではないかと思います。

Advanced Parsing 予想図.png

RetieveAndGenerate

それでは、 RetieveAndGenerate API でモデル実行してみます。まずはベーシックなパターンから。

01.py
import boto3

KNOWLEDGEBASE_ID = "XXXXXXXX"  # ナレッジベース ID を記載
model_arn = "arn:aws:bedrock:us-west-2:123456789012:inference-profile/us.anthropic.claude-3-5-sonnet-20240620-v1:0"
client = boto3.client("bedrock-agent-runtime", region_name="us-west-2")

prompt = (
    "20代社員は何名ほどでしょうか?"
)

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

print(response["output"]["text"],flush=False)

ログの内容です。推論先のリージョンは us-east-1 だったり us-west-2 だったりしました。

ModelInvocationLog.json
{
    "schemaType": "ModelInvocationLog",
    "schemaVersion": "1.0",
    "timestamp": "2024-09-16T03:59:26Z",
    "accountId": "123456789012",
    "identity": {
        "arn": "arn:aws:sts::123456789012:assumed-role/takakuni/botocore-session-1726457820"
    },
    "region": "us-west-2",
    "requestId": "70d12b27-3527-4c09-809d-d9ff410d12d0",
    "operation": "Converse",
    "modelId": "arn:aws:bedrock:us-west-2:123456789012:inference-profile/us.anthropic.claude-3-5-sonnet-20240620-v1:0",
    "input": {
        "inputContentType": "application/json",
        "inputBodyJson": {
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "text": "20代社員は何名ほどでしょうか?"
                        }
                    ]
                }
            ],
            "system": [
                {
                    "text": "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 to the question. Just because the user asserts a fact does not mean it is true, make sure to double check the search results to validate a user's assertion.\n\nHere are the search results in numbered order:\n<search_results>\n<search_result>\n<content>\n# 会社概要\n</content>\n<source>\n1\n</source>\n</search_result>\n<search_result>\n<content>\n# 会社概要\n会社名: サンプルサンプル株式会社\n所在地: サンプル県サンプル市サンプル村 サンプル 1-2-3\n代表: サン・プル\n設立: 2010年04月\n従業員数: 12,345名 (グループ全体)\n事業:\n* 宇宙旅行事業\n* AI搭載ロボットペット事業\n* 宇宙マーケティング戦略\n</content>\n<source>\n2\n</source>\n</search_result>\n<search_result>\n<content>\n# 社員数\nThe image shows a bar chart depicting the employee count over multiple fiscal years. The y-axis represents the employee count, while the x-axis shows the fiscal years from FY2010 to FY2023.\nThe bars in the chart display a consistent upward trend, indicating a steady increase in the employee count over the years. The chart also includes a note stating that the employee count surpassed 10,000 in FY2022, as indicated by the sharp rise in the bar for that fiscal year.\n</content>\n<source>\n3\n</source>\n</search_result>\n<search_result>\n<content>\n# 社員構成(年代別)\nA pie chart showing the employee composition by age group. The chart is divided into the following sections:\n| Age Group | Percentage |\n|-|-|\n| 20代 | 24.0% |\n| 30代 | 40.0% |\n| 40代 | 24.0% |\n| 50代 | 8.0% |\n| 60代以上 | 4.0% |\n\n</content>\n<source>\n4\n</source>\n</search_result>\n\n</search_results>\n\nYou should provide your answer without any inline citations or references to specific sources within the answer text itself. Do not include phrases like \"according to source X\", \"[1]\", \"[source 2, 3]\", etc within your <text> tags.\n\nHowever, you should include <sources> tags at the end of each <answer_part> to specify which source(s) the information came from.\nNote that <sources> may contain multiple <source> if you include information from multiple results in your answer.\n\nDo NOT directly quote the <search_results> in your answer. Your job is to answer the user's question as concisely as possible.\n\nYou must output your answer in the following format. Pay attention and follow the formatting and spacing exactly:\n<answer>\n<answer_part>\n<text>\nfirst answer text\n</text>\n<sources>\n<source>source ID</source>\n</sources>\n</answer_part>\n<answer_part>\n<text>\nsecond answer text\n</text>\n<sources>\n<source>source ID</source>\n</sources>\n</answer_part>\n</answer>\n"
                }
            ],
            "inferenceConfig": {
                "maxTokens": 2048,
                "temperature": 0,
                "topP": 1,
                "stopSequences": [
                    "\nObservation"
                ]
            }
        },
        "inputTokenCount": 892
    },
    "output": {
        "outputContentType": "application/json",
        "outputBodyJson": {
            "output": {
                "message": {
                    "role": "assistant",
                    "content": [
                        {
                            "text": "<answer>\n<answer_part>\n<text>\n20代の社員数を正確に計算するには、まず全従業員数と20代の割合の情報が必要です。\n\n会社の従業員数は12,345名(グループ全体)です。また、社員構成の内訳によると、20代の社員は全体の24.0%を占めています。\n\nこれらの情報を基に計算すると、20代の社員数は約2,963名(12,345 × 0.24)となります。\n\nただし、この数字は概算であり、実際の数字は若干異なる可能性があります。また、この計算はグループ全体の従業員数を基にしているため、特定の部門や地域での20代の社員数とは異なる可能性があることにご注意ください。\n</text>\n<sources>\n<source>2</source>\n<source>4</source>\n</sources>\n</answer_part>\n</answer>"
                        }
                    ]
                }
            },
            "stopReason": "end_turn",
            "metrics": {
                "latencyMs": 5848
            },
            "usage": {
                "inputTokens": 892,
                "outputTokens": 270,
                "totalTokens": 1162
            }
        },
        "outputTokenCount": 270
    },
    "inferenceRegion": "us-east-1"
}

続いて、 sessionId を引き継ぎ、クエリ書き換えをしてみます。

02.py
import boto3

KNOWLEDGEBASE_ID = "XXXXXXXX"  # ナレッジベース ID を記載
model_arn = "arn:aws:bedrock:us-west-2:123456789012:inference-profile/us.anthropic.claude-3-5-sonnet-20240620-v1:0"
client = boto3.client("bedrock-agent-runtime", region_name="us-west-2")
+ session_id = ""

prompt = (
    "20代社員は何名ほどでしょうか?"
)

response = client.retrieve_and_generate(
    input={"text": prompt},
    retrieveAndGenerateConfiguration={
        "type": "KNOWLEDGE_BASE",
        "knowledgeBaseConfiguration": {
            "knowledgeBaseId": KNOWLEDGEBASE_ID,
            "modelArn": model_arn,
        },
    },
)
print(response["output"]["text"],flush=False)
+ session_id = response["sessionId"]
+
+ # 書き換え
+ prompt = (
+     "30代社員は何名ほどでしょうか?"
+ )
+
+ response_with_session_id =  client.retrieve_and_generate(
+     input={"text": prompt},
+     retrieveAndGenerateConfiguration={
+         "type": "KNOWLEDGE_BASE",
+         "knowledgeBaseConfiguration": {
+             "knowledgeBaseId": KNOWLEDGEBASE_ID,
+             "modelArn": model_arn,
+         },
+     },
+     sessionId=session_id
+ )
+
+ print(response_with_session_id["output"]["text"],flush=False)

確認すると、クエリ書き換えと回答生成は別のリージョンが使われている(分離している)ことがわかりました。

2024-09-16 at 13.37.18-ログ  CloudWatch  us-west-2.png

クエリ例
fields @timestamp, @message, @logStream, @log
| sort @timestamp desc
| filter modelId != "arn:aws:bedrock:us-west-2::foundation-model/amazon.titan-embed-text-v2:0"
| display @message, inferenceRegion
| limit 200

最後にクエリ分解機能によるプロンプトのクエリ分解を行います。

03.py
import boto3

KNOWLEDGEBASE_ID = "XXXXXXXX"  # ナレッジベース ID を記載
model_arn = "arn:aws:bedrock:us-west-2:123456789012:inference-profile/us.anthropic.claude-3-5-sonnet-20240620-v1:0"
client = boto3.client("bedrock-agent-runtime", region_name="us-west-2")

prompt = (
    "20代社員は何名ほどでしょうか?組織の何%を占めていますか?"
)

response = client.retrieve_and_generate(
    input={"text": prompt},
    retrieveAndGenerateConfiguration={
        "type": "KNOWLEDGE_BASE",
        "knowledgeBaseConfiguration": {
            "knowledgeBaseId": KNOWLEDGEBASE_ID,
            "modelArn": model_arn,
+           "orchestrationConfiguration": {
+               "queryTransformationConfiguration": {"type": "QUERY_DECOMPOSITION"}
+           },
        },
    },
)
print(response["output"]["text"],flush=False)

こちらもクエリの分解と、回答生成で利用するリージョンが異なるケースが確認できました。

2024-09-16 at 13.49.11-CloudWatch  us-west-2.png

クエリ例
fields @timestamp, @message, @logStream, @log
| sort @timestamp desc
| filter modelId != "arn:aws:bedrock:us-west-2::foundation-model/amazon.titan-embed-text-v2:0"
| display @message, inferenceRegion, input.inputBodyJson.messages.0.content.0.text
| limit 200

まとめ

以上、「Amazon Bedrock Knowledge bases がクロスリージョン推論をサポートしました」でした。

モデルの実行ごとに、クロスリージョン推論を実行しているため、レスポンスが少し気になるところですが、クォータ性能に強いシステムを作りたいケースでは、非常に有効的な手段だと思いました。このブログがどなたかの参考になれば幸いです。

AWS 事業本部コンサルティング部のたかくに(@takakuni_)でした!

この記事をシェアする

facebook logohatena logotwitter logo

© Classmethod, Inc. All rights reserved.