refactor: improve block modularity on server (#758)

app/(chat)/api/chat/route.ts

@@ -120,7 +120,7 @@ export async function POST(request: Request) {
         sendReasoning: true,
       });
     },
-    onError: (error) => {
+    onError: () => {
       return 'Oops, an error occured!';
     },
   });

blocks/code/server.ts

@@ -0,0 +1,73 @@
+import { z } from 'zod';
+import { streamObject } from 'ai';
+import { myProvider } from '@/lib/ai/models';
+import { codePrompt, updateDocumentPrompt } from '@/lib/ai/prompts';
+import { createDocumentHandler } from '@/lib/blocks/server';
+
+export const codeDocumentHandler = createDocumentHandler<'code'>({
+  kind: 'code',
+  onCreateDocument: async ({ title, dataStream }) => {
+    let draftContent = '';
+
+    const { fullStream } = streamObject({
+      model: myProvider.languageModel('block-model'),
+      system: codePrompt,
+      prompt: title,
+      schema: z.object({
+        code: z.string(),
+      }),
+    });
+
+    for await (const delta of fullStream) {
+      const { type } = delta;
+
+      if (type === 'object') {
+        const { object } = delta;
+        const { code } = object;
+
+        if (code) {
+          dataStream.writeData({
+            type: 'code-delta',
+            content: code ?? '',
+          });
+
+          draftContent = code;
+        }
+      }
+    }
+
+    return draftContent;
+  },
+  onUpdateDocument: async ({ document, description, dataStream }) => {
+    let draftContent = '';
+
+    const { fullStream } = streamObject({
+      model: myProvider.languageModel('block-model'),
+      system: updateDocumentPrompt(document.content, 'code'),
+      prompt: description,
+      schema: z.object({
+        code: z.string(),
+      }),
+    });
+
+    for await (const delta of fullStream) {
+      const { type } = delta;
+
+      if (type === 'object') {
+        const { object } = delta;
+        const { code } = object;
+
+        if (code) {
+          dataStream.writeData({
+            type: 'code-delta',
+            content: code ?? '',
+          });
+
+          draftContent = code;
+        }
+      }
+    }
+
+    return draftContent;
+  },
+});

blocks/image/server.ts

@@ -0,0 +1,54 @@
+import { myProvider } from '@/lib/ai/models';
+import { createDocumentHandler } from '@/lib/blocks/server';
+import { saveDocument } from '@/lib/db/queries';
+import { experimental_generateImage } from 'ai';
+
+export const imageDocumentHandler = createDocumentHandler<'image'>({
+  kind: 'image',
+  onCreateDocument: async ({ id, title, dataStream, session }) => {
+    let draftContent = '';
+
+    const { image } = await experimental_generateImage({
+      model: myProvider.imageModel('small-model'),
+      prompt: title,
+      n: 1,
+    });
+
+    draftContent = image.base64;
+
+    dataStream.writeData({
+      type: 'image-delta',
+      content: image.base64,
+    });
+
+    if (session?.user?.id) {
+      await saveDocument({
+        id,
+        title,
+        kind: 'image',
+        content: draftContent,
+        userId: session.user.id,
+      });
+    }
+
+    return draftContent;
+  },
+  onUpdateDocument: async ({ description, dataStream }) => {
+    let draftContent = '';
+
+    const { image } = await experimental_generateImage({
+      model: myProvider.imageModel('image-model'),
+      prompt: description,
+      n: 1,
+    });
+
+    draftContent = image.base64;
+
+    dataStream.writeData({
+      type: 'image-delta',
+      content: image.base64,
+    });
+
+    return draftContent;
+  },
+});

blocks/sheet/server.ts

@@ -0,0 +1,78 @@
+import { myProvider } from '@/lib/ai/models';
+import { sheetPrompt, updateDocumentPrompt } from '@/lib/ai/prompts';
+import { createDocumentHandler } from '@/lib/blocks/server';
+import { streamObject } from 'ai';
+import { z } from 'zod';
+
+export const sheetDocumentHandler = createDocumentHandler<'sheet'>({
+  kind: 'sheet',
+  onCreateDocument: async ({ title, dataStream }) => {
+    let draftContent = '';
+
+    const { fullStream } = streamObject({
+      model: myProvider.languageModel('block-model'),
+      system: sheetPrompt,
+      prompt: title,
+      schema: z.object({
+        csv: z.string().describe('CSV data'),
+      }),
+    });
+
+    for await (const delta of fullStream) {
+      const { type } = delta;
+
+      if (type === 'object') {
+        const { object } = delta;
+        const { csv } = object;
+
+        if (csv) {
+          dataStream.writeData({
+            type: 'sheet-delta',
+            content: csv,
+          });
+
+          draftContent = csv;
+        }
+      }
+    }
+
+    dataStream.writeData({
+      type: 'sheet-delta',
+      content: draftContent,
+    });
+
+    return draftContent;
+  },
+  onUpdateDocument: async ({ document, description, dataStream }) => {
+    let draftContent = '';
+
+    const { fullStream } = streamObject({
+      model: myProvider.languageModel('block-model'),
+      system: updateDocumentPrompt(document.content, 'sheet'),
+      prompt: description,
+      schema: z.object({
+        csv: z.string(),
+      }),
+    });
+
+    for await (const delta of fullStream) {
+      const { type } = delta;
+
+      if (type === 'object') {
+        const { object } = delta;
+        const { csv } = object;
+
+        if (csv) {
+          dataStream.writeData({
+            type: 'sheet-delta',
+            content: csv,
+          });
+
+          draftContent = csv;
+        }
+      }
+    }
+
+    return draftContent;
+  },
+});

