AIエージェントのチャット中断機能の実装 — AbortSignalからポーリング方式へ

はじめに

こんにちは、RightTouchのプロダクトエンジニアの大日向です。QANT Webエージェントというプロダクトを担当しています。

QANT Webエージェントは、一言で言うとユーザの困りごとを解決するための対話型AIエージェントです。

ユーザは自然言語入力でエージェントに質問でき、会話を通してエージェントがユーザの困りごとを特定します。そして、エージェントはあらかじめ用意されたFAQなどを元に案内します。

また、ユーザはエージェントの返答中でもチャットを中断できます。本記事では、このチャット中断機能の実装について紹介します。

webagent_demo.gif

TL;DR

  1. AbortSignalを使ってチャット中断ロジックを作成
  2. 本番環境では通信経路上のLBにより、AbortSignalが末端まで届かないと判明
  3. これによって、画面上中断したとしても、サーバ側の処理は中断しない&中断したことをDBに記録しないので、実際に起きた会話と記録された会話に差分が生じた
  4. サーバ側でも正しく中断するように、APIベースのロジックに変更

チャット処理のロジック

まず、チャットを中断せずに終わらせたときの処理の流れは以下のようになります。

  1. チャット画面でチャット送信ボタンを押すと、サーバ側の /chat エンドポイントに送信
  2. 外部LLMサービスを呼び出して回答を生成
  3. 生成された回答をDBに保存し、レスポンスをチャット画面に返す
  4. チャット履歴画面では、DBから取り出したチャット履歴をほぼそのまま表示

修正前ローカル.png

処理の流れは上のようになりますが、本番環境では負荷分散のため、チャット画面とサーバの間にロードバランサー(以下LB)を配置しています。

修正前本番.png

中断ロジックv1

チャット中断機能を実現するために、AbortSignalを使って実装しました。

フロントからAbortSignalを用いてサーバまでの通信経路を切り、サーバ側でそれを検出して必要な中断処理を行うような流れです。

v1ローカル.png

ローカル環境では問題なく動きました。しかし、本番環境にデプロイするとLBがあるため、AbortSignalでチャット画面からLBまでの通信は切れたが、LBから /chat までの通信は切れませんでした。

これによって、

  1. AbortSignalが /chat まで届かないので、AbortSignalを検出するような実装をしていてもそれを検出できない
  2. 必要な中断処理が行われないため、外部LLMサービスで生成し切った不要な回答をそのままDBに保存してしまう
  3. チャット履歴画面では、中断によって本当はユーザに表示されていない回答が、あたかも表示されていたかのように見える

という問題が発生していました。

v1本番.png

中断ロジックv2

問題を解決するために以下の検討をしました。

  1. WebSocket化
    1. WebSocketでTCP上に双方向の永続コネクションを張ることで、フロントからの切断リクエストをサーバ側で直接検知できるようになる
    2. 解決したい課題に対して実装コストが高いため見送り
  2. HTTP/2化
    1. サーバはHTTP/1.1を使っているが、これをHTTP/2に変える
    2. HTTP/2にすることで、フロント〜サーバ間がHTTP/2ストリームで張られるようになり、フロントからの切断リクエストがコンテナまで伝播するので、サーバ側で中断を直接検知できるようになる可能性がある
    3. 影響範囲が不明なうえ、効果も不確実なため見送り
  3. /chat ストリーミング化
    1. レスポンスをchunked/SSEで流し始めると、サーバ側(正確にはCloud Run)がストリーミングモードに入り、クライアント切断がコンテナに届く可能性がある
    2. 効果が不確実なため見送り
  4. 中断用のエンドポイント+中断管理テーブルを新規作成
    1. サーバ側で中断を検出する代わりに、チャットごとの中断を管理&参照できるようにする
    2. 確実に解決できるため採用

以上を踏まえて、4の方針に従ってv2では以下のような修正を入れました。

  • チャット中断状態をIDで管理するテーブル ChatAbortStatus を作成
  • チャット中断用のエンドポイント /chat/abort を作成
  • 回答生成中にチャット中断状態をポーリング処理で監視

これらを使って、本番環境でのチャット中断処理を下のようにすることで、サーバ側の処理を正しく中断させるようにしました

  1. チャット中、API経由でチャット中断リクエスト
  2. 中断管理テーブルに状態登録
  3. ポーリング処理経由で中断検知
  4. AbortSignalを送信し、回答生成を中断
  5. 中断した事実込みでテーブル保存
  6. 中断した事実が正しく履歴に反映される

v2本番.png

さらなる改善点

各図にあるDBはAlloyDBを指していて、ChatAbortStatusもAlloyDBに作成しました。

プロダクトエンジニアとして「チャットの記録差分という問題を早く解き切る」ことを優先し、既に本番環境で利用しているAlloyDBにそのままテーブルを追加しました。

しかし、ChatAbortStatus は性質上、

  • 何かの不具合が生じたときに特定のレコードが残り続けるリスクがあるため、レコードにTTLが欲しい
  • ポーリング処理により頻繁にアクセスされる

ので、AlloyDBのようなRDBよりもインメモリDBの方が適切です。

そのため今後は、ValkeyなどのインメモリDBへの切り替えを予定しています。

おわりに

本記事では、QANT Webエージェントのチャット中断処理の実装とその経緯について紹介しました。 環境差分によって、開発時には気づけなかった問題をリリース後に発見するのは理想的とは言えませんが、今回の出来事を通して、

  • 開発時はより本番環境を意識するようになった
  • プロダクトエンジニアとして課題を解く楽しさを感じた
  • AIエージェントのチャット中断の実装例として知見を公開したいと感じた

ので、自分にとってすごく貴重な体験となりました。今回の経験を活かしながら、今後もQANT Webエージェントを改善していきたいと思います。

また、この記事を読んでもっといいやり方を思いついた方や、「自分も課題を解くプロダクトエンジニアになってみたい」と思った方がいましたら、ぜひ弊社にご応募ください!

採用情報

RightTouchでは、Product Engineerをはじめ、プロダクト価値を一緒に育てる仲間を積極採用中です。カジュアル面談も歓迎しています。ご興味があれば、ぜひ採用ページをご覧ください。