Bedrockのアプリ開発におけるChat Memoryの活用

はじめに

先日、GAになったばかりのSpring AIを使ってBedrockを対象としたアプリを実装した記事を記しました。

miyohide.hatenablog.com

今回は、上の記事の考察に記したChat Memoryを実装してみました。

Chat Memoryを使わない場合

Chat Memoryを使わない場合、前回のプロンプトの結果を踏まえた内容を聞いても反映されません。以下の例では「結果に10を足して」と聞いていますが、計算が実行できない旨が回答されます。

Chat Memoryを使う

Spring AIの公式ページに説明があるのでそれを元に実装します。

docs.spring.io

具体的には、前回作ったControllerのコンストラクタを改修して以下のように実装します。

  public MyController(ChatClient.Builder builder) {
    ChatMemory chatMemory =
        MessageWindowChatMemory.builder().build();
    this.chatClient =
        builder.defaultAdvisors(MessageChatMemoryAdvisor.builder(chatMemory).build()).build();
  }

ポイントはchatClientを作る際にdefaultAdvisorsとしてMessageChatMemoryAdvisorを指定してあげることです。他にもPromptChatMemoryAdvisorVectorStoreChatMemoryAdvisorを使う方法があるのですが今回は未検証。defaultAdvisorsについては以下の説明を参照してください。

docs.spring.io

実装は以上で終了です。他にはデータベースに保存したり、Apache Cassandraに保存したりする方法もあるのですが、今回はすぐさま確認できるメモリ内保存の方法である方法を使います。

これでプロンプトを投げると前回の結果を踏まえた演算を実施してくれます。

Bedrockのログで確認する

挙動はChat Memoryの使用の有無で差が発生することは確認できるのですが、Bedrockのログでどのような挙動が行われているかも確認したいと思います。Bedrockのログは明示的に設定する必要があります。「モデル呼び出しのログ記録」をオンにして、S3など保存先を指定すれば良いだけです。

今回、S3に保存するようにすると結果としてJSONファイルが出力されます。詳細は以下の公式ドキュメントを参照すると良いかと。

docs.aws.amazon.com

今回ログを取ると、Chat Memoryを設定していない場合は以下のようになっていました。

{省略"input":{"inputContentType":"application/json","inputBodyJson":{"messages":[{"role":"user","content":[{"text":"2+4は"}]}],"system":[],"inferenceConfig":{"maxTokens":300,"temperature":0.7},"additionalModelRequestFields":{}},"inputTokenCount":5},"output":{"outputContentType":"application/json","outputBodyJson":{"output":{"message":{"role":"assistant","content":[{"text":"2と4を足すと、その結果は6になります。\n\n\\[ 2 + 4 = 6 \\]\n\nなので、2+4は6です。"}]}},"stopReason":"end_turn","metrics":{"latencyMs":289},"usage":{"inputTokens":5,"outputTokens":36,"totalTokens":41}},"outputTokenCount":36}}
{省略"input":{"inputContentType":"application/json","inputBodyJson":{"messages":[{"role":"user","content":[{"text":"結果に10を足して"}]}],"system":[],"inferenceConfig":{"maxTokens":300,"temperature":0.7},"additionalModelRequestFields":{}},"inputTokenCount":8},"output":{"outputContentType":"application/json","outputBodyJson":{"output":{"message":{"role":"assistant","content":[{"text":"「結果に10を足して」という指示は、具体的な数値や式が提供されていないため、明確な数値を導くことはできません。しかし、一般的な方法を説明します。\n\n例えば、ある数値の結果が「X」である場合、その結果に10を足すには、以下のように計算します:\n\n\\[ \\text{新しい結果} = X + 10 \\]\n\n例えば、もし「X」が5である場合:\n\n\\[ \\text{新しい結果} = 5 + 10 = 15 \\]\n\n具体的な数値や式があれば、それを代入して計算することができます。具体的な例を教えてください、もしあれば、それに基づいて計算します。"}]}},"stopReason":"end_turn","metrics":{"latencyMs":976},"usage":{"inputTokens":8,"outputTokens":161,"totalTokens":169}},"outputTokenCount":161}}

これだけだと分かりにくいのですが、Chat Memoryを設定している場合と比較すると違いが明らかです。

Chat Memoryを設定している場合は以下のようになっていました。

{省略"input":{"inputContentType":"application/json","inputBodyJson":{"messages":[{"role":"user","content":[{"text":"2+4は"}]}],"system":[],"inferenceConfig":{"maxTokens":300,"temperature":0.7},"additionalModelRequestFields":{}},"inputTokenCount":5},"output":{"outputContentType":"application/json","outputBodyJson":{"output":{"message":{"role":"assistant","content":[{"text":"2と4を足すと、その結果は6になります。計算式は以下の通りです:\n\n\\[ 2 + 4 = 6 \\]\n\nしたがって、2+4は6です。"}]}},"stopReason":"end_turn","metrics":{"latencyMs":326},"usage":{"inputTokens":5,"outputTokens":44,"totalTokens":49}},"outputTokenCount":44}}
{省略"input":{"inputContentType":"application/json","inputBodyJson":{"messages":[{"role":"user","content":[{"text":"2+4は"}]},{"role":"assistant","content":[{"text":"2と4を足すと、その結果は6になります。計算式は以下の通りです:\n\n\\[ 2 + 4 = 6 \\]\n\nしたがって、2+4は6です。"}]},{"role":"user","content":[{"text":"結果に10を足して"}]}],"system":[],"inferenceConfig":{"maxTokens":300,"temperature":0.7},"additionalModelRequestFields":{}},"inputTokenCount":62},"output":{"outputContentType":"application/json","outputBodyJson":{"output":{"message":{"role":"assistant","content":[{"text":"2+4の結果である6に10を足すと、次のようになります:\n\n\\[ 6 + 10 = 16 \\]\n\nしたがって、結果に10を足すと16になります。"}]}},"stopReason":"end_turn","metrics":{"latencyMs":304},"usage":{"inputTokens":62,"outputTokens":48,"totalTokens":110}},"outputTokenCount":48}}

Chat Memoryを設定していた場合、Bedrockに対して送信しているプロンプトに前回の内容(InputとOutputの両方)が含まれていることが確認できます。このため、計算結果に基づいた処理が可能となっています。

考察

Spring AIのChat Memoryを使った実装を行い、チャットの動作確認を行いました。今回、Chat Memoryの有無によって挙動の違いが分かりやすく出たのですが、詳細に確認するためにBedrockのログを活用しました。SDKを使っているとこのような機能を実装するのは思いのほか大変なので、フレームワーク(Spring AI)が提供している機能は使うとよいかなと思います。 一方で、入力と出力の結果をプロンプトに都度含まれることになるので、トークン数が大きく増える要因にもなり結果として料金も多くかかります。この点は注意すべきかと思います。