blocks/text.tsx → blocks/text/client.tsx

@@ -12,7 +12,7 @@ import {
 } from '@/components/icons';
 import { Suggestion } from '@/lib/db/schema';
 import { toast } from 'sonner';
-import { getSuggestions } from './actions';
+import { getSuggestions } from '../actions';
 
 interface TextBlockMetadata {
   suggestions: Array<Suggestion>;

blocks/text/server.ts

@@ -0,0 +1,70 @@
+import { smoothStream, streamText } from 'ai';
+import { myProvider } from '@/lib/ai/models';
+import { createDocumentHandler } from '@/lib/blocks/server';
+import { updateDocumentPrompt } from '@/lib/ai/prompts';
+
+export const textDocumentHandler = createDocumentHandler<'text'>({
+  kind: 'text',
+  onCreateDocument: async ({ title, dataStream }) => {
+    let draftContent = '';
+
+    const { fullStream } = streamText({
+      model: myProvider.languageModel('block-model'),
+      system:
+        'Write about the given topic. Markdown is supported. Use headings wherever appropriate.',
+      experimental_transform: smoothStream({ chunking: 'word' }),
+      prompt: title,
+    });
+
+    for await (const delta of fullStream) {
+      const { type } = delta;
+
+      if (type === 'text-delta') {
+        const { textDelta } = delta;
+
+        draftContent += textDelta;
+
+        dataStream.writeData({
+          type: 'text-delta',
+          content: textDelta,
+        });
+      }
+    }
+
+    return draftContent;
+  },
+  onUpdateDocument: async ({ document, description, dataStream }) => {
+    let draftContent = '';
+
+    const { fullStream } = streamText({
+      model: myProvider.languageModel('block-model'),
+      system: updateDocumentPrompt(document.content, 'text'),
+      experimental_transform: smoothStream({ chunking: 'word' }),
+      prompt: description,
+      experimental_providerMetadata: {
+        openai: {
+          prediction: {
+            type: 'content',
+            content: document.content,
+          },
+        },
+      },
+    });
+
+    for await (const delta of fullStream) {
+      const { type } = delta;
+
+      if (type === 'text-delta') {
+        const { textDelta } = delta;
+
+        draftContent += textDelta;
+        dataStream.writeData({
+          type: 'text-delta',
+          content: textDelta,
+        });
+      }
+    }
+
+    return draftContent;
+  },
+});

components/block-actions.tsx

@@ -91,6 +91,7 @@ export const BlockActions = memo(PureBlockActions, (prevProps, nextProps) => {
   if (prevProps.currentVersionIndex !== nextProps.currentVersionIndex)
     return false;
   if (prevProps.isCurrentVersion !== nextProps.isCurrentVersion) return false;
+  if (prevProps.block.content !== nextProps.block.content) return false;
 
   return true;
 });

components/block.tsx

@@ -26,11 +26,11 @@ import { BlockCloseButton } from './block-close-button';
 import { BlockMessages } from './block-messages';
 import { useSidebar } from './ui/sidebar';
 import { useBlock } from '@/hooks/use-block';
-import { textBlock } from '@/blocks/text';
-import { imageBlock } from '@/blocks/image';
-import { codeBlock } from '@/blocks/code';
+import { imageBlock } from '@/blocks/image/client';
+import { codeBlock } from '@/blocks/code/client';
+import { sheetBlock } from '@/blocks/sheet/client';
+import { textBlock } from '@/blocks/text/client';
 import equal from 'fast-deep-equal';
-import { sheetBlock } from '@/blocks/sheet';
 
 export const blockDefinitions = [textBlock, codeBlock, imageBlock, sheetBlock];
 export type BlockKind = (typeof blockDefinitions)[number]['kind'];