ADKでstateを使った会話の状態管理

はじめに

ADKでは会話をセッション単位で管理しており、セッション内で状態を保持したい場合はstateという機能を使用できます。今回はstateについて試してみたいと思います。

stateとは

stateとは、セッション内で状態を保持するための仕組みです。キーと値のペア（辞書形式）で情報を管理し、エージェントが会話の文脈や進行状況を記憶できるようにします。
エージェントが会話を効果的に進めるために必要な動的な情報を保存・更新するために利用されます。

stateの使用方法

登録したstateは{key}の形式で埋め込むことが可能です。
以下は、animalというキーで登録したstateを利用する例です。

stateの使用

root_agent = Agent(
    model=os.environ.get("MODEL"),
    name='memory_agent',
    instruction="挨拶を返してください。{animal}のような口調で会話してください。",

埋め込まれたキーはstateの値に変換されます。

挨拶を返してください。犬のような口調で会話してください。

stateの登録方法

公式ドキュメントによると、stateの登録方法は3種類あるそうです。

• output_keyを利用する方法
• EventActions.state_deltaを利用する方法
• CallbackContextやToolContextを利用する方法

今回は3つ目の「CallbackContextやToolContextを利用する方法」を試してみようと思います。（それ以外は試せてないので、またの機会に試してみようと思います）

CallbackContextを利用してstateを更新する方法

Agentに渡す引数にはbefore_agent_callbackなどのフックが用意されており、任意のタイミングで処理を割り込ませることができます。その処理にはCallbackContext型の引数を渡すことができ、その中のstateにキーを値を指定することで登録/更新することができます。

stateの更新

def update_session(callback_context: CallbackContext):
    # セッション更新
    callback_context.state["animal"] = "犬"

root_agent = Agent(
    model=os.environ.get("MODEL"),
    name='memory_agent',
    before_agent_callback=update_session,
    instruction="挨拶を返してください。{animal}のような口調で会話してください。",
)

試してみる

ということで試してみましょう。
例で使用したコードを少し修正し、「ユーザとのやり取りの回数に応じて、犬や猫のような口調になるエージェント」を作成してみました。

犬や猫のような口調になるエージェント

import os

from google.adk.agents.llm_agent import Agent
from google.adk.agents.callback_context import CallbackContext
from google.adk.tools.load_memory_tool import load_memory_tool
from google.adk.tools.preload_memory_tool import preload_memory_tool

from dotenv import load_dotenv

load_dotenv()

def update_session(callback_context: CallbackContext):
    count = callback_context.state.get("user_action_count", 0) + 1
    # セッション更新
    callback_context.state["user_action_count"] = count
    callback_context.state["animal"] = "犬" if count % 2 == 1 else "猫"

root_agent = Agent(
    model=os.environ.get("MODEL"),
    name='memory_agent',
    before_agent_callback=update_session,
    instruction="挨拶を返してください。{animal}のような口調で会話してください。",
    tools=[
        load_memory_tool,
        preload_memory_tool,
    ],
)

adk webで実行してみました。

Agentのtoolsにload_memory_toolとpreload_memory_toolを利用しています。詳しい役割は分かっていませんが、参考にしたサンプルコードにもついていたので、必須なのかもしれないということでつけておきました。（なくしてもstateの参照はできました）。

ちなみに、web ui 上にはstateというタブがあり、そこには現在のstateの一覧が表示されています。ここを見てもstateが更新されていることを確認できました。

おわりに

今回はADKのstateの利用を試してみました。Vertex AI Agent Engine上で試した際、前の会話の内容を参照していないように思えたので、stateを使って改善できないかは試してみようと思います。
ここまでご覧いただきありがとうございました！