Google は AI 技術を使用して、コンテンツをご希望の言語に翻訳しています。AI 翻訳には誤りが含まれる場合があります。

RMI ADK エージェントの高度なツールインターセプト

概要

大規模なデータセットを扱う場合、モデルに生データを直接渡すと、レスポンスが遅くなり、トークン費用が増加し、推論エラーが発生します。RMI データは数百万行に及ぶ可能性があるため、エージェントが無制限の量をクエリして分析のためにすべて取り込むと、トークン費用が高くなり、コンテキストウィンドウが肥大化し、後続のターンごとにエージェントのパフォーマンスが低下します。

たとえば、RMI historical_travel_time テーブルに対するクエリは、数千の道路セグメント行を返す可能性があります。結果セット全体を会話履歴に戻すと、後続のターンごとに処理が遅くなり、費用が増加し、モデルが最終的な回答を合成できなくなります。

プロンプトの指示だけでは、これらの問題を解決できません。唯一の確実な方法は、ツール呼び出しをインターセプト して、大きなデータペイロードをモデルのコンテキストウィンドウの外にルーティングすることです。

ADK ツールのコールバック

ADK には、ツールの実行をインターセプトするためのコールバックフックが用意されています。大きな結果セットを処理するには、after_tool_callback を使用します。

before_tool_callback(tool, args, tool_context): モデルがツールを選択した後、ツールが実行される前に実行されます。
after_tool_callback(tool, args, tool_context, tool_response): ツールが完了した後、結果がモデルに戻る前に実行されます。これを使用してツールのレスポンスを変更 します。

コールバックは、セッション state にアクセスできる tool_context オブジェクトを受け取ります。これにより、大規模なデータセットをモデルのプロンプトに含めることなく、会話のターン全体でバックグラウンドメモリに保存できます。

# root_agent = llm_agent.Agent(
#     ...,
#     after_tool_callback=stash_and_truncate_results,
# )

実行後の切り捨てとスタッシュ

目的: 意思決定のためにクエリ結果の小さなサンプルをモデルに表示し、最終出力のために完全なデータセットを保存します。

RMI エージェントの場合、after_tool_callback は execute_sql の結果をインターセプトし、行の完全なリストをセッション状態に保存し、合計行数とともに短いサンプルを返します。

def stash_and_truncate_results(tool, args, tool_context, tool_response):
  if tool.name == "execute_sql" and "rows" in tool_response:
    rows = tool_response["rows"]
    tool_context.state["last_sql_result"] = rows      # keep everything
    tool_response["total_rows_fetched"] = len(rows)
    tool_response["rows"] = rows[:SAMPLE_LIMIT]        # show a sample
    tool_response["message"] = (
        f"Showing {SAMPLE_LIMIT} of {len(rows)} rows. "
        "Full result stashed in background state."
    )
  return tool_response

サンプルを返すことで、後続のターン全体で入力トークンが削減され、費用とレスポンス時間が短縮されます。また、大きなクエリ結果がコンテキストウィンドウを圧迫し、モデルが指示に従う能力を損なうこともありません。

エージェントのプロンプトで、ツールが設計上サンプルを返すことを明示的に説明してください。このコンテキストがないと、モデルが部分的なデータで混乱し、欠落しているレコードを調整しようとしてコンピューティングトークンと推論トークンを無駄にしたり、残りのレコードを取得するために冗長なツール呼び出しを開始したりする可能性があります。

例: 大規模な結果セットの処理

ユーザーが「現在、遅延率が 2 を超えるルートは？」と質問します。
エージェントの historical_travel_time に対するクエリは、数千行を返します。
after_tool_callback は、完全な結果をセッション状態に保存し、小さなサンプルと合計行数のみを返します。
モデルはそのサンプルに基づいて推論を行い、要約を作成します。
アプリケーションまたはダウンストリームツールは、セッション状態から完全なデータセットを直接読み取り、完全なテーブルをレンダリングするか、ユーザー向けにエクスポートします。

重要ポイント

大きなペイロードをプロンプトから除外する: 大きなデータセットは、プロンプト履歴を介してやり取りするのではなく、セッション状態に保存します。
大きなツール出力を切り捨てる: after_tool_callback を使用して、完全な結果をセッション状態に保存しながら、小さなサンプルをモデルに返します。
サンプルについてモデルに通知する: 結果がサンプリングされることを説明するプロンプトルールを追加して、モデルの混乱、推論トークンの無駄、冗長なツール呼び出しを防ぎます。

次のステップ

ADK コールバックのリファレンス: モデルレベル、エージェントレベル、ツールレベルのインターセプトの詳細については、公式の ADK コールバックガイドをご覧ください。

寄稿者

ナサニエル・トーマス | Google Maps Platform、ソフトウェアエンジニアリングインターン