AIトークンストリーミング:SSE対WebSocketだけの問題ではない
本記事では、本番環境でのトークンストリーミングの真の課題はSSEとWebSocketの選択ではなく、再接続、復旧、キャンセル、トークン圧縮、マルチユーザーサポートなどの実装にあると論じています。プロンプトリクエストとレスポンスストリームを分離するアーキテクチャを提案し、デモと本番のギャップを強調しています。
AIアプリケーションにおけるトークンストリーミングは、モデルが生成したトークンをリアルタイムでクライアントにプッシュする重要な技術です。多くの開発者はまずSSE(サーバー送信イベント)とWebSocketのプロトコル比較に目を向けますが、本記事の著者は真の困難はそこにはないと指摘します。
記事は本番環境に適したアーキテクチャを提案します。ユーザーのプロンプトリクエストとレスポンスストリームを分離し、クライアントはまずPOSTリクエストでプロンプトを送信し、サーバーは即座にストリームIDを返します。その後、クライアントは別のGETリクエストでそのストリームに接続してトークンを受信します。すべてのトークンは生成されると同時にトークンキャッシュデータベースに保存され、接続が途切れても最後の位置から再開できます。この設計は水平スケーラブルなステートレスサーバーに自然に適合し、どのサーバーでもリクエストとレスポンスを処理できます。
WebSocketと比較すると、SSEはこのアーキテクチャにおいてよりシンプルです。WebSocketはクライアントが長期間のコネクションを維持し、独自のメッセージ形式を定義する必要がありますが、SSEはHTTP標準に準拠し、Last-Event-IDによる再接続機能を内蔵しています。ただし、どちらのプロトコルを選んでも、以下の本番環境向け機能を実装しなければなりません。再接続と復旧、接続断の検出、ユーザーによるキャンセル要求の処理、トークン圧縮(保存されたトークンを完全なレスポンスに統合)、そしてマルチユーザーまたはマルチデバイスの同期です。
キャンセル機能は特に厄介です。ユーザーが停止ボタンを押すと、サーバーはトークンの生成を中断する必要があり、トークンキャッシュは読み書きの両方をチェックできる必要があります。トークン圧縮は再接続時とレスポンス完了時の両方で必要です。再接続時に大量の欠落トークンは一つずつ送信するのではなく一括で送信すべきであり、レスポンス完了後は履歴クエリのために統合保存する必要があります。
最も複雑なのはマルチユーザーシナリオです。複数のユーザーが同じチャットセッションを共有している場合、一人がプロンプトを送信したときに他のユーザーが自動的にレスポンスを受け取るにはどうすればよいでしょうか?一つの解決策は、チャットセッション全体を一つのレスポンスストリームとみなし、すべてのプロンプトとレスポンスをそのストリーム上でやり取りすることです。
著者は、簡単なデモやブログ記事と実際の本番デプロイメントの間には大きなギャップがあると強調します。多くの開発者はこのギャップに気づいていないか、埋めるための作業量を過小評価しています。解決策として、著者が所属するAbly社は、上記のすべての問題を自動的に処理するパブリッシュ/サブスクライブシステムに基づいた製品を提供しており、わずか数行のコードで統合できると述べています。