AWS Lambda と ServerlessAdvent Calendar 2024

【鍵はLambda】Amazon Bedrock Agentsのカスタムオーケストレーション機能を理解する

Posted at 2024-12-05

はじめに

当記事はAWS Lambda と Serverless Advent Calendar 2024 3日目の記事です！

Amazon Bedrock Agentsのアップデート

2024年11月28日、Amazon Bedrock Agents がカスタムオーケストレーションをサポートするようになりましたというアップデートがありました！

このカスタムオーケストレーションを使用するには、オーケストレーションロジックの概要を示すLambda関数が必要になる、とのこと。

であれば、良いアドベントカレンダーのネタになりそうだ！ということでキャッチアップしてみることにしました！

カスタムオーケストレーションとは？

公式ドキュメントによると、このカスタムオーケストレーションを実装することで

エージェントが複数のステップのタスクを処理し、決定を下し、ワークフローを実行する方法を完全に制御できる

ようになるそうです。

ざっくり言うと、エージェントの行動順序や判断プロセスをカスタマイズできるようになった、ということです。

以下のドキュメントに、カスタムオーケストレーションについて詳しく書かれていたのでこちらを読み進めていきます。

従来のBedrock Agent

従来のBedrock Agentでは、いわゆるReAct(Reasoning&Acting)によってオーケストレーションされていました。

ReActとは、LLMを用いて推論と行動を同時に行うための手法です。

具体的に言うと、以下の3つを繰り返すフローがReActです。

Reasoning：自身の行動と理由を推論すること
→推論は、人間で言う「考える」という行動に似ています。
Acting：その推論に基づいて行動すること
Observation：行動結果を観察し推論と行動を繰り返して、最終回答を出力する

ということで、LLM自身が考え、行動し、最終的に何かを決定して回答を出力してくれる、というようなものになっています。

Bedrock Agentでも同じことが行われています。
LLMが各ステップを推論し、そのステップごとに収集された情報に基づいて次のアクションを決定する、という反復的な意思決定プロセスです。

この手法は、各ステップごとに推論が行われます。
したがって、段階的に行動を考えていくようなワークフローに適しています。
ステップ後の結果を鑑みて計画を変更したり、行動が変化したり、というようなこともあるようです。

その一方で、かなり逐次的な構造であるため、複雑な計画を実行する場合に遅延を引き起こす可能性もあります。
1回1回の行動ごとに思案する時間が入るようなものなので、それは確かに時間がかかりますよね。

ReActの個人的なイメージとしては、Googleマップで調べながら自転車を漕いでいるようなイメージを持ちました。

行き先を決定してGoogleマップに登録する(実現したいことを決定する)
→自転車を漕ぐ(Acting)
→曲がり角でマップを確認する(Observation)
→別ルートのほうが近道そうだと考える(Reasoning)
→マップとは異なる道を選んで自転車を漕ぐ(Acting)
→曲がり角でマップを確認する(Observation)
→間違えていそうなので道を戻る(Reasoning)

曲がり角が来る度にマップを確認(Observation)していては、確かに時間はかかってしまいそうですよね。
ただし、エージェントが継続的に判断を見直しながら進められるという柔軟性が備わっています。

このReActが、これまでのBedrock Agentでは用いられていました。

参考：

Amazon Bedrock 生成AIアプリ開発入門 [AWS深掘りガイド]

カスタムオーケストレーション機能

それに対してカスタムオーケストレーションを用いると、ReWOO(Reasoning WithOut Observations)によってオーケストレーションされることになります。

ReWOOでは、ReActとは異なり、タスク実行後にLLMによるObservationが行われません。
タスク計画を事前に作成しておき、推論や行動後の出力はチェックせずに一連のワークフローを実行します。

別サイトには、こんな図も載っていました。こちらの方がイメージしやすいかもしれません。

引用：

このReWOOのメリットとしては、大きく2つあります。

LLMの呼び出しを最小限にできる
タスク完了までの遅延を削減できる

ReActでは、各ステップごとにObservationが行われていました。
このObservationはLLMが行っているため、1ステップごとにLLMを呼び出す必要があります。
最終的には、Nステップに対して、N+1回のLLM呼び出しが必要になります。

その一方で、ReWOOではObservationを行いません。
したがって、　1ステップごとにLLMを呼び出す必要がありません。

