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/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'];

lib/ai/tools/create-document.ts

@@ -1,17 +1,8 @@
 import { generateUUID } from '@/lib/utils';
-import {
-  DataStreamWriter,
-  experimental_generateImage,
-  smoothStream,
-  streamObject,
-  streamText,
-  tool,
-} from 'ai';
+import { DataStreamWriter, tool } from 'ai';
 import { z } from 'zod';
-import { codePrompt, sheetPrompt } from '../prompts';
-import { saveDocument } from '@/lib/db/queries';
 import { Session } from 'next-auth';
-import { myProvider } from '../models';
+import { blockKinds, documentHandlersByBlockKind } from '@/lib/blocks/server';
 
 interface CreateDocumentProps {
   session: Session;
@@ -24,11 +15,15 @@ export const createDocument = ({ session, dataStream }: CreateDocumentProps) =>
       'Create a document for a writing or content creation activities. This tool will call other functions that will generate the contents of the document based on the title and kind.',
     parameters: z.object({
       title: z.string(),
-      kind: z.enum(['text', 'code', 'image', 'sheet']),
+      kind: z.enum(blockKinds),
     }),
     execute: async ({ title, kind }) => {
       const id = generateUUID();
-      let draftText = '';
+
+      dataStream.writeData({
+        type: 'kind',
+        content: kind,
+      });
 
       dataStream.writeData({
         type: 'id',
@@ -40,129 +35,28 @@ export const createDocument = ({ session, dataStream }: CreateDocumentProps) =>
         content: title,
       });
 
-      dataStream.writeData({
-        type: 'kind',
-        content: kind,
-      });
-
       dataStream.writeData({
         type: 'clear',
         content: '',
       });
 
-      if (kind === 'text') {
-        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,
-        });
+      const documentHandler = documentHandlersByBlockKind.find(
+        (documentHandlerByBlockKind) =>
+          documentHandlerByBlockKind.kind === kind,
+      );
 
-        for await (const delta of fullStream) {
-          const { type } = delta;
-
-          if (type === 'text-delta') {
-            const { textDelta } = delta;
-
-            draftText += textDelta;
-            dataStream.writeData({
-              type: 'text-delta',
-              content: textDelta,
-            });
-          }
+      if (!documentHandler) {
+        throw new Error(`No document handler found for kind: ${kind}`);
       }
 
-        dataStream.writeData({ type: 'finish', content: '' });
-      } else if (kind === 'code') {
-        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 ?? '',
-              });
-
-              draftText = code;
-            }
-          }
-        }
-
-        dataStream.writeData({ type: 'finish', content: '' });
-      } else if (kind === 'image') {
-        const { image } = await experimental_generateImage({
-          model: myProvider.imageModel('small-model'),
-          prompt: title,
-          n: 1,
-        });
-
-        draftText = image.base64;
-
-        dataStream.writeData({
-          type: 'image-delta',
-          content: image.base64,
-        });
-
-        dataStream.writeData({ type: 'finish', content: '' });
-      } else if (kind === 'sheet') {
-        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,
-              });
-
-              draftText = csv;
-            }
-          }
-        }
-
-        dataStream.writeData({
-          type: 'sheet-delta',
-          content: draftText,
-        });
-
-        dataStream.writeData({ type: 'finish', content: '' });
-      }
-
-      if (session.user?.id) {
-        await saveDocument({
+      await documentHandler.onCreateDocument({
         id,
         title,
-          kind,
-          content: draftText,
-          userId: session.user.id,
+        dataStream,
+        session,
       });
-      }
+
+      dataStream.writeData({ type: 'finish', content: '' });
 
       return {
         id,

lib/ai/tools/update-document.ts

@@ -1,16 +1,8 @@
-import {
-  DataStreamWriter,
-  experimental_generateImage,
-  smoothStream,
-  streamObject,
-  streamText,
-  tool,
-} from 'ai';
+import { DataStreamWriter, tool } from 'ai';
 import { Session } from 'next-auth';
 import { z } from 'zod';
 import { getDocumentById, saveDocument } from '@/lib/db/queries';
-import { updateDocumentPrompt } from '../prompts';
-import { myProvider } from '../models';
+import { documentHandlersByBlockKind } from '@/lib/blocks/server';
 
 interface UpdateDocumentProps {
   session: Session;
@@ -35,129 +27,28 @@ export const updateDocument = ({ session, dataStream }: UpdateDocumentProps) =>
         };
       }
 
-      const { content: currentContent } = document;
-      let draftText = '';
-
       dataStream.writeData({
         type: 'clear',
         content: document.title,
       });
 
-      if (document.kind === 'text') {
-        const { fullStream } = streamText({
-          model: myProvider.languageModel('block-model'),
-          system: updateDocumentPrompt(currentContent, 'text'),
-          experimental_transform: smoothStream({ chunking: 'word' }),
-          prompt: description,
-          experimental_providerMetadata: {
-            openai: {
-              prediction: {
-                type: 'content',
-                content: currentContent,
-              },
-            },
-          },
-        });
+      const documentHandler = documentHandlersByBlockKind.find(
+        (documentHandlerByBlockKind) =>
+          documentHandlerByBlockKind.kind === document.kind,
+      );
 
-        for await (const delta of fullStream) {
-          const { type } = delta;
-
-          if (type === 'text-delta') {
-            const { textDelta } = delta;
-
-            draftText += textDelta;
-            dataStream.writeData({
-              type: 'text-delta',
-              content: textDelta,
-            });
-          }
+      if (!documentHandler) {
+        throw new Error(`No document handler found for kind: ${document.kind}`);
       }
 
-        dataStream.writeData({ type: 'finish', content: '' });
-      } else if (document.kind === 'code') {
-        const { fullStream } = streamObject({
-          model: myProvider.languageModel('block-model'),
-          system: updateDocumentPrompt(currentContent, '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 ?? '',
-              });
-
-              draftText = code;
-            }
-          }
-        }
-
-        dataStream.writeData({ type: 'finish', content: '' });
-      } else if (document.kind === 'image') {
-        const { image } = await experimental_generateImage({
-          model: myProvider.imageModel('image-model'),
-          prompt: description,
-          n: 1,
-        });
-
-        draftText = image.base64;
-
-        dataStream.writeData({
-          type: 'image-delta',
-          content: image.base64,
+      await documentHandler.onUpdateDocument({
+        document,
+        description,
+        dataStream,
+        session,
       });
 
       dataStream.writeData({ type: 'finish', content: '' });
-      } else if (document.kind === 'sheet') {
-        const { fullStream } = streamObject({
-          model: myProvider.languageModel('block-model'),
-          system: updateDocumentPrompt(currentContent, '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,
-              });
-
-              draftText = csv;
-            }
-          }
-        }
-
-        dataStream.writeData({ type: 'finish', content: '' });
-      }
-
-      if (session.user?.id) {
-        await saveDocument({
-          id,
-          title: document.title,
-          content: draftText,
-          kind: document.kind,
-          userId: session.user.id,
-        });
-      }
 
       return {
         id,

lib/blocks/server.ts

@@ -0,0 +1,99 @@
+import { codeDocumentHandler } from '@/blocks/code/server';
+import { imageDocumentHandler } from '@/blocks/image/server';
+import { sheetDocumentHandler } from '@/blocks/sheet/server';
+import { textDocumentHandler } from '@/blocks/text/server';
+import { BlockKind } from '@/components/block';
+import { DataStreamWriter } from 'ai';
+import { Document } from '../db/schema';
+import { saveDocument } from '../db/queries';
+import { Session } from 'next-auth';
+
+export interface SaveDocumentProps {
+  id: string;
+  title: string;
+  kind: BlockKind;
+  content: string;
+  userId: string;
+}
+
+export interface CreateDocumentCallbackProps {
+  id: string;
+  title: string;
+  dataStream: DataStreamWriter;
+  session: Session;
+}
+
+export interface UpdateDocumentCallbackProps {
+  document: Document;
+  description: string;
+  dataStream: DataStreamWriter;
+  session: Session;
+}
+
+export interface DocumentHandler<T = BlockKind> {
+  kind: T;
+  onCreateDocument: (args: CreateDocumentCallbackProps) => Promise<void>;
+  onUpdateDocument: (args: UpdateDocumentCallbackProps) => Promise<void>;
+}
+
+export function createDocumentHandler<T extends BlockKind>(config: {
+  kind: T;
+  onCreateDocument: (params: CreateDocumentCallbackProps) => Promise<string>;
+  onUpdateDocument: (params: UpdateDocumentCallbackProps) => Promise<string>;
+}): DocumentHandler<T> {
+  return {
+    kind: config.kind,
+    onCreateDocument: async (args: CreateDocumentCallbackProps) => {
+      const draftContent = await config.onCreateDocument({
+        id: args.id,
+        title: args.title,
+        dataStream: args.dataStream,
+        session: args.session,
+      });
+
+      if (args.session?.user?.id) {
+        await saveDocument({
+          id: args.id,
+          title: args.title,
+          content: draftContent,
+          kind: config.kind,
+          userId: args.session.user.id,
+        });
+      }
+
+      return;
+    },
+    onUpdateDocument: async (args: UpdateDocumentCallbackProps) => {
+      const draftContent = await config.onUpdateDocument({
+        document: args.document,
+        description: args.description,
+        dataStream: args.dataStream,
+        session: args.session,
+      });
+
+      if (args.session?.user?.id) {
+        await saveDocument({
+          id: args.document.id,
+          title: args.document.title,
+          content: draftContent,
+          kind: config.kind,
+          userId: args.session.user.id,
+        });
+      }
+
+      return;
+    },
+  };
+}
+
+/*
+ * Use this array to define the document handlers for each block kind.
+ */
+export const documentHandlersByBlockKind: Array<DocumentHandler> = [
+  textDocumentHandler,
+  codeDocumentHandler,
+  imageDocumentHandler,
+  sheetDocumentHandler,
+];
+
+export const blockKinds = ['text', 'code', 'image', 'sheet'] as const;