feat: modularize block system (#718)

components/data-stream-handler.tsx

@@ -1,14 +1,13 @@
 'use client';
 
 import { useChat } from 'ai/react';
-import { useEffect, useRef, useState } from 'react';
-import { BlockKind } from './block';
+import { useEffect, useRef } from 'react';
+import { blockDefinitions, BlockKind } from './block';
 import { Suggestion } from '@/lib/db/schema';
 import { initialBlockData, useBlock } from '@/hooks/use-block';
 import { useUserMessageId } from '@/hooks/use-user-message-id';
-import { useSWRConfig } from 'swr';
 
-type DataStreamDelta = {
+export type DataStreamDelta = {
   type:
     | 'text-delta'
     | 'code-delta'
@@ -26,22 +25,9 @@ type DataStreamDelta = {
 export function DataStreamHandler({ id }: { id: string }) {
   const { data: dataStream } = useChat({ id });
   const { setUserMessageIdFromServer } = useUserMessageId();
-  const { setBlock } = useBlock();
+  const { block, setBlock, setMetadata } = useBlock();
   const lastProcessedIndex = useRef(-1);
 
-  const { mutate } = useSWRConfig();
-  const [optimisticSuggestions, setOptimisticSuggestions] = useState<
-    Array<Suggestion>
-  >([]);
-
-  useEffect(() => {
-    if (optimisticSuggestions && optimisticSuggestions.length > 0) {
-      const [optimisticSuggestion] = optimisticSuggestions;
-      const url = `/api/suggestions?documentId=${optimisticSuggestion.documentId}`;
-      mutate(url, optimisticSuggestions, false);
-    }
-  }, [optimisticSuggestions, mutate]);
-
   useEffect(() => {
     if (!dataStream?.length) return;
 
@@ -54,6 +40,18 @@ export function DataStreamHandler({ id }: { id: string }) {
         return;
       }
 
+      const blockDefinition = blockDefinitions.find(
+        (blockDefinition) => blockDefinition.kind === block.kind,
+      );
+
+      if (blockDefinition?.onStreamPart) {
+        blockDefinition.onStreamPart({
+          streamPart: delta,
+          setBlock,
+          setMetadata,
+        });
+      }
+
       setBlock((draftBlock) => {
         if (!draftBlock) {
           return { ...initialBlockData, status: 'streaming' };
@@ -81,50 +79,6 @@ export function DataStreamHandler({ id }: { id: string }) {
               status: 'streaming',
             };
 
-          case 'text-delta':
-            return {
-              ...draftBlock,
-              content: draftBlock.content + (delta.content as string),
-              isVisible:
-                draftBlock.status === 'streaming' &&
-                draftBlock.content.length > 400 &&
-                draftBlock.content.length < 450
-                  ? true
-                  : draftBlock.isVisible,
-              status: 'streaming',
-            };
-
-          case 'code-delta':
-            return {
-              ...draftBlock,
-              content: delta.content as string,
-              isVisible:
-                draftBlock.status === 'streaming' &&
-                draftBlock.content.length > 300 &&
-                draftBlock.content.length < 310
-                  ? true
-                  : draftBlock.isVisible,
-              status: 'streaming',
-            };
-
-          case 'image-delta':
-            return {
-              ...draftBlock,
-              content: delta.content as string,
-              isVisible: true,
-              status: 'streaming',
-            };
-
-          case 'suggestion':
-            setTimeout(() => {
-              setOptimisticSuggestions((currentSuggestions) => [
-                ...currentSuggestions,
-                delta.content as Suggestion,
-              ]);
-            }, 0);
-
-            return draftBlock;
-
           case 'clear':
             return {
               ...draftBlock,
@@ -143,7 +97,7 @@ export function DataStreamHandler({ id }: { id: string }) {
         }
       });
     });
-  }, [dataStream, setBlock, setUserMessageIdFromServer]);
+  }, [dataStream, setBlock, setUserMessageIdFromServer, setMetadata, block]);
 
   return null;
 }