大規模言語モデル(LLM)使ってアプリケーションを開発する際、開発者が直面する重要な課題の1つは、LLMの出力を構造化されたオブジェクトとしてストリーミングすることです。このブログ記事では、VercelのAI SDKを使って、LLMの出力を構造化しながらストリーミングする方法をサンプルコードと共に紹介します。
LLMの出力の課題
LLMを使用する際、生成結果をユーザーに早く見せるために出力をストリーミングすることがよくあります。しかし、ストリーミングデータは構造化されていないので以下のような課題に遭遇します。
「日本の四季それぞれについて短い説明をJSONで構造化して出力してください。」とLLMに入力し、出力をストリーミングすると、まず、以下のようなデータが得られます:
{ "日本の四季": [
次に得られるのは以下のようなデータです。
{ "季節": "春", "期間": "3月〜5月",
ストリーミングされた全ての文字列を結合すると以下のように、構造化されたJSONデータが得られます。
{ "日本の四季": [ { "季節": "春", "期間": "3月〜5月", "特徴": "桜の開花、新学期の始まり", "気候": "穏やかで過ごしやすい", "代表的な行事": "花見、入学式" }, { "季節": "夏", "期間": "6月〜8月", "特徴": "高温多湿、梅雨", "気候": "蒸し暑く、時折激しい雨", "代表的な行事": "花火大会、夏祭り" }, { "季節": "秋", "期間": "9月〜11月", "特徴": "紅葉、収穫の季節", "気候": "涼しく、カラッとした快適な気候", "代表的な行事": "月見、紅葉狩り" }, { "季節": "冬", "期間": "12月〜2月", "特徴": "雪景色、こたつの季節", "気候": "寒冷で乾燥、北日本では大雪", "代表的な行事": "初詣、節分" } ] }
ただ、ストリーム中の値は結合しても、JSONとしてパースできる形式ではないので、例えば、「春の情報が全て出力されたら、春だけは画面上に表示する」ようなUIを実装することはこのままでは難しいです。
{ "日本の四季": [ { "季節": "春", "期間": "3月〜5月", "特徴": "桜の開花、新学期の始まり", "気候": "穏やかで過ごしやすい", // <-- カンマで終わっている // } がない // ] がない // } がない
この問題に対処するにはいくつかの方法がありますが、今回はVercelのAI SDKを使って対応してみます。
Vercel AI SDKの紹介
Vercel AI SDKには色々な機能があるのですが、今回はstreamObject1というAPIを使います。
streamObjectは、LLMを用いて、与えられたプロンプトとスキーマに対して型付けされた構造化オブジェクトをストリームできるAPIで、以下のように使います。
import { openai } from '@ai-sdk/openai'; import { streamObject } from 'ai'; import { z } from 'zod'; const { partialObjectStream } = await streamObject({ model: openai('gpt-4-turbo'), schema: z.object({ recipe: z.object({ name: z.string(), ingredients: z.array(z.string()), steps: z.array(z.string()), }), }), prompt: 'Generate a lasagna recipe.', }); for await (const partialObject of partialObjectStream) { console.clear(); console.log(partialObject); }
このコードをZed REPLで実行すると以下のような出力が得られます。
Zed REPLについては、以下の記事で紹介しています。
動画だとちょっとわかりにくいですが、ストリーミングされているデータは、常にJSONとしてパースできる形式になっています。以下、切り抜きを2つ貼ります。
これは嬉しいですね。そして、これをNext.jsのServer Actionとして実行し、クライアントで構造化された値として取得することができます。
// app/action.ts "use server"; import { openai } from "@ai-sdk/openai"; import { jsonSchema, streamObject } from "ai"; import { createStreamableValue } from "ai/rsc"; import { schema } from "./schema"; export async function generate(input: string) { const stream = createStreamableValue(); (async () => { const { partialObjectStream } = await streamObject({ model: openai("gpt-4o-mini"), system: "You generate an answer to a question. ", prompt: input, schema, }); for await (const partialObject of partialObjectStream) { stream.update(partialObject); } stream.done(); })(); return { object: stream.value }; }
// app/page.tsx "use client"; import { readStreamableValue } from "ai/rsc"; import { useCallback, useState } from "react"; import { generate } from "./actions"; import { Result } from "./components/result"; import { Welcome } from "./components/welcome"; import type { PartialGeneratedObject } from "./schema"; // Allow streaming responses up to 30 seconds export const maxDuration = 30; export default function Home() { const [query, setQuery] = useState(""); const [completed, setCompleted] = useState(false); const [generatedObject, setGeneratedObject] = useState<PartialGeneratedObject>({}); const handleRequest = useCallback(async (query: string) => { setQuery(query); const { object } = await generate(query); for await (const partialObject of readStreamableValue(object)) { if (partialObject) { setGeneratedObject(partialObject as PartialGeneratedObject); } } setCompleted(true); }, []); if (query !== "") { return ( <Result query={query} generatedObject={generatedObject} completed={completed} onReset={() => { setQuery(""); setCompleted(false); setGeneratedObject({}); }} /> ); } return <Welcome onRequest={handleRequest} />; }
LLMの出力をJSONに変換している仕組み
AI SDKのstreamObjectは、LLMの出力をJSONとしてパースできる形式に変換してからストリーミングすることで、開発者がデータを簡単に処理できるようにしています。
この変換の仕組みは、fix-json.ts
というファイルで実装されています。
テストケースを見るとわかりやすのですが、JSONとして不完全な文字列をJSONとしてパースできる文字列に変換しています。
test('should handle keys without values', () => { assert.strictEqual(fixJson('{"key":'), '{}'); }); test('should handle closing brace after number in object', () => { assert.strictEqual( fixJson('{"a": {"b": 1}, "c": {"d": 2'), '{"a": {"b": 1}, "c": {"d": 2}}', ); }); test('should handle closing brace after string in object', () => { assert.strictEqual( fixJson('{"a": {"b": "1"}, "c": {"d": 2'), '{"a": {"b": "1"}, "c": {"d": 2}}', ); });
streamObjectを使った実装例
最後に、streamObjectを使った実装例を紹介します。
https://github.com/toyamarinyon/artifact-sample
streamObjectを使って、ClaudeのArtifactを模したUIを実装しました。実行後に会話を続けることはできませんが、LLMの出力を構造化しつつストリーミングすることで、UIをリアルタイムに更新しながらLLMの出力を表示することができました。
以上、VercleのAI SDKを使ったLLMの出力を構造化しつつストリーミングする方法の紹介でした。streamObjectを使うことで、LLMの出力を簡単にストリーミングできるようになります。ぜひお試しください。