Bedrock Agentsのカスタムオーケストレーション
こんにちは、ふくちと申します。
2024年11月28日、Amazon Bedrock Agents がカスタムオーケストレーションをサポートするようになりましたというアップデートがありました！
このカスタムオーケストレーション機能とは、「Bedrockエージェントの行動順序や判断プロセスをワークフローとして定義できるもの」となっております。
通常のBedrockエージェントであれば、タスクをどのように行うかはエージェント自身で決定します。
エージェントにタスク実行権を付与した場合は、柔軟な対応が可能ですが、毎回の動作や出力が安定しない可能性があります。
一方でこのカスタムオーケストレーション機能を用いれば、エージェントが毎回同じような流れでタスクを実行してくれるようになります。
システマチックにAIエージェントを動かすことができるというのは、2025年における1つの注目技術です。
それをAWS上で実現できるのが、このカスタムオーケストレーション機能です。
実装の中身を詳しく知りたい方は以下のブログをご参照ください。
実際に触ってみる
今回は以下のサンプルを実際に触ってみて、どのような挙動になっているのか試してみます！
こちらのサンプルでは、レストランのメニュー提示や予約を自律的に行ってくれるエージェントを構築していきます。
その際、従来のBedrockエージェントと、カスタムオーケストレーションを搭載したBedrockエージェントの2つを作成して比較を行います。
実際の出力を確認することで、両者にどのような違いがあるのか、学んでいきましょう！
そして今回構築するアーキテクチャとしては、以下のようになります。
- Bedrockエージェント
- カスタムオーケストレーション用Lambda関数
- エージェント用ツール
- 予約作業用Lambda関数
- 予約データ格納用DynamoDB
- お店のメニューなどを格納したS3と、それを参照可能なBedrockナレッジベース
質問内容に応じ、適切なツールを使い分けるようなエージェントを構築できます。
色々ファイルがありますが、使用するのは「custom_orchestration_example.ipynb」というファイル(以下、ノートブックと記載)です。これがドキュメント兼ソースコードを兼ね備えています。
やり方としては、VSCodeやCursorなどでリポジトリをクローンした後、ローカルでJupyter Notebookを立ち上げて各コマンドを実行していくような流れです。
(詳しいやり方は別途纏めます。)
ノートブック上で実行した後は、コンソール上からも同じプロンプトを投げて、トレースを確認します。
※トレース：エージェントが会話のその時点で返すレスポンスに至るまでの推論プロセスを出力したもの。詳しい見方は以下をご参照ください。
https://docs.aws.amazon.com/ja_jp/bedrock/latest/userguide/trace-events.html
コンソールで使用するモデルを有効化する
まずはBedrockコンソールにおいて、「Titan Text Embeddings V2」「Claude 3.5 Sonnet v2」のアクセスを有効化しておきましょう。
今回はデフォルトでこちらのモデルを使用していきます。
従来(ReAct)のエージェントを使ってみる
基本的にはノートブックの手順通りに進めるだけなので、細かい構築手順は省略します。
ここでは、実際のエージェントの挙動を見ていきましょう。
エージェントに与えられたプロンプトとしては、以下の通りです。(本来は英語ですが、日本語訳してあります)
あなたは'The Regrettable Experience'の予約を担当するレストランアシスタントです。
メニューについて話し、新しい予約を作成し、既存の予約の詳細を取得したり、
既存の予約を削除したりすることができます。
常に丁寧に応答し、返信の中でレストランの名前に言及してください。
新しい会話の開始時にレストランの名前を決して省略しないでください。
お客様が回答できないことについて質問された場合は、よりパーソナライズされた体験のために
以下の電話番号をご案内ください：+1 999 999 99 9999
お客様の質問に答えるのに役立つ情報：
The Regrettable Experience 住所：101W 87th Street, 100024, New York, New York
営業時間：
- 月曜日～金曜日：午前11時～午後2時、午後5時～午後10時
- 土曜日：午前11時～午後11時
- 日曜日：午前11時～午後8時
まず、「今夜7時に、2人で予約できますか？」と質問してみましょう。
Certainly! I'm pleased to confirm that
I've made a reservation for you at The Regrettable Experience.
Here are the details:
- Date: Tonight, December 30, 2024
- Time: 7:00 PM
- Number of guests: 2
- Reservation name: John
- Booking ID: 7914fb20
Your table is all set for a delightful evening at The Regrettable Experience.
If you need to make any changes or have any questions,
please don't hesitate to ask. Enjoy your meal!
CPU times: user 71.3 ms, sys: 14.1 ms, total: 85.4 ms
Wall time: 11.2 s
ということで、予約することができたようです。
本当に予約できているか、DynamoDBに確認しておきましょう。
きちんと予約されているようです！
また、ノートブック上ではなくコンソール上からエージェントに同様の依頼をすると、トレースを確認することができます。
そこで裏側の動作を確認すると、アクショングループが呼び出されていることがわかります。
これは、エージェントが「予約を行うにはアクショングループを用いなければいけない」ということを自ら考えて実行してくれているため、実現できています。
続いて、このレストランにはどんなメニューがあるのか聞いてみましょう。
Welcome to The Regrettable Experience!
For dinner, we offer a variety of delicious options.
Two of our popular entrees are:
1. Buffalo Chicken Wings:
These are served with celery sticks and blue cheese dressing.
2. Shrimp and Grits:
This dish features shrimp sautéed in a flavorful sauce,
served over creamy cheese grits.
Please note that these are just a couple of our dinner offerings.
We have more options available on our full menu.
If you'd like more detailed information about our dinner menu or
have any specific dietary requirements, please don't hesitate to ask.
CPU times: user 50.7 ms, sys: 5.96 ms, total: 56.6 ms
Wall time: 15.2 s
2つのメニューが人気だと教えてくれましたが、本当なのでしょうか。
ディナーメニューに関するドキュメントをS3に格納しているので確認してみましょう。
↓
前菜の最初に来ていた2品を教えてくれました。
間違ってはいないのですが…せっかくならメインディッシュを教えてほしいですよね。
ということで、少し質問を変えて再度聞いてみます。
ついでにトレースを確認するため、コンソール上で質問します。
回答内容、変わってないですね…
トレースを確認したところ、きちんとナレッジベースを参照していました。
ただ、相変わらず前菜しか教えてくれず…ここはナレッジベースの精度の問題もありそうです。
最後に、より複雑な依頼をしてみます。メニューの質問と予約依頼を同時に行ってみましょう。
Welcome to The Regrettable Experience!
Let me tell you about our dinner menu and your reservation:
Dinner Menu Highlights:
- Buffalo Chicken Wings: Classic wings with blue cheese dressing
- Shrimp and Grits: Succulent shrimp over creamy cheese grits
I've successfully made a reservation for you:
- Date: Tonight (December 30, 2024)
- Time: 9:00 PM
- Number of Guests: 4
- Booking ID: be92a26e
We look forward to serving you this evening!
If you need any further assistance, please don't hesitate to ask.
CPU times: user 115 ms, sys: 17.6 ms, total: 133 ms
Wall time: 14.8 s
ということでメニューの提示と予約の実行を実行してくれました！
DynamoDBにもきちんと予約されていました。
トレースも確認していきましょう。
まずはナレッジベースに問い合わせています。
続いて、アクショングループを用いて予約処理。
最後にまとめて回答をしてくれています。
ということできちんと回答を得ることができました！
なぜかこの複雑なプロンプトの際だけ、リクエストレートが高すぎますと怒られることがあります。
これの原因としては、従来のBedrockエージェントで用いられているReActにあります(恐らく)。
後ほど詳しく解説します。
カスタムオーケストレーションを搭載したBedrockエージェントを使ってみる
こちらも基本的にはノートブックの手順に沿って進めれば構築できます。
カスタムオーケストレーション用Lambda関数について
ただ、先ほどと違うのは、カスタムオーケストレーション用のLambda関数があること。
ではそんなカスタムオーケストレーションLambdaはどんな処理フローになっているかを見ていきましょう。
1. `lambda_handler` (エントリーポイント)
- イベントを受信し、状態を抽出
- `nextEvent`関数を呼び出して次のアクションを決定
2. `nextEvent` (状態管理)
各状態での処理:
- `START`
- `create_prompt`でプランニング用プロンプトを生成
- モデル呼び出し用のペイロードを準備
- `MODEL_INVOKED`
- `is_end_state`で会話終了判定
- 続行時は`execute_plan_on_generation`でツール実行を準備
- `TOOL_INVOKED`
- `continue_execution`で次のツール実行を判断
- ツール実行か、モデル再呼び出しを決定
3. プラン処理
`execute_plan_on_generation`:
- モデル出力からプランを抽出
- `get_tool_to_execute`でツール実行情報を取得
`get_tool_to_execute`:
- XMLプランをパース
- ステップタイプ（関数呼び出し/ループ）を判別
- ツール実行情報を生成
4. メッセージ管理
`construct_messages`:
- セッション履歴から会話を再構築
- プロンプト種別に応じてメッセージを構造化
`merge_conversation_turn`:
- 連続する同一ロールのメッセージを統合
- Bedrock Converse API用に整形
5. プロンプト制御
- プランニングプロンプト：実行計画の生成を指示
- サマリープロンプト：会話履歴に基づく応答生成を指示
ざっくり言うと、
- 最初に実行計画を生成
- 状態(ステート)に応じたアクションを実行する
- 最終回答を生成する
というような流れになっています。
出力確認
さて、ここからは実際の出力を確認していきましょう。
まずは先程と同様に、「今夜7時に、2人で予約できますか？」と質問してみます。
"Certainly! I've made a reservation for you for 2 people at 7pm tonight.
Your booking has been confirmed with the following details:
- Number of guests: 2
- Date: Tonight (2024-12-30)
- Time: 7:00 PM (19:00)
- Booking ID: e4e184dd
Is there anything else you'd like to know about your reservation or any changes you'd like to make?"
CPU times: user 60.6 ms, sys: 13.3 ms, total: 73.9 ms
Wall time: 13.9 s
ということで、2名で予約できたようです。こちらもDynamoDBで確認しておきます。
最後にトレースを確認します。
エージェントがアクショングループを用いていることがわかりますね。
また、Step1、Step2と、ワークフローに沿った形でエージェントがタスクを実行していっているのがわかると思います。
コンソール上でプロンプトだけ投げると、エラーとなっていました。
ノートブックに記載されているクエリを丸ごとコピペすると、私は上手くいきました。
続いて、レストランのメニューを聞いてみます。こちらもメインディッシュを聞いてみましょう。
"Certainly!
I'd be happy to tell you about the main dishes for dinner at The Regrettable Experience.
Based on the information from our menu, here are the entrees available for dinner:
1. Buffalo Chicken Wings
- Description: Classic buffalo wings served with celery sticks and blue cheese dressing.
- Allergens: Dairy (in blue cheese dressing), Gluten (in the coating), possible Soy (in the sauce).
2. Shrimp and Grits
- Description: Succulent shrimp saut\u00e9ed in a flavorful sauce, served over creamy cheese grits.
- Allergens: Dairy (in the cheese grits).
3. BBQ Ribs (information partially cut off, but included in the menu)
4. Steak with Garlic Butter
- Description: Grilled sirloin steak seasoned with garlic butter, accompanied by loaded mashed potatoes and saut\u00e9ed vegetables.
- Allergens: Dairy (in garlic butter), potential Soy (in marinade).
5. Vegetable Stir-Fry with Tofu
- Description: Stir-fried tofu and a colorful medley of vegetables in a savory soy-ginger sauce, served over jasmine rice.
- Allergens: Soy (in tofu and sauce).
These main dishes offer a variety of options, including meat, seafood, and vegetarian choices.
Each dish comes with its own unique flavors and accompaniments.
Is there any particular dish you'd like more information about, or do you have any questions regarding allergens or dietary preferences?"
CPU times: user 203 ms, sys: 18.9 ms, total: 221 ms
Wall time: 22.9 s
先ほどのReActよりも回答数が多くなりましたね。相変わらず1.と2.は前菜なのに出力されていますが…
そして3.のBBQ Ribsはドキュメントに存在しません。ちょっとナレッジベースの検索精度自体はかなりイマイチですね…
トレースも確認しておきましょう。ステップごとに区切られた上で、ナレッジベースに検索しているのがわかります。
カスタムオーケストレーションに従って動作しているのも確認できました。
そして最後に、メニューの質問と予約依頼を同時に行ってみましょう。
"Certainly!
I'd be happy to provide information about our dinner menu and make a reservation for you.
Regarding our dinner menu, The Regrettable Experience offers the following entrees:
1. BUFFALO CHICKEN WINGS
- Description: Classic buffalo wings served with celery sticks and blue cheese dressing.
- Allergens: Dairy (in blue cheese dressing), Gluten (in the coating), possible Soy (in the sauce).
2. SHRIMP AND GRITS
- Description: Succulent shrimp saut\u00e9ed in a flavorful sauce, served over creamy cheese grits.
- Allergens: Dairy (in the cheese grits).
These are just a couple of examples from our dinner menu.
We likely have more options available, but the full menu wasn't provided in the information I have access to.
Now, regarding your reservation request, I've successfully made a booking for you.
Here are the details:
- Number of guests: 4
- Date: Tonight (assuming today's date)
- Time: 9:00 PM
- Booking ID: cd481680
Your reservation has been confirmed.
Please make sure to arrive on time,
and if you need to make any changes or have any special requests,
don't hesitate to contact the restaurant directly.
They may ask for your booking ID, so please keep it handy.
Is there anything else you'd like to know about our menu or your reservation?"
CPU times: user 149 ms, sys: 18.3 ms, total: 167 ms
Wall time: 18.9 s
両方にきちんと回答してくれました！
DynamoDBにも、きちんと予約登録されていました。
コンソール上でも動作確認しておきます。
ノートブックと全く同じ形式でクエリを投げると、上手いこと動作しました。
こちらのトレースは、後ほど詳しく確認します。
両者の処理を比較してみる
さて、ここまでで2つのエージェントの出力を比較してきました。
御覧頂いた通り、概ね出力自体に差はありませんでした。
ただ、実は出力までの過程が大きく異なっています。
まず従来のBedrockエージェントですが、こちらはReActという枠組みが用いられています。
一方でカスタムオーケストレーションを搭載したBedrockエージェントには、ReWOOという枠組みが用いられています。
この違いについて、実際のトレースを確認しながら見ていきましょう！
通常のBedrockエージェントの処理
まずReActを用いているエージェントの処理についてです。
こちらは、毎回自身の実行結果を振り返る、という動作が入ります。
先ほどのように、「メニューを教えて。あと、今夜9時から4人で予約したい」というタスクを与えたとしましょう。
トレースステップ 1
ステップ1では、まずは以下の流れになります。図でいうと、赤枠で囲った部分を実行します。
- Thought：タスク計画として、ディナーメニューを調査しようと考えます
- Action：ナレッジベースに問い合わせ、その結果を出力します
ここからは実際のトレースを確認していきます。
まず、ユーザーからの入力をmodelInvocationInput、すなわち基盤モデル呼び出しインプットとしてエージェントが受け取ります。
まず、Thoughtはrationale(入力に基づいたエージェントの推論・思考プロセス) として出力されます。
下記では、まずナレッジベースを検索してその後予約を行う、という形でタスク計画を作成していますね。
続いてのActionは、トレースではinvocationInput(呼び出しまたはクエリするアクショングループまたはナレッジベースに入力される情報) として出力されます。
ここではどのナレッジベースを呼び出すのか、そしてどのような検索を行いたいのかといった情報が入力されています。
Action実行後、ナレッジベースから返ってきた返答がobservation(アクショングループまたはナレッジベースの結果または出力、またはユーザーへのレスポンス) として出力されます。
ここではナレッジベースにおいて、どのバケットにあるドキュメントのどこを参考したのかについて教えてくれています。
↓
(実際はもっと多いですが、省略します。)
トレースステップ 2
続いて、ナレッジベースからの出力を観察します。
観察後、再度プラン作成へと戻ります。
ここで言う観察とは、LLM(Bedrock)へナレッジベースからの出力をすべて渡して、そこからユーザーへの回答を生成してもらう、というような挙動です。
modelInvocationInputとして、先ほどのナレッジベースの出力がすべて含まれています。以下画像のtext配下がそれに該当します。
↓
そして最後には、「検索結果を直接引用せずに、ユーザーの質問にできるだけ簡潔に回答してください。」と記載されており、またフォーマットに則って回答してください、と記載されています。
上記を踏まえると、ナレッジベースからの出力をそのままユーザーに返すのではなく、質問に対してどれが最も適切な答えなのかをBedrockが考えた上で回答を生成してくれるような処理がなされている、ということがわかります。
トレースステップ 3
ここでは、ナレッジベースからの出力を受け取った後の処理が行われます。
そこでエージェントは、次のように考えてタスクを実行していきます。
- Thought：(ナレッジベースへの)検索結果には夕食メニューについての情報はありますが、今夜9時に4人で予約可能かどうかの表示がないため、予約作業が必要だと考えます
- Action：予約用ツールを用います
- Observe Results：予約IDを出力します
実際のトレースを見ていきましょう。
まず、これまでのやり取りをすべてエージェントが把握し、管理します。
modelInvocationInputに対してこれまでのやり取りがすべて入力されます。
入力自体はトレースステップ2の最後と同じなのですが、typeがORCHESTRATIONになっています。
↓
そしてその後、トレースステップ 1と同じように、rationaleとしてエージェントの思考回路が出力されます。
ここでは、次に予約を行う、というような出力がなされていますね。
続いてinvocationInputとして、今度はアクショングループへのインプット情報が出力されています。
最後にobservationとして予約IDが出力されています。
トレースステップ 4
最後に、アクショングループの出力を踏まえて、最終的な回答を生成します。
- Thought：質問回答に必要な情報はすべて揃ったと判断します
- Finished：思考を終了し、回答を返却します
先ほどのアクショングループの実行結果もすべてmodelInvocationInputとして入力されます。
↓
そしてobservationとして、最終的な回答を出力します。
ReActの特徴
以上を踏まえて、ReActでのエージェントの思考回路がなんとなくご理解いただけたかなと思います。
ここで1つ重要な点は、Bedrockの呼び出し回数です。
思い返すと、各トレースステップごとにBedrockを呼び出していたので、4回呼び出していた計算になります。
しかも、後半のトレースステップになればなるほど、Bedrockに渡すテキスト量は増えていました。
今回はあまり複雑な仕組みではなかったですが、エージェントにもっと多くのことを任せた時、インプットの量が膨大になってしまう恐れがあります。
そうすると、コストの増大や処理時間の遅延、回答精度低下などに繋がる恐れがあります。
特に、毎回Bedrockを呼び出して入力をすべて渡すというところで、先ほど述べたリクエストレート高すぎ問題が発生することがあります。
これは恐らく、毎トレースでBedrockを呼び出していることと、毎回のインプットが多すぎることに起因しています(確証はないですが…)。
上記のような問題を、カスタムオーケストレーションでは防ぐことができます。
カスタムオーケストレーション搭載Bedrockエージェントの処理
こちらはReWOOを用いているエージェントの処理についてです。
こちらは、毎回自身の実行結果を振り返る、という動作が入りません。
あらかじめ決められたワークフローに沿った形で、エージェントが動作していきます。
先ほど同様、「メニューを教えて。あと、今夜9時から4人で予約したい」というタスクを与えたとしましょう。
トレースステップ 1
まずは、タスク計画を作成します。
トレースを確認すると、まずはユーザーからの依頼やシステムプロンプトがmodelInvocationInputとしてBedrockへ入力されています。
ただし、ここの段階ですでにカスタムオーケストレーションのワークフローの中身も渡されているようです。
↓
↓
そして、続くrationale(入力に基づいたエージェントの推論・思考プロセス)のところでは、どのような流れでタスクを実行するかが明確に書かれていました。
簡単に要約すると、以下の通りです。
- フェーズ1. ナレッジベースでディナーメニュー検索
- フェーズ2. 検索結果の条件分岐
→メニュー情報がある：メニュー情報を保存
→メニュー情報がない：電話番号を案内 - フェーズ3. 予約日の設定
- フェーズ4. アクショングループのLambda関数で予約処理
- フェーズ5. レスポンス生成
→メニュー情報と予約結果の返却
重要なのは、これをトレースステップ1の段階ですべて決定していることです。
これ以降はこの計画に沿って処理を進めるだけです。
トレースステップ 2
ここでは、ナレッジベースでのディナーメニュー検索を行います。
フェーズ1と2に当たる部分ですね。
ただし、このステップ2と3は、どちらも「計画実行フェーズ」として一括りにされるようです。
invocationInput(呼び出しまたはクエリするアクショングループまたはナレッジベースに入力される情報)に対して、利用するナレッジベースとプロンプトを入力しています。
そしてナレッジベースからのレスポンスが、以下のobservationとなっています。
トレースステップ 3
ここでは、予約作業を行います。
フェーズ3, 4に当たります。
前述の通り、ナレッジベースの検索も、アクショングループの実行も、どちらも「計画実行フェーズ」として扱われます。
invocationInputに対して、利用するアクショングループと、パラメータを入力しています。
そして、observationとしてアクショングループの実行結果が返ってきています。
ここでは予約IDですね。
トレースステップ 4
これまでのタスク実行による出力を踏まえ、最終的な回答を生成します。
トレースでは、予定通りタスクが実行できたので、実行結果を纏めてtextとしてmodelInvocationInputに入力しています。
ナレッジベース・アクショングループ両方の出力を渡しています。
そしてその出力を踏まえて、最終的に回答を生成します。
ここではrationaleとして、回答が出力されていました。(なぜobservationでないのかはわかりませんでした…)
トレースステップ 5はログが出力されていただけでした。
カスタムオーケストレーション(ReWOO)の特徴
以上を踏まえて、こちらでのBedrockの呼び出し回数を踏まえると、トレースステップ1, 4の2回だけでした。
カスタムオーケストレーションを用いることで、トレースステップ 1で決定したタスク計画に沿ってエージェントが動くだけになります。
すなわち、動作する度にBedrockへ入力を渡す必要がないです。
したがって、処理の高速化と安定化、そしてコスト削減を実現することができます。
まとめ
エージェントとワークフローに関する話は、2024年の年末にかなり話題になったトピックでした。
そのワークフローをAWSで実装できるのは面白いサービスだと思います。
また、そのメリットを身を以て体感できたのも良かったです。
次はカスタムオーケストレーションLambdaを自分で書いてみるか、もしくは同じようなことができるらしいLangGraphに入門するかですね。
2025年は多分LangGraphとLLMOpsを中心に色々やると思います！