chatbot-template/components/use-block-stream.tsx

import type { JSONValue } from 'ai';
import { type Dispatch, type SetStateAction, useEffect, useState } from 'react';
import { useSWRConfig } from 'swr';

import type { Suggestion } from '@/lib/db/schema';

import type { UIBlock } from './block';
import { useUserMessageId } from '@/hooks/use-user-message-id';

type StreamingDelta = {
  type:
    | 'text-delta'
    | 'title'
    | 'id'
    | 'suggestion'
    | 'clear'
    | 'finish'
    | 'user-message-id';

  content: string | Suggestion;
};

export function useBlockStream({
  streamingData,
  setBlock,
}: {
  streamingData: JSONValue[] | undefined;
  setBlock: Dispatch<SetStateAction<UIBlock>>;
}) {
  const { mutate } = useSWRConfig();
  const [optimisticSuggestions, setOptimisticSuggestions] = useState<
    Array<Suggestion>
  >([]);

  const { setUserMessageIdFromServer } = useUserMessageId();

  useEffect(() => {
    if (optimisticSuggestions && optimisticSuggestions.length > 0) {
      const [optimisticSuggestion] = optimisticSuggestions;
      const url = `/api/suggestions?documentId=${optimisticSuggestion.documentId}`;
      mutate(url, optimisticSuggestions, false);
    }
  }, [optimisticSuggestions, mutate]);

  useEffect(() => {
    const mostRecentDelta = streamingData?.at(-1);
    if (!mostRecentDelta) return;

    const delta = mostRecentDelta as StreamingDelta;

    if (delta.type === 'user-message-id') {
      setUserMessageIdFromServer(delta.content as string);
      return;
    }

    setBlock((draftBlock) => {
      switch (delta.type) {
        case 'id':
          return {
            ...draftBlock,
            documentId: delta.content as string,
          };

        case 'title':
          return {
            ...draftBlock,
            title: delta.content as string,
          };

        case 'text-delta':
          return {
            ...draftBlock,
            content: draftBlock.content + (delta.content as string),
            isVisible:
              draftBlock.status === 'streaming' &&
              draftBlock.content.length > 200 &&
              draftBlock.content.length < 250
                ? true
                : draftBlock.isVisible,
            status: 'streaming',
          };

        case 'suggestion':
          setTimeout(() => {
            setOptimisticSuggestions((currentSuggestions) => [
              ...currentSuggestions,
              delta.content as Suggestion,
            ]);
          }, 0);

          return draftBlock;

        case 'clear':
          return {
            ...draftBlock,
            content: '',
            status: 'streaming',
          };

        case 'finish':
          return {
            ...draftBlock,
            status: 'idle',
          };

        default:
          return draftBlock;
      }
    });
  }, [streamingData, setBlock]);
}