커스텀 데이터 스트리밍하기

모델의 응답과 함께 추가 데이터를 보내야 하는 경우가 종종 있어요. 예를 들어 처리 상태 정보, 저장 후의 메시지 id, 언어 모델이 참조하는 콘텐츠의 출처(reference) 같은 것을 함께 보내고 싶을 수 있죠. AI SDK는 이런 데이터를 클라이언트로 스트리밍해 UIMessage의 parts 배열에 붙여주는 여러 도우미를 제공합니다. 정확히 어떤 헬퍼가 있고 어떻게 쓰는지 살펴볼게요.

출처: 공식문서

본문

AI SDK가 제공하는 도우미는 세 가지예요.

createUIMessageStream — 데이터 스트림을 생성합니다.
createUIMessageStreamResponse — 데이터를 스트리밍하는 응답 객체를 만듭니다.
pipeUIMessageStreamToResponse — 데이터 스트림을 서버 응답 객체에 연결(pipe)합니다.

데이터는 Server-Sent Events를 사용해 응답 스트림의 일부로 흘러가요.

타입 안전한 데이터 스트리밍 설정

먼저 커스텀 메시지 타입을 데이터 파트 스키마와 함께 정의해 타입 안전성을 확보합니다.

import { UIMessage } from 'ai';

// 데이터 파트 스키마로 커스텀 메시지 타입 정의
export type MyUIMessage = UIMessage<
  never, // metadata type
  {
    weather: {
      city: string;
      weather?: string;
      status: 'loading' | 'success';
    };
    notification: {
      message: string;
      level: 'info' | 'warning' | 'error';
    };
  } // data parts type
>;

서버에서 데이터 스트리밍하기

서버 라우트 핸들러에서 UIMessageStream을 만들고 createUIMessageStreamResponse에 넘겨줍니다. writer.write(...)로 상태, 출처, 데이터 파트를 차례로 보낼 수 있어요.

import { openai } from '@ai-sdk/openai';
import {
  convertToModelMessages,
  createUIMessageStream,
  createUIMessageStreamResponse,
  streamText,
  toUIMessageStream,
} from 'ai';
import type { MyUIMessage } from '@/ai/types';

export async function POST(req: Request) {
  const { messages } = await req.json();

  const stream = createUIMessageStream<MyUIMessage>({
    execute: ({ writer }) => {
      // 1. 초기 상태 전송 (transient — 메시지 이력에 남지 않음)
      writer.write({
        type: 'data-notification',
        data: { message: 'Processing your request...', level: 'info' },
        transient: true, // 이 파트는 메시지 이력에 추가되지 않음
      });

      // 2. 출처 전송 (RAG 활용에 유용)
      writer.write({
        type: 'source',
        value: {
          type: 'source',
          sourceType: 'url',
          id: 'source-1',
          url: 'https://weather.com',
          title: 'Weather Data Source',
        },
      });

      // 3. 로딩 상태를 가진 데이터 파트 전송
      writer.write({
        type: 'data-weather',
        id: 'weather-1',
        data: { city: 'San Francisco', status: 'loading' },
      });

      const result = streamText({
        model: "xai/grok-4.5",
        messages: await convertToModelMessages(messages),
        onEnd() {
          // 4. 같은 데이터 파트 갱신 (reconciliation)
          writer.write({
            type: 'data-weather',
            id: 'weather-1', // 같은 ID = 기존 파트 갱신
            data: {
              city: 'San Francisco',
              weather: 'sunny',
              status: 'success',
            },
          });
        },
      });
    },
  });

  return createUIMessageStreamResponse({ stream });
}

핵심 포인트 두 가지를 짚어볼게요. transient: true를 붙인 파트는 메시지 이력에 남지 않고, 같은 id로 다시 쓰면 기존 데이터 파트를 갱신(reconcile)해요. 덕분에 로딩 상태 → 성공 상태로 자연스럽게 바뀌는 UI를 만들 수 있습니다.

더 알아보기

Server-Sent Events(SSE) — 스트리밍 데이터 전송의 기반 형식.
RAG처럼 출처를 표시해야 하면 source 타입의 파트 활용.
UIMessage의 parts 배열 구조와 스트림 프로토콜 페이지 참고.