VercelのAI SDKを用いてストリーミング可能な動的UIを実現する

この記事は every Tech Blog Advent Calendar 2024 9 日目の記事です。

はじめに

こんにちは。DELISH KITCHEN開発部の村上です。

DELISH KITCHENでは、これまでの『レシピ動画アプリ』から『AI料理アシスタント』を目指すべく、これまで以上にAI領域に力を入れています。詳しくはこちらにも記載があるので、ぜひご覧ください。

AI/LLMでtoC向けサービスはどう変わるのか？『DELISH KITCHEN』は、「レシピ動画アプリ」から「AI料理アシスタント」へ

このAI活用は社内での業務改善にも進んでおり、直近でOpenAI APIを用いた社内システムの開発をする機会がありました。その中で今回はVercelのAI SDKを使う機会があったのでAI SDKを用いたストリーミング可能なUIをwebアプリケーション内で実現する方法を紹介します。

AI SDKとは

VercelのAI SDKはAI/LLMを用いたwebアプリケーション開発を支援するためのツールです。AI/LLMを用いた開発ではOpenAI, Claudeなど外部APIへの繋ぎこみやチャットUIの実装、チャット履歴の保存、ストリーミング機能、RAGの利用などの機能が求められたりしますが、これを全て自分で開発しようと思うと、たとえwebフレームワークを使っていてもかなり手間がかかってしまいます。AI SDKを利用することでこうした実装工数を削減し、より周辺の機能開発に時間を割くことができます。

現状AI SDKは以下の3つから構成されています。

1. AI SDK Core

   • テキスト生成、構造化オブジェクトの生成、LLM（大規模言語モデル）を使用したツール呼び出しを行うためのプロバイダーに依存しない統一APIの提供

2. AI SDK UI

   • チャットやその他ユーザーインターフェースを構築するためのツールの提供

3. AI SDK RSC

   • React Server Components (RSC) を使用してユーザーインターフェースをストリーミングする機能。※現在は実験的な開発段階。

AI SDKは同じ開発元から出ているNext.jsはもちろんNuxtやSvelteなど他環境への対応もしていますが、AI SDK RSCはNext.jsのApp Routerだけをサポートしていたりするので、環境別で何が使えるかは公式を参照するのがおすすめです。

https://sdk.Vercel.ai/docs/getting-started/navigating-the-library#environment-compatibility

AI/LLMを用いたアプリケーションにおけるストリーミング機能

特にLLMを用いたwebアプリケーション開発をしていく場合、開発上ケアしておきたいポイントはいくつかありますが、その一つにストリーミング機能があります。

まず、単純な一問一答的なものを実現しようと思っても質問内容や回答によってはAPIでLLMを用いて回答を生成する段階でその出力までユーザーに待ち時間が発生してしまいます。さらに単純な1回のAPI呼び出しだけで終わればいいですが、多くの場合では複数のAPI呼び出しや処理のステップを経て、出力結果を作っていくのでその待ち時間はユーザー体験として無視できないものになってきます。

そこで身近なところであれば、ChatGPTでも全ての回答を生成し終わる前に回答の出力が段階的に行われ、ユーザーが体感する待ち時間を軽減していると思いますが、出力や処理の途中でもユーザーにフィードバックできるようなストリーミング機能が求められます。

AI SDKではAI SDK RSCの中でその機能をサポートしているので以降ではいくつかの種類に分けて機能を紹介していきます。

テキストの出力結果をストリーミングで表示する

Server

Server側ではまず createStreamableValue でServerからClientにストリーミングで送るためのデータの格納先を準備し、streamText を使ってOpenAI APIなどProviderからストリーミングされた出力結果で更新します。

'use server';

import { streamText } from 'ai';
import { openai } from '@ai-sdk/openai';
import { createStreamableValue } from 'ai/rsc';

export async function generate(input: string) {
  const stream = createStreamableValue('');

  (async () => {
    const { textStream } = streamText({
      model: openai('gpt-4o-mini'),
      prompt: input,
    });

    for await (const delta of textStream) {
      stream.update(delta);
    }

    stream.done();
  })();

  return { output: stream.value };
}

Client

Client側ではServer側が createStreamableValue で生成されたデータを readStreamableValue を用いることで簡単に読み取ることができるので受け取ったものを処理するhooksを定義します。

import { StreamableValue, readStreamableValue } from 'ai/rsc'
import { useEffect, useState } from 'react'

export const useStreamableText = (
  content: string | StreamableValue<string>
) => {
  const [rawContent, setRawContent] = useState(
    typeof content === 'string' ? content : ''
  )

  useEffect(() => {
    ;(async () => {
      if (typeof content === 'object') {
        let value = ''
        for await (const delta of readStreamableValue(content)) {
          if (typeof delta === 'string') {
            setRawContent((value = value + delta))
          }
        }
      }
    })()
  }, [content])

  return rawContent
}

表示するコンポーネント側ではServerからの結果を上記で定義した useStreambleText を使って表示するだけで簡単に実現できます。

'use client';

