VercelのAI SDKを使ってLLMの出力を構造化しつつストリーミングする

大規模言語モデル(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でstreamObjectを実行している様子

Zed REPLについては、以下の記事で紹介しています。

tech.route06.co.jp

動画だとちょっとわかりにくいですが、ストリーミングされているデータは、常にJSONとしてパースできる形式になっています。以下、切り抜きを2つ貼ります。

streamObjectの実行途中1

streamObjectの実行途中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というファイルで実装されています。

https://github.com/vercel/ai/blob/d3aa5486529e3d1a38b30e3972b4f4c63ea4ae9a/packages/ui-utils/src/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 sample app demo

streamObjectを使って、ClaudeのArtifactを模したUIを実装しました。実行後に会話を続けることはできませんが、LLMの出力を構造化しつつストリーミングすることで、UIをリアルタイムに更新しながらLLMの出力を表示することができました。


以上、VercleのAI SDKを使ったLLMの出力を構造化しつつストリーミングする方法の紹介でした。streamObjectを使うことで、LLMの出力を簡単にストリーミングできるようになります。ぜひお試しください。