feat: add sheet block (#743)

Co-authored-by: Nicolas <[email protected]>

blocks/sheet.tsx

@@ -0,0 +1,115 @@
+import { Block } from '@/components/create-block';
+import {
+  CopyIcon,
+  LineChartIcon,
+  RedoIcon,
+  SparklesIcon,
+  UndoIcon,
+} from '@/components/icons';
+import { SpreadsheetEditor } from '@/components/sheet-editor';
+import { parse, unparse } from 'papaparse';
+import { toast } from 'sonner';
+
+type Metadata = any;
+
+export const sheetBlock = new Block<'sheet', Metadata>({
+  kind: 'sheet',
+  description: 'Useful for working with spreadsheets',
+  initialize: async () => {},
+  onStreamPart: ({ setBlock, streamPart }) => {
+    if (streamPart.type === 'sheet-delta') {
+      setBlock((draftBlock) => ({
+        ...draftBlock,
+        content: streamPart.content as string,
+        isVisible: true,
+        status: 'streaming',
+      }));
+    }
+  },
+  content: ({
+    content,
+    currentVersionIndex,
+    isCurrentVersion,
+    onSaveContent,
+    status,
+  }) => {
+    return (
+      <SpreadsheetEditor
+        content={content}
+        currentVersionIndex={currentVersionIndex}
+        isCurrentVersion={isCurrentVersion}
+        saveContent={onSaveContent}
+        status={status}
+      />
+    );
+  },
+  actions: [
+    {
+      icon: <UndoIcon size={18} />,
+      description: 'View Previous version',
+      onClick: ({ handleVersionChange }) => {
+        handleVersionChange('prev');
+      },
+      isDisabled: ({ currentVersionIndex }) => {
+        if (currentVersionIndex === 0) {
+          return true;
+        }
+
+        return false;
+      },
+    },
+    {
+      icon: <RedoIcon size={18} />,
+      description: 'View Next version',
+      onClick: ({ handleVersionChange }) => {
+        handleVersionChange('next');
+      },
+      isDisabled: ({ isCurrentVersion }) => {
+        if (isCurrentVersion) {
+          return true;
+        }
+
+        return false;
+      },
+    },
+    {
+      icon: <CopyIcon />,
+      description: 'Copy as .csv',
+      onClick: ({ content }) => {
+        const parsed = parse<string[]>(content, { skipEmptyLines: true });
+
+        const nonEmptyRows = parsed.data.filter((row) =>
+          row.some((cell) => cell.trim() !== ''),
+        );
+
+        const cleanedCsv = unparse(nonEmptyRows);
+
+        navigator.clipboard.writeText(cleanedCsv);
+        toast.success('Copied csv to clipboard!');
+      },
+    },
+  ],
+  toolbar: [
+    {
+      description: 'Format and clean data',
+      icon: <SparklesIcon />,
+      onClick: ({ appendMessage }) => {
+        appendMessage({
+          role: 'user',
+          content: 'Can you please format and clean the data?',
+        });
+      },
+    },
+    {
+      description: 'Analyze and visualize data',
+      icon: <LineChartIcon />,
+      onClick: ({ appendMessage }) => {
+        appendMessage({
+          role: 'user',
+          content:
+            'Can you please analyze and visualize the data by creating a new code block in python?',
+        });
+      },
+    },
+  ],
+});

components/block.tsx

@@ -30,8 +30,9 @@ import { textBlock } from '@/blocks/text';
 import { imageBlock } from '@/blocks/image';
 import { codeBlock } from '@/blocks/code';
 import equal from 'fast-deep-equal';
+import { sheetBlock } from '@/blocks/sheet';
 
-export const blockDefinitions = [textBlock, codeBlock, imageBlock] as const;
+export const blockDefinitions = [textBlock, codeBlock, imageBlock, sheetBlock];
 export type BlockKind = (typeof blockDefinitions)[number]['kind'];
 
 export interface UIBlock {

components/data-stream-handler.tsx

@@ -10,6 +10,7 @@ export type DataStreamDelta = {
   type:
     | 'text-delta'
     | 'code-delta'
+    | 'sheet-delta'
     | 'image-delta'
     | 'title'
     | 'id'

components/document-preview.tsx

@@ -19,6 +19,7 @@ import { DocumentToolCall, DocumentToolResult } from './document';
 import { CodeEditor } from './code-editor';
 import { useBlock } from '@/hooks/use-block';
 import equal from 'fast-deep-equal';
+import { SpreadsheetEditor } from './sheet-editor';
 import { ImageEditor } from './image-editor';
 
 interface DocumentPreviewProps {
@@ -43,6 +44,7 @@ export function DocumentPreview({
 
   useEffect(() => {
     const boundingBox = hitboxRef.current?.getBoundingClientRect();
+
     if (block.documentId && boundingBox) {
       setBlock((block) => ({
         ...block,
@@ -256,6 +258,12 @@ const DocumentContent = ({ document }: { document: Document }) => {
             <CodeEditor {...commonProps} onSaveContent={() => {}} />
           </div>
         </div>
+      ) : document.kind === 'sheet' ? (
+        <div className="flex flex-1 relative size-full p-4">
+          <div className="absolute inset-0">
+            <SpreadsheetEditor {...commonProps} />
+          </div>
+        </div>
       ) : document.kind === 'image' ? (
         <ImageEditor
           title={document.title}

components/icons.tsx

@@ -1119,3 +1119,37 @@ export const FullscreenIcon = ({ size = 16 }: { size?: number }) => (
     ></path>
   </svg>
 );
+
+export const DownloadIcon = ({ size = 16 }: { size?: number }) => (
+  <svg
+    height={size}
+    strokeLinejoin="round"
+    viewBox="0 0 16 16"
+    width={size}
+    style={{ color: 'currentcolor' }}
+  >
+    <path
+      fillRule="evenodd"
+      clipRule="evenodd"
+      d="M8.75 1V1.75V8.68934L10.7197 6.71967L11.25 6.18934L12.3107 7.25L11.7803 7.78033L8.70711 10.8536C8.31658 11.2441 7.68342 11.2441 7.29289 10.8536L4.21967 7.78033L3.68934 7.25L4.75 6.18934L5.28033 6.71967L7.25 8.68934V1.75V1H8.75ZM13.5 9.25V13.5H2.5V9.25V8.5H1V9.25V14C1 14.5523 1.44771 15 2 15H14C14.5523 15 15 14.5523 15 14V9.25V8.5H13.5V9.25Z"
+      fill="currentColor"
+    ></path>
+  </svg>
+);
+
+export const LineChartIcon = ({ size = 16 }: { size?: number }) => (
+  <svg
+    height={size}
+    strokeLinejoin="round"
+    viewBox="0 0 16 16"
+    width={size}
+    style={{ color: 'currentcolor' }}
+  >
+    <path
+      fill="currentColor"
+      fillRule="evenodd"
+      d="M1 1v11.75A2.25 2.25 0 0 0 3.25 15H15v-1.5H3.25a.75.75 0 0 1-.75-.75V1H1Zm13.297 5.013.513-.547-1.094-1.026-.513.547-3.22 3.434-2.276-2.275a1 1 0 0 0-1.414 0L4.22 8.22l-.53.53 1.06 1.06.53-.53L7 7.56l2.287 2.287a1 1 0 0 0 1.437-.023l3.573-3.811Z"
+      clipRule="evenodd"
+    ></path>
+  </svg>
+);

components/sheet-editor.tsx

@@ -0,0 +1,143 @@
+'use client';
+
+import React, { memo, useEffect, useMemo, useState } from 'react';
+import DataGrid, { textEditor } from 'react-data-grid';
+import { parse, unparse } from 'papaparse';
+import { useTheme } from 'next-themes';
+import { cn } from '@/lib/utils';
+
+import 'react-data-grid/lib/styles.css';
+
+type SheetEditorProps = {
+  content: string;
+  saveContent: (content: string, isCurrentVersion: boolean) => void;
+  status: string;
+  isCurrentVersion: boolean;
+  currentVersionIndex: number;
+};
+
+const MIN_ROWS = 50;
+const MIN_COLS = 26;
+
+const PureSpreadsheetEditor = ({
+  content,
+  saveContent,
+  status,
+  isCurrentVersion,
+}: SheetEditorProps) => {
+  const { theme } = useTheme();
+
+  const parseData = useMemo(() => {
+    if (!content) return Array(MIN_ROWS).fill(Array(MIN_COLS).fill(''));
+    const result = parse<string[]>(content, { skipEmptyLines: true });
+
+    const paddedData = result.data.map((row) => {
+      const paddedRow = [...row];
+      while (paddedRow.length < MIN_COLS) {
+        paddedRow.push('');
+      }
+      return paddedRow;
+    });
+
+    while (paddedData.length < MIN_ROWS) {
+      paddedData.push(Array(MIN_COLS).fill(''));
+    }
+
+    return paddedData;
+  }, [content]);
+
+  const columns = useMemo(() => {
+    const rowNumberColumn = {
+      key: 'rowNumber',
+      name: '',
+      frozen: true,
+      width: 50,
+      renderCell: ({ rowIdx }: { rowIdx: number }) => rowIdx + 1,
+      cellClass: 'border-t border-r dark:bg-zinc-950 dark:text-zinc-50',
+      headerCellClass: 'border-t border-r dark:bg-zinc-900 dark:text-zinc-50',
+    };
+
+    const dataColumns = Array.from({ length: MIN_COLS }, (_, i) => ({
+      key: i.toString(),
+      name: String.fromCharCode(65 + i),
+      renderEditCell: textEditor,
+      width: 120,
+      cellClass: cn(`border-t dark:bg-zinc-950 dark:text-zinc-50`, {
+        'border-l': i !== 0,
+      }),
+      headerCellClass: cn(`border-t dark:bg-zinc-900 dark:text-zinc-50`, {
+        'border-l': i !== 0,
+      }),
+    }));
+
+    return [rowNumberColumn, ...dataColumns];
+  }, []);
+
+  const initialRows = useMemo(() => {
+    return parseData.map((row, rowIndex) => {
+      const rowData: any = {
+        id: rowIndex,
+        rowNumber: rowIndex + 1,
+      };
+
+      columns.slice(1).forEach((col, colIndex) => {
+        rowData[col.key] = row[colIndex] || '';
+      });
+
+      return rowData;
+    });
+  }, [parseData, columns]);
+
+  const [localRows, setLocalRows] = useState(initialRows);
+
+  useEffect(() => {
+    setLocalRows(initialRows);
+  }, [initialRows]);
+
+  const generateCsv = (data: any[][]) => {
+    return unparse(data);
+  };
+
+  const handleRowsChange = (newRows: any[]) => {
+    setLocalRows(newRows);
+
+    const updatedData = newRows.map((row) => {
+      return columns.slice(1).map((col) => row[col.key] || '');
+    });
+
+    const newCsvContent = generateCsv(updatedData);
+    saveContent(newCsvContent, true);
+  };
+
+  return (
+    <DataGrid
+      className={theme === 'dark' ? 'rdg-dark' : 'rdg-light'}
+      columns={columns}
+      rows={localRows}
+      enableVirtualization
+      onRowsChange={handleRowsChange}
+      onCellClick={(args) => {
+        if (args.column.key !== 'rowNumber') {
+          args.selectCell(true);
+        }
+      }}
+      style={{ height: '100%' }}
+      defaultColumnOptions={{
+        resizable: true,
+        sortable: true,
+      }}
+    />
+  );
+};
+
+function areEqual(prevProps: SheetEditorProps, nextProps: SheetEditorProps) {
+  return (
+    prevProps.currentVersionIndex === nextProps.currentVersionIndex &&
+    prevProps.isCurrentVersion === nextProps.isCurrentVersion &&
+    !(prevProps.status === 'streaming' && nextProps.status === 'streaming') &&
+    prevProps.content === nextProps.content &&
+    prevProps.saveContent === nextProps.saveContent
+  );
+}
+
+export const SpreadsheetEditor = memo(PureSpreadsheetEditor, areEqual);

components/toolbar.tsx

@@ -33,6 +33,7 @@ import {
   LogsIcon,
   MessageIcon,
   PenIcon,
+  SparklesIcon,
   StopIcon,
   SummarizeIcon,
 } from './icons';

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}
+`
+        : '';