その結果、ReWOOではLLMの呼び出しを最大2回にまで削減できます。
ここまで呼び出し回数を減らせるとなると、タスクの実行時間も大きく減らすことが出来ますね。

では、なぜそんなことが可能なのか。なぜObservationが不要なのか。

それは、タスク開始前に、タスクの実行計画を作成しているからです。その計画に則って作業を実施するだけなので、LLMを逐一呼び出す必要がなくなっています。

先程の図でも、最初に CreatePlan と記載されています。ここで計画を作成しているというわけです。

そのタスク実行計画にあたるのが、Lambda関数です。

このLambda関数を元に、Bedrock Agentの行動順序や判断プロセスをワークフロー化できます。
その分、トレードオフとして、プロンプトの指定が難しくなるようです。

何はともあれ、Lambda関数を用いることで、Bedrock Agentの機能や動作を調整し、精度・適応性・効率を向上させることができるようになります。

それが、カスタムオーケストレーション機能です。

カスタムオーケストレーションLambda関数の設定

Bedrock AgentとLambda関数間のやり取りフローとしては以下の通りです。

このLambda関数は、Bedrockのランタイムプロセスに対してモデルやアクションツールを呼び出すタイミングや方法を指示し、最終的な応答を決定することによって、エージェントが入力に応答する方法を制御します。

と、言葉でだけ説明されてもよくわかりませんよね。

ドキュメントに記載されていたサンプルコードを例に、どんなことを行っているのか確認しましょう。

sample.js

function react_chain_of_thought_orchestration(event) {
    const incomingState = event.state;

    let payloadData = '';
    let responseEvent = '';
    let responseTrace = '';
    let responseAttribution = '';

    if (incomingState == 'START') {
        // 1. Invoke model in start
        responseEvent = 'INVOKE_MODEL';
        payloadData = JSON.stringify(intermediatePayload(event));
    } else if (incomingState == 'MODEL_INVOKED') {
        const stopReason = modelInvocationStopReason(event);
        if (stopReason == "tool_use") {
            // 2.a. If invoke model predicts tool call, then we send INVOKE_TOOL event
            responseEvent = 'INVOKE_TOOL';
            payloadData = toolUsePayload(event);
        } else if (stopReason == "end_turn") {
            // 2.b. If invoke model predicts an end turn, then we send FINISH event
            responseEvent = 'FINISH';
            payloadData = getEndTurnPayload(event);
        }
    } else if (incomingState == 'TOOL_INVOKED') {
        // 3. After a tool invocation, we again ask LLM to predict what should be the next step
        responseEvent = 'INVOKE_MODEL';
        payloadData = intermediatePayload(event);
    } else {
        // Invalid incoming state
        throw new Error('Invalid state provided!');
    }
                    
    // 4. Create the final payload to send back to BedrockAgent
    const payload = createPayload(payloadData, responseEvent, responseTrace, ...);
    return JSON.stringify(payload);
}

1.ステート管理

この関数内では、3つの主要な状態(state)を管理しています。

START：会話の開始状態
MODEL_INVOKED：LLMが呼び出された後の状態
→恐らくですが、ここで言う「LLM」はエージェントを指しているのだと思います
→以降、LLMはエージェントのことだとして進めます
TOOL_INVOKED：ツールが呼び出された後の状態
→ここで言うツールとは、Knowledge baseやアクショングループを指します

2.イベント実行

この関数内では、3つのイベントを実行しています。

INVOKE_MODEL：LLMを呼び出す
INVOKE_TOOL：特定のツールを使用する
FINISH：会話を終了する

3.処理の流れ

やっていることとしては、ステートごとにif文で処理を指定しています。これがワークフローとなるわけですね。

(私はあまりキャッチアップできていないのですが、AWS Heroのみのるんさんは「LangGraph風のフローを書く」と仰っていたので、そちらの知見を深めることでも理解が深まりそうです。)

ここからは具体的に、先程のサンプルコードを上から順番に見ていきます。

a)開始時のステート(`START`)

sample.js

if (incomingState == 'START') {
    // 1. Invoke model in start
    responseEvent = 'INVOKE_MODEL';
    payloadData = JSON.stringify(intermediatePayload(event));
}

