feat: add sheet block (#743)

Co-authored-by: Nicolas <nicolascamara29@gmail.com>

lib/ai/prompts.ts

@@ -64,6 +64,10 @@ print(f"Factorial of 5 is: {factorial(5)}")
 \`\`\`
 `;
 
+export const sheetPrompt = `
+You are a spreadsheet creation assistant. Create a spreadsheet in csv format based on the given prompt. The spreadsheet should contain meaningful column headers and data.
+`;
+
 export const updateDocumentPrompt = (
   currentContent: string | null,
   type: BlockKind,
@@ -80,4 +84,10 @@ Improve the following code snippet based on the given prompt.
 
 ${currentContent}
 `
-      : '';
+      : type === 'sheet'
+        ? `\
+Improve the following spreadsheet based on the given prompt.
+
+${currentContent}
+`
+        : '';

lib/ai/tools/create-document.ts

@@ -9,7 +9,7 @@ import {
 } from 'ai';
 import { z } from 'zod';
 import { customModel, imageGenerationModel } from '..';
-import { codePrompt } from '../prompts';
+import { codePrompt, sheetPrompt } from '../prompts';
 import { saveDocument } from '@/lib/db/queries';
 import { Session } from 'next-auth';
 import { Model } from '../models';
@@ -27,10 +27,10 @@ export const createDocument = ({
 }: CreateDocumentProps) =>
   tool({
     description:
-      'Create a document for a writing or content creation activities like image generation. This tool will call other functions that will generate the contents of the document based on the title and kind.',
+      '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']),
+      kind: z.enum(['text', 'code', 'image', 'sheet']),
     }),
     execute: async ({ title, kind }) => {
       const id = generateUUID();
@@ -123,6 +123,40 @@ export const createDocument = ({
           content: image.base64,
         });
 
+        dataStream.writeData({ type: 'finish', content: '' });
+      } else if (kind === 'sheet') {
+        const { fullStream } = streamObject({
+          model: customModel(model.apiIdentifier),
+          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: '' });
       }
 

lib/ai/tools/update-document.ts

@@ -123,6 +123,35 @@ export const updateDocument = ({
           content: image.base64,
         });
 
+        dataStream.writeData({ type: 'finish', content: '' });
+      } else if (document.kind === 'sheet') {
+        const { fullStream } = streamObject({
+          model: customModel(model.apiIdentifier),
+          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: '' });
       }
 

lib/db/schema.ts

@@ -72,7 +72,7 @@ export const document = pgTable(
     createdAt: timestamp('createdAt').notNull(),
     title: text('title').notNull(),
     content: text('content'),
-    kind: varchar('text', { enum: ['text', 'code', 'image'] })
+    kind: varchar('text', { enum: ['text', 'code', 'image', 'sheet'] })
       .notNull()
       .default('text'),
     userId: uuid('userId')