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>; }) { const { mutate } = useSWRConfig(); const [optimisticSuggestions, setOptimisticSuggestions] = useState< Array >([]); 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]); }