modularize stream parts processing

vercel · Jan 27, 2025 · bd435b7 · bd435b7
1 parent a36412f
commit bd435b7
Show file tree

Hide file tree

Showing 10 changed files with 258 additions and 126 deletions.
diff --git a/blocks/actions.ts b/blocks/actions.ts
@@ -0,0 +1,8 @@
+'use server';
+
+import { getSuggestionsByDocumentId } from '@/lib/db/queries';
+
+export async function getSuggestions({ documentId }: { documentId: string }) {
+  const suggestions = await getSuggestionsByDocumentId({ documentId });
+  return suggestions;
+}
diff --git a/blocks/code.tsx b/blocks/code.tsx
@@ -9,23 +9,101 @@ import {
   UndoIcon,
 } from '@/components/icons';
 import { toast } from 'sonner';
+import { generateUUID } from '@/lib/utils';
+import { Console, ConsoleOutput } from '@/components/console';
 
-export const codeBlock = new Block({
+interface Metadata {
+  outputs: Array<ConsoleOutput>;
+}
+
+export const codeBlock = new Block<'code', Metadata>({
   kind: 'code',
   description:
     'Useful for code generation; Code execution is only available for python code.',
-  content: ({ ...props }) => <CodeEditor {...props} />,
+  initialize: () => ({
+    outputs: [],
+  }),
+  onStreamPart: ({ streamPart, setBlock }) => {
+    if (streamPart.type === 'code-delta') {
+      setBlock((draftBlock) => ({
+        ...draftBlock,
+        content: streamPart.content as string,
+        isVisible:
+          draftBlock.status === 'streaming' &&
+          draftBlock.content.length > 300 &&
+          draftBlock.content.length < 310
+            ? true
+            : draftBlock.isVisible,
+        status: 'streaming',
+      }));
+    }
+  },
+  content: ({ metadata, setMetadata, ...props }) => {
+    return (
+      <div className="flex flex-col gap-4 bg-red-500">
+        <CodeEditor {...props} />
+
+        {metadata?.outputs && (
+          <Console
+            consoleOutputs={metadata.outputs}
+            setConsoleOutputs={() => {
+              setMetadata({
+                ...metadata,
+                outputs: [],
+              });
+            }}
+          />
+        )}
+      </div>
+    );
+  },
   actions: [
     {
       icon: <PlayIcon size={18} />,
       label: 'Run',
       description: 'Execute code',
-      onClick: ({ content }) => {
-        // TODO
-        // 1. Initialize pyodide
-        // 2. Run code
-        // 3. Loading third party packages
-        // 4. Display output
+      onClick: async ({ content, setMetadata }) => {
+        const runId = generateUUID();
+        const outputs: any[] = [];
+
+        // @ts-expect-error - loadPyodide is not defined
+        const currentPyodideInstance = await globalThis.loadPyodide({
+          indexURL: 'https://cdn.jsdelivr.net/pyodide/v0.23.4/full/',
+        });
+
+        currentPyodideInstance.setStdout({
+          batched: (output: string) => {
+            outputs.push({
+              id: runId,
+              contents: [
+                {
+                  type: output.startsWith('data:image/png;base64')
+                    ? 'image'
+                    : 'text',
+                  value: output,
+                },
+              ],
+              status: 'completed',
+            });
+          },
+        });
+
+        await currentPyodideInstance.loadPackagesFromImports(content, {
+          messageCallback: (message: string) => {
+            outputs.push({
+              id: runId,
+              contents: [{ type: 'text', value: message }],
+              status: 'loading_packages',
+            });
+          },
+        });
+
+        await currentPyodideInstance.runPythonAsync(content);
+
+        setMetadata((metadata: any) => ({
+          ...metadata,
+          outputs,
+        }));
       },
     },
     {

diff --git a/blocks/image.tsx b/blocks/image.tsx
@@ -5,6 +5,16 @@ import { ImageEditor } from '@/components/image-editor';
 export const imageBlock = new Block({
   kind: 'image',
   description: 'Useful for image generation',
+  onStreamPart: ({ streamPart, setBlock }) => {
+    if (streamPart.type === 'image-delta') {
+      setBlock((draftBlock) => ({
+        ...draftBlock,
+        content: streamPart.content as string,
+        isVisible: true,
+        status: 'streaming',
+      }));
+    }
+  },
   content: ImageEditor,
   actions: [
     {

diff --git a/blocks/text.tsx b/blocks/text.tsx
@@ -10,24 +10,65 @@ import {
   RedoIcon,
   UndoIcon,
 } from '@/components/icons';
+import { Suggestion } from '@/lib/db/schema';
 import { toast } from 'sonner';
+import { getSuggestions } from './actions';
 
-export const textBlock = new Block({
+interface TextBlockMetadata {
+  suggestions: Array<Suggestion>;
+}
+
+export const textBlock = new Block<'text', TextBlockMetadata>({
   kind: 'text',
   description: 'Useful for text content, like drafting essays and emails.',
+  initialize: async ({ documentId, setMetadata }) => {
+    const suggestions = await getSuggestions({ documentId });
+
+    setMetadata({
+      suggestions,
+    });
+  },
+  onStreamPart: ({ streamPart, setMetadata, setBlock }) => {
+    if (streamPart.type === 'suggestion') {
+      setMetadata((metadata) => {
+        return {
+          suggestions: [
+            ...metadata.suggestions,
+            streamPart.content as Suggestion,
+          ],
+        };
+      });
+    }
+
+    if (streamPart.type === 'text-delta') {
+      setBlock((draftBlock) => {
+        return {
+          ...draftBlock,
+          content: draftBlock.content + (streamPart.content as string),
+          isVisible:
+            draftBlock.status === 'streaming' &&
+            draftBlock.content.length > 400 &&
+            draftBlock.content.length < 450
+              ? true
+              : draftBlock.isVisible,
+          status: 'streaming',
+        };
+      });
+    }
+  },
   content: ({
     mode,
     status,
     content,
     isCurrentVersion,
     currentVersionIndex,
     onSaveContent,
-    suggestions,
     getDocumentContentById,
     isLoading,
+    metadata,
   }) => {
     if (isLoading) {
-      <DocumentSkeleton blockKind="text" />;
+      return <DocumentSkeleton blockKind="text" />;
     }
 
     if (mode === 'diff') {
@@ -40,7 +81,7 @@ export const textBlock = new Block({
     return (
       <Editor
         content={content}
-        suggestions={suggestions}
+        suggestions={metadata ? metadata.suggestions : []}
         isCurrentVersion={isCurrentVersion}
         currentVersionIndex={currentVersionIndex}
         status={status}
@@ -55,7 +96,7 @@ export const textBlock = new Block({
       onClick: ({ handleVersionChange }) => {
         handleVersionChange('toggle');
       },
-      isDisabled: ({ currentVersionIndex }) => {
+      isDisabled: ({ currentVersionIndex, setMetadata }) => {
         if (currentVersionIndex === 0) {
           return true;
         }

diff --git a/components/block-actions.tsx b/components/block-actions.tsx
@@ -11,7 +11,8 @@ interface BlockActionsProps {
   currentVersionIndex: number;
   isCurrentVersion: boolean;
   mode: 'edit' | 'diff';
-  setConsoleOutputs: Dispatch<SetStateAction<Array<ConsoleOutput>>>;
+  metadata: any;
+  setMetadata: Dispatch<SetStateAction<any>>;
 }
 
 function PureBlockActions({
@@ -20,7 +21,8 @@ function PureBlockActions({
   currentVersionIndex,
   isCurrentVersion,
   mode,
-  setConsoleOutputs,
+  metadata,
+  setMetadata,
 }: BlockActionsProps) {
   const blockDefinition = blockDefinitions.find(
     (definition) => definition.kind === block.kind,
@@ -36,6 +38,8 @@ function PureBlockActions({
     currentVersionIndex,
     isCurrentVersion,
     mode,
+    metadata,
+    setMetadata,
   };
 
   return (

diff --git a/components/block.tsx b/components/block.tsx
@@ -16,15 +16,14 @@ import {
 } from 'react';
 import useSWR, { useSWRConfig } from 'swr';
 import { useDebounceCallback, useWindowSize } from 'usehooks-ts';
-import type { Document, Suggestion, Vote } from '@/lib/db/schema';
+import type { Document, Vote } from '@/lib/db/schema';
 import { cn, fetcher } from '@/lib/utils';
 import { MultimodalInput } from './multimodal-input';
 import { Toolbar } from './toolbar';
 import { VersionFooter } from './version-footer';
 import { BlockActions } from './block-actions';
 import { BlockCloseButton } from './block-close-button';
 import { BlockMessages } from './block-messages';
-import { Console } from './console';
 import { useSidebar } from './ui/sidebar';
 import { useBlock } from '@/hooks/use-block';
 import equal from 'fast-deep-equal';
@@ -50,17 +49,6 @@ export interface UIBlock {
   };
 }
 
-export interface ConsoleOutputContent {
-  type: 'text' | 'image';
-  value: string;
-}
-
-export interface ConsoleOutput {
-  id: string;
-  status: 'in_progress' | 'loading_packages' | 'completed' | 'failed';
-  contents: Array<ConsoleOutputContent>;
-}
-
 function PureBlock({
   chatId,
   input,
@@ -102,7 +90,7 @@ function PureBlock({
   ) => Promise<string | null | undefined>;
   isReadonly: boolean;
 }) {
-  const { block, setBlock } = useBlock();
+  const { block, setBlock, metadata, setMetadata } = useBlock();
 
   const {
     data: documents,
@@ -115,22 +103,9 @@ function PureBlock({
     fetcher,
   );
 
-  const { data: suggestions } = useSWR<Array<Suggestion>>(
-    documents && block && block.status !== 'streaming'
-      ? `/api/suggestions?documentId=${block.documentId}`
-      : null,
-    fetcher,
-    {
-      dedupingInterval: 5000,
-    },
-  );
-
   const [mode, setMode] = useState<'edit' | 'diff'>('edit');
   const [document, setDocument] = useState<Document | null>(null);
   const [currentVersionIndex, setCurrentVersionIndex] = useState(-1);
-  const [consoleOutputs, setConsoleOutputs] = useState<Array<ConsoleOutput>>(
-    [],
-  );
 
   const { open: isSidebarOpen } = useSidebar();
 
@@ -273,7 +248,14 @@ function PureBlock({
     throw new Error('Block definition not found!');
   }
 
-  const toolsByBlockKind = blockDefinition.toolbar;
+  useEffect(() => {
+    if (block && block.documentId !== 'init') {
+      blockDefinition.initialize({
+        documentId: block.documentId,
+        setMetadata,
+      });
+    }
+  }, [block, blockDefinition, setMetadata]);
 
   return (
     <AnimatePresence>
@@ -464,7 +446,8 @@ function PureBlock({
                 handleVersionChange={handleVersionChange}
                 isCurrentVersion={isCurrentVersion}
                 mode={mode}
-                setConsoleOutputs={setConsoleOutputs}
+                metadata={metadata}
+                setMetadata={setMetadata}
               />
             </div>
 
@@ -493,17 +476,19 @@ function PureBlock({
                   mode={mode}
                   status={block.status}
                   currentVersionIndex={currentVersionIndex}
-                  suggestions={suggestions ?? []}
+                  suggestions={[]}
                   onSaveContent={saveContent}
                   isInline={false}
                   isCurrentVersion={isCurrentVersion}
                   getDocumentContentById={getDocumentContentById}
                   isLoading={isDocumentsFetching && !block.content}
+                  metadata={metadata}
+                  setMetadata={setMetadata}
                 />
 
-                {suggestions && suggestions.length > 0 ? (
+                {/* {suggestions && suggestions.length > 0 ? (
                   <div className="md:hidden h-dvh w-12 shrink-0" />
-                ) : null}
+                ) : null} */}
 
                 <AnimatePresence>
                   {isCurrentVersion && (
@@ -530,13 +515,6 @@ function PureBlock({
                 />
               )}
             </AnimatePresence>
-
-            <AnimatePresence>
-              <Console
-                consoleOutputs={consoleOutputs}
-                setConsoleOutputs={setConsoleOutputs}
-              />
-            </AnimatePresence>
           </motion.div>
         </motion.div>
       )}