import { useState } from 'react';
import { generate } from '@/lib/actions';
import { useStreamableText } from '@/lib/hooks';
import { StreamableValue } from 'ai/rsc';

export default function QuestionAnswer() {
  const [answer, setAnswer] = useState<string | StreamableValue<string>>('');

  return (
    <div>
      <button
        onClick={async () => {
          const { output } = await generate('簡単に作れるお弁当レシピを教えてください。');
          setAnswer(output);
        }}
      >
        Ask
      </button>

      {answer && <AssistantMessage answer={answer} />}
    </div>
  );
}

export function AssistantMessage({
  answer,
}: {
  answer: string | StreamableValue<string>;
}) {
  const text = useStreamableText(content);

  return (
    <div>
      {text}
    </div>
  );
}

オブジェクトの出力結果をストリーミングで表示する

Server

Server側では、同じように createStreamableValue を利用するところは同じですが、ここではオブジェクト形式の出力に対応する streamObject を利用して、プロバイダーから逐次送信される構造化されたデータで更新します。AI SDKではOpenAI APIのStructured Outputsにも対応しているので、 structuredOutputs パラメーターで指定します。

'use server';

import { streamObject } from 'ai';
import { openai } from '@ai-sdk/openai';
import { createStreamableValue } from 'ai/rsc';
import { z } from 'zod';

export async function generateObject(input: string) {
  const stream = createStreamableValue({ answer: '', quotation_links: [] });

  (async () => {
    const { objectStream } = streamObject({
      model: openai('gpt-4o-mini', {
        structuredOutputs: true,
      }),
      schema: z.object({
        answer: z.string(),
        quotation_links: z.array(
          z.object({
            title: z.string(),
            link: z.string(),
          })
        ),
      }),
      prompt: input,
    });

    for await (const delta of objectStream) {
      stream.update(delta);
    }

    stream.done();
  })();

  return { output: stream.value };
}

上記の例では、objectStream には常に { answer: '', quotation_links: [] } の形式が保たれた形で随時そのテキスト情報や配列要素が追加されていくので、特に複雑な加工処理をすることなく、Client側で利用可能な状態になります。

Client

Client側では、テキストをストリーミングする時と同じように readStreamableValue を使用してストリーミングされたオブジェクトを受け取り、動的に更新します。

import { StreamableValue, readStreamableValue } from 'ai/rsc';
import { useEffect, useState } from 'react';

type AnswerObject = {
  answer: string;
  quotation_links: { title: string; link: string }[];
};

export const useStreamableObject = (
  content: AnswerObject | StreamableValue<AnswerObject>
) => {
  const [rawContent, setRawContent] = useState<AnswerObject | null>(
    typeof content === 'object' && !('subscribe' in content)
      ? content
      : { answer: '', quotation_links: [] }
  );

  useEffect(() => {
    (async () => {
      if (typeof content === 'object' && 'subscribe' in content) {
        let value: AnswerObject | null = null;
        for await (const delta of readStreamableValue(content)) {
          if (typeof delta === 'object') {
            setRawContent((value = { ...value, ...delta }));
          }
        }
      }
    })();
  }, [content]);

  return rawContent;
};

定義した useStreamableObject を使用して、ストリーミングで受け取ったオブジェクトデータを表示します。

'use client';

import { useState } from 'react';
import { generateObject } from '@/lib/actions';
import { useStreamableObject } from '@/lib/hooks';
import { StreamableValue } from 'ai/rsc';

type AnswerObject = {
  answer: string;
  quotation_links: { title: string; link: string }[];
};

export default function ObjectDisplay() {
  const [answer, setAnswer] = useState<
    AnswerObject | StreamableValue<AnswerObject> | null
  >(null);

  return (
    <div>
      <button
        onClick={async () => {
          const { output } = await generateObject('環境問題に関する最新のレポートを教えてください。');
          setAnswer(output);
        }}
      >
        Ask
      </button>

      {answer && <AssistantObjectMessage answer={answer} />}
    </div>
  );
}

export function AssistantObjectMessage({
  answer,
}: {
  answer: AnswerObject | StreamableValue<AnswerObject>;
}) {
  const data = useStreamableObject(answer);

  return (
    <div>
      {data ? (
        <div>
          <p>{data.answer}</p>
          <ul>
            {data.quotation_links.map((item, index) => (
              <li key={index}>
                <a href={item.link} target="_blank" rel="noopener noreferrer">
                  {item.title}
                </a>
              </li>
            ))}
          </ul>
        </div>
      ) : (
        'Loading...'
      )}
    </div>
  );
}

オブジェクト形式のストリーミングのメリットは上記のようにそれぞれ異なる要素に対して個別のスタイリングや処理を行うことができる点です。今までのテキストでのストリーミングは表現方法が限定的になるか、やろうと思ってもテキスト情報を変換するような複雑な処理をしないといけず不安定になってしまいますが、オブジェクト形式で受け取れることによって、アプリケーションによって独自の見せ方が可能になり、自由度が上がりました。

処理に合わせてUI自体をストリーミングで表示する