新しい会話が始まると、最初はSTARTステートです。
この時は、INVOKE_MODELイベントでLLMを呼び出します。
そしてユーザーの入力に基づいて、次のアクションを決定するように依頼します。

b)LLM呼び出し後のステート(`MODEL_INVOKED`)

sample.js

else if (incomingState == 'MODEL_INVOKED') {
    const stopReason = modelInvocationStopReason(event);
    if (stopReason == "tool_use") {
        // 2.a. If invoke model predicts tool call, then we send INVOKE_TOOL event
        responseEvent = 'INVOKE_TOOL';
        payloadData = toolUsePayload(event);
    } else if (stopReason == "end_turn") {
        // 2.b. If invoke model predicts an end turn, then we send FINISH event
        responseEvent = 'FINISH';
        payloadData = getEndTurnPayload(event);
    }
}

a)の処理において、INVOKE_MODELイベントが実行されたので、ステートはMODEL_INVOKEDに変化しました。その時の処理が書かれています。

LLMの応答に基づいて、2つのパスがあります。

ツールの活用が必要な場合、INVOKE_TOOLイベントを実行します
会話を終了する場合、FINISHイベントを実行します

1.だった場合は、ステートがTOOL_INVOKEDに変化し、それに応じた次のアクションへと移ります。

2.だった場合は、最終回等がユーザーに返却されて会話終了です。

c)ツール呼び出し後のステート(`TOOL_INVOKED`)

sample.js

else if (incomingState == 'TOOL_INVOKED') {
    // 3. After a tool invocation, we again ask LLM to predict what should be the next step
    responseEvent = 'INVOKE_MODEL';
    payloadData = intermediatePayload(event);
}

b)の処理において、INVOKE_TOOLイベントが実行された場合には、TOOL_INVOKEDステートへ変化します。その時に実行する処理が記載されています。

ツールの実行結果を受け取った後、再びINVOKE_MODELイベントを実行してLLMを呼び出します。
この時、ステートはMODEL_INVOKEDに変化します。

したがって、ツールの実行結果を受け取ったら、b)に戻るということですね。
ただ順次実行しているわけではなく、ワークフローが定義されているのだということがここでわかると思います。

具体例

簡単な例をご紹介します。
ユーザーから「明日の東京の天気を教えて」という入力があったとしましょう。

その際の処理フローとしては、

STARTステート → INVOKE_MODELイベントを実行
(LLMが、天気情報が必要と判断する)
INVOKED_MODELステート → INVOKE_TOOLイベントを実行
(天気APIツールを呼び出す)
TOOL_INVOKEDステート → INVOKE_MODELイベントを実行
(天気情報を受け取り、応答を生成する)
INVOKED_MODELステート → FINISHイベントを実行
(ユーザーに回答を返す)

という形になります。
ワークフローに則って処理されているのがわかると思います。

まとめ

このようにフローを作成することで、各ステップで次のアクションを動的に判断しながらLLMの挙動を制御できる。それがカスタムオーケストレーションです。

上記は簡単な例でしたが、使い込むためにはかなり大変なLambda関数を作成しないといけなそうな印象を持ちました。

ただ、それでもこの機能がどんなものなのか・ReWOOとはなにか、を知っておくだけでもだいぶ違うのではないかな、と思います！

次にやってみること

以下のBedrockサンプルリポジトリにカスタムオーケストレーションを試せるものが追加されていたので、これを試してみます！

また、LangGraphもインプットして、どのように違うのかについても纏めてみます！

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up

【鍵はLambda】Amazon Bedrock Agentsのカスタムオーケストレーション機能を理解する

はじめに

Amazon Bedrock Agentsのアップデート

カスタムオーケストレーションとは？

従来のBedrock Agent

カスタムオーケストレーション機能

カスタムオーケストレーションLambda関数の設定

1.ステート管理

2.イベント実行

3.処理の流れ

a)開始時のステート(START)

b)LLM呼び出し後のステート(MODEL_INVOKED)

c)ツール呼び出し後のステート(TOOL_INVOKED)

具体例

まとめ

次にやってみること

a)開始時のステート(`START`)

b)LLM呼び出し後のステート(`MODEL_INVOKED`)

c)ツール呼び出し後のステート(`TOOL_INVOKED`)