Vercel AI SDKを使って、SSE(ServerSentEvent)のハンドリングを簡潔に書く

株式会社エクスプラザ（公式）

こんにちは、Explazaでエンジニアをしています @_mkazutaka です。

Vercel AI SDKをみなさまご存知でしょうか。1ヶ月ぐらい前に公開されたオープンソースのライブラリで、チャットアプリケーションを構築するための便利なフックやクラスが実装されています。

チャットアプリケーションは、ChatGPTをはじめ多くの場合、SSE(Server Sent Events)を使ってストリーミングの処理が実現されています。下記のようなイメージです。このSSEの処理は、1セッションあたりで二箇所で用いられます。

[OpenAI] <- SSE -> [Web Backend] <-- SSE --> [Web Frontend]

Vercel AI SDKを使うことで、二箇所で持ちられるSSEのハンドリングの処理を簡潔に書くことができます。ドキュメントにあるコードを用いて具体的な例を見ていきます。

OpenAIとWebBackend間は、以下のコードだけで済みます。SSEの処理をうまくOpenAIStreamというクラスが隠蔽してくれています。StreamingTextResponseのインスタンスを返すことで、WebFrontendとWebBackend間の処理も簡潔に行なえます

# イメージ
import { StreamingTextResponse, OpenAIStream } from 'ai'

export async function POST(req: Request) {
  const { messages } = await req.json()
  const response = await openai.createChatCompletion({
    model: 'gpt-4',
    stream: true,
    messages
  })
  const stream = OpenAIStream(response)
  return new StreamingTextResponse(stream)
}

Webフロントエンド側の処理を見ていきます。useChatを使ってWebBackendとのSSEでの通信を行います。フック一つで済むので、実装は簡潔になります。

// イメージ

import { useChat } from 'ai/react'

export default function Chat() {
  const { messages, input, handleInputChange, handleSubmit } = useChat()

  return (
    <div className="flex flex-col w-full max-w-md py-24 mx-auto stretch">
      {messages.length > 0
        ? messages.map(m => (
            <div key={m.id} className="whitespace-pre-wrap">
              {m.role === 'user' ? 'User: ' : 'AI: '}
              {m.content}
            </div>
          ))
        : null}

      <form onSubmit={handleSubmit}>
        <input
          className="fixed bottom-0 w-full max-w-md p-2 mb-8 border border-gray-300 rounded shadow-xl"
          value={input}
          placeholder="Say something..."
          onChange={handleInputChange}
        />
      </form>
    </div>
  )
}

またVercel AI SDKは、OpenAIだけでなく自社で用意したカスタムサーバにも対応しています。本記事では、Vercel AI SDKを使ってカスタムサーバとチャットのやり取りを簡潔に書く方法を紹介します。

やってみる

(事前準備) プロジェクトの作成と起動

今回は、https://github.com/vercel-labs/ai/tree/main/examples/next-openai 上のプロジェクトを使います。これは、Vercel AI SDKを使ってOpenAIと通信する実装が書かれた例です。これの向き先をOpenAIではなく、ローカルホストににしてSSEでの通信が確認できるかを検証します。

npx create-next-app --example https://github.com/vercel-labs/ai/tree/main/examples/next-openai vercel-ai
cd vercel-ai

cp .env.local.example .env.local

# .env.localにOPENAI_API_KEYを設定する
vim .env.local

npm run dev

起動すると以下のような画面になります。チャットを送信すると返答が帰ってきます。

vercel-aiのexampleを動かした図

WebBakendのコードを変更する

接続先をOpenAIからLocalhostに変更します。ローカルのアプリケーションは、「LangChainのストリーミングレスポンスをFastAPIを介してクライアントに返す」にて構築したアプリケーションを利用します。

ローカルホストに向ける場合、Vercel AI SDKが提供するAIStreamクラスを利用します。AIStreamクラスについては下記をご参照ください。Next.JSのOpenAIと通信している部分(コード)を書き直していきます。

AIStream is a helper function for creating a readable stream for AI responses. This is based on the responses returned by fetch and serves as the basis for the OpenAIStream and AnthropicStream. It allows you to handle AI response streams in a controlled and customized manner that will work with useChat and useCompletion.

AIStreamは、AIのレスポンス用の読み取り可能なストリームを作成するためのヘルパー関数です。これは fetch によって返されるレスポンスをベースにしており、OpenAIStream と AnthropicStream のベースとなっています。useChatやuseCompletionで動作するように制御され、カスタマイズされた方法でAIレスポンスストリームを扱うことができます。

https://sdk.vercel.ai/docs/api-reference/ai-stream
import { AIStream } from 'ai'
import type { AIStreamParser, AIStreamCallbacks } from 'ai'

function parseLocalStream(): AIStreamParser {
  return data => {
    const json = JSON.parse(data) as {
      choices:[
        {
          delta: {
            content: string
          }
          finish_reason?: string,
        }
      ]
    }
    return json.choices[0].delta.content
  }
}

export function LocalStream(
  res: Response,
  cb?: AIStreamCallbacks
): ReadableStream {
  return AIStream(res, parseLocalStream(), cb)
}

以上でほぼ完成です。作ったクラスを利用するようにバックエンドのコードを変更します。最初のコードを比べると、`openai.createChatCompletion` の部分がfetchに変わって、Streamのクラスが変わったぐらいです。

export async function POST(req: Request) {
  // Extract the `prompt` from the body of the request
  const { messages } = await req.json()

  const fetchResponse = await fetch(
    'http://localhost:8000/v1/chat/completions',
    {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json'
      },
      body: JSON.stringify({
        messages: messages
      })
    }
  )

  const stream = LocalStream(fetchResponse)

  return new StreamingTextResponse(stream)
}

WebFrontendのコードを変更する

必要なしです。useChatをそのまま使うことができます。

試してみる

以上の変更を加えた上で、メッセージを送ります。API側でログが出てるので動作させることができました

messages=[ChatMessage(role='user', content='こんにちわ')]
INFO:     127.0.0.1:62833 - "POST /v1/chat/completions HTTP/1.1" 200 OK
結果


まとめ

Vercel AI SDKを使うとSSEのハンドリングを完結に書くことができます。
またローカルサーバ等に向き先を変更することもできるで、おすすめのライブラリです。

この記事が気に入ったらサポートをしてみませんか？