Vercelではテキストやオブジェクトだけではなく、UI自体をストリーミングすることができます。この機能を使うことでテキストやオブジェクトだけではなく、LLMの回答結果をUIで表示することも可能になります。 今回は、AIアプリケーションでありがちな回答を出力するまで進捗を表示するUIを例に紹介します。

一番最初に利用した AssistantMessage と以下の WorkflowProgress コンポーネントをServerとClientでやりとりします。

進捗を表示するWorkflowProgressコンポーネント

export function WorkflowProgress({ workflowSteps }) {
  const completedSteps = workflowSteps.filter((step) => step.status === 'completed').length;
  const totalSteps = workflowSteps.length;
  const progressValue = (completedSteps / totalSteps) * 100;

  return (
    <div className="p-4 bg-gray-100 rounded-lg shadow">
      <h3 className="font-semibold text-lg">Progress</h3>
      <progress value={progressValue} max="100" className="w-full mb-2"></progress>
      <ul>
        {workflowSteps.map((step) => (
            <li key={step.id} className="mb-2">
            <strong>{step.name}:</strong> {step.status}
            </li>
        ))}
      </ul>
    </div>
  );
}

Server

サーバー側では createStreamableUI を利用して、ストリーミングするコンポーネントを追加、更新することができます。今回はAI SDKの機能紹介がメインのため、step毎の具体的な処理については省略します。

'use server';

import { streamText } from 'ai';
import { openai } from '@ai-sdk/openai';
import { createStreamableUI } from 'ai/rsc';
import { WorkflowProgress } from '@/components/WorkflowProgress';
import { AssistantMessage } from '@/components/AssistantMessage';

export async function generateWithSteps(input: string) {
  const workflowSteps = [
    { id: 'step1', name: '質問を解析中', status: 'in-progress', tasks: [] },
    { id: 'step2', name: '探索方法を検討', status: 'pending', tasks: [] },
    { id: 'step3', name: '関連データを取得', status: 'pending', tasks: [] },
  ];

  const displayUI = createStreamableUI(
    <WorkflowProgress workflowSteps={workflowSteps} />
  );

  (async () => {
    // Step 1: 質問を解析
    await new Promise((resolve) => setTimeout(resolve, 2000));
    workflowSteps[0].status = 'completed';
    workflowSteps[1].status = 'in-progress';
    displayUI.update(
      <WorkflowProgress workflowSteps={workflowSteps} />
    );

    // Step 2: 探索方法を決定
    await new Promise((resolve) => setTimeout(resolve, 2000));
    workflowSteps[1].status = 'completed';
    workflowSteps[2].status = 'in-progress';
    displayUI.update(
      <WorkflowProgress workflowSteps={workflowSteps} />
    );

    // Step 3: データを取得
    await new Promise((resolve) => setTimeout(resolve, 2000));
    workflowSteps[2].status = 'completed';
    displayUI.update(
      <WorkflowProgress workflowSteps={workflowSteps} />
    );

    // 結果を生成
    const { textStream } = streamText({
      model: openai('gpt-4o-mini'),
      prompt: input,
    });

    let generatedText = '';
    for await (const delta of textStream) {
      generatedText += delta;

      // 進捗状況のコンポーネントから回答結果のコンポーネントに変える
      displayUI.update(
        <AssistantMessage answer={generatedText} />
      );
    }

    displayUI.done();
  })();

  return { display: displayUI.value };
}

Client

クライアント側では、定義したアクションを呼び出し、その結果をそのまま表示することで簡単に動的なUIが作れます。実際の挙動はまず最初にProgressを表示し、回答結果がLLMから出力され始めるとProgressは非表示になり、回答が生成されていくような形になります。

'use client';

import { useState } from 'react';
import { generateWithSteps } from '@/lib/actions';

export default function DynamicProgressDisplay() {
  const [display, setDisplay] = useState<React.ReactNode | null>(null);

  return (
    <div>
      <button
        onClick={async () => {
          const { display } = await generateWithSteps('環境問題に関する最新のレポートを教えてください。');
          setDisplay(display);
        }}
        className="mb-4 px-4 py-2 bg-blue-500 text-white rounded"
      >
        Ask
      </button>

      {display}
    </div>
  );
}

最後に

今回はVercelのAI SDKを使って、いくつかの手法でストリーミング可能なUIをアプリケーション上で実装する方法を紹介しました。この他にも会話履歴の保存やその状態の管理を簡単に扱えるように豊富な機能が提供されています。AI SDK RSCはまだベータ的な位置付けですが、webでLLMを使ったアプリケーションをする際には比較的簡単にリッチなアプリケーションが作れるので、ぜひ社内ツールや簡単に動くものを作りたい場合にはVercelのAI SDKを使ってみてください。

冒頭でも紹介したようにDELISH KITCHENではこれまでの『レシピ動画アプリ』から『AI料理アシスタント』への変化を起こそうとしています。ぜひ、この取り組みに興味を持った方は一度カジュアル面談でお話しましょう！

corp.every.tv