132 lines
3.3 KiB
TypeScript
132 lines
3.3 KiB
TypeScript
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 { BlockKind, UIBlock } from './block';
|
|
import { useUserMessageId } from '@/hooks/use-user-message-id';
|
|
|
|
type StreamingDelta = {
|
|
type:
|
|
| 'text-delta'
|
|
| 'code-delta'
|
|
| 'title'
|
|
| 'id'
|
|
| 'suggestion'
|
|
| 'clear'
|
|
| 'finish'
|
|
| 'user-message-id'
|
|
| 'kind';
|
|
|
|
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 'kind':
|
|
return {
|
|
...draftBlock,
|
|
kind: delta.content as BlockKind,
|
|
};
|
|
|
|
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 'code-delta':
|
|
return {
|
|
...draftBlock,
|
|
content: delta.content as string,
|
|
isVisible:
|
|
draftBlock.status === 'streaming' &&
|
|
draftBlock.content.length > 20 &&
|
|
draftBlock.content.length < 30
|
|
? 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, setUserMessageIdFromServer]);
|
|
}
|