feat: add image block type (vercel#709)

spookymargiela007 · Jan 15, 2025 · df449a4 · df449a4
1 parent 6c23a8a
commit df449a4
Show file tree

Hide file tree

Showing 13 changed files with 229 additions and 92 deletions.
diff --git a/app/(chat)/api/chat/route.ts b/app/(chat)/api/chat/route.ts
@@ -2,13 +2,14 @@ import {
   type Message,
   convertToCoreMessages,
   createDataStreamResponse,
+  experimental_generateImage,
   streamObject,
   streamText,
 } from 'ai';
 import { z } from 'zod';
 
 import { auth } from '@/app/(auth)/auth';
-import { customModel } from '@/lib/ai';
+import { customModel, imageGenerationModel } from '@/lib/ai';
 import { models } from '@/lib/ai/models';
 import {
   codePrompt,
@@ -124,10 +125,10 @@ export async function POST(request: Request) {
           },
           createDocument: {
             description:
-              'Create a document for a writing activity. 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 like image generation. 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']),
+              kind: z.enum(['text', 'code', 'image']),
             }),
             execute: async ({ title, kind }) => {
               const id = generateUUID();
@@ -204,6 +205,21 @@ export async function POST(request: Request) {
                   }
                 }
 
+                dataStream.writeData({ type: 'finish', content: '' });
+              } else if (kind === 'image') {
+                const { image } = await experimental_generateImage({
+                  model: imageGenerationModel,
+                  prompt: title,
+                  n: 1,
+                });
+
+                draftText = image.base64;
+
+                dataStream.writeData({
+                  type: 'image-delta',
+                  content: image.base64,
+                });
+
                 dataStream.writeData({ type: 'finish', content: '' });
               }
 
@@ -309,6 +325,21 @@ export async function POST(request: Request) {
                   }
                 }
 
+                dataStream.writeData({ type: 'finish', content: '' });
+              } else if (document.kind === 'image') {
+                const { image } = await experimental_generateImage({
+                  model: imageGenerationModel,
+                  prompt: description,
+                  n: 1,
+                });
+
+                draftText = image.base64;
+
+                dataStream.writeData({
+                  type: 'image-delta',
+                  content: image.base64,
+                });
+
                 dataStream.writeData({ type: 'finish', content: '' });
               }
 

diff --git a/components/block-actions.tsx b/components/block-actions.tsx
@@ -2,11 +2,11 @@ import { cn } from '@/lib/utils';
 import { ClockRewind, CopyIcon, RedoIcon, UndoIcon } from './icons';
 import { Button } from './ui/button';
 import { Tooltip, TooltipContent, TooltipTrigger } from './ui/tooltip';
-import { useCopyToClipboard } from 'usehooks-ts';
 import { toast } from 'sonner';
 import { ConsoleOutput, UIBlock } from './block';
 import { Dispatch, memo, SetStateAction } from 'react';
 import { RunCodeButton } from './run-code-button';
+import { useMultimodalCopyToClipboard } from '@/hooks/use-multimodal-copy-to-clipboard';
 
 interface BlockActionsProps {
   block: UIBlock;
@@ -25,7 +25,8 @@ function PureBlockActions({
   mode,
   setConsoleOutputs,
 }: BlockActionsProps) {
-  const [_, copyToClipboard] = useCopyToClipboard();
+  const { copyTextToClipboard, copyImageToClipboard } =
+    useMultimodalCopyToClipboard();
 
   return (
     <div className="flex flex-row gap-1">
@@ -96,7 +97,12 @@ function PureBlockActions({
             variant="outline"
             className="p-2 h-fit dark:hover:bg-zinc-700"
             onClick={() => {
-              copyToClipboard(block.content);
+              if (block.kind === 'image') {
+                copyImageToClipboard(block.content);
+              } else {
+                copyTextToClipboard(block.content);
+              }
+
               toast.success('Copied to clipboard!');
             }}
             disabled={block.status === 'streaming'}

diff --git a/components/block.tsx b/components/block.tsx
@@ -34,8 +34,9 @@ import { Console } from './console';
 import { useSidebar } from './ui/sidebar';
 import { useBlock } from '@/hooks/use-block';
 import equal from 'fast-deep-equal';
+import { ImageEditor } from './image-editor';
 
-export type BlockKind = 'text' | 'code';
+export type BlockKind = 'text' | 'code' | 'image';
 
 export interface UIBlock {
   title: string;
@@ -476,7 +477,7 @@ function PureBlock({
                 })}
               >
                 {isDocumentsFetching && !block.content ? (
-                  <DocumentSkeleton />
+                  <DocumentSkeleton blockKind={block.kind} />
                 ) : block.kind === 'code' ? (
                   <CodeEditor
                     content={
@@ -512,9 +513,22 @@ function PureBlock({
                       newContent={getDocumentContentById(currentVersionIndex)}
                     />
                   )
+                ) : block.kind === 'image' ? (
+                  <ImageEditor
+                    title={block.title}
+                    content={
+                      isCurrentVersion
+                        ? block.content
+                        : getDocumentContentById(currentVersionIndex)
+                    }
+                    isCurrentVersion={isCurrentVersion}
+                    currentVersionIndex={currentVersionIndex}
+                    status={block.status}
+                    isInline={false}
+                  />
                 ) : null}
 
-                {suggestions ? (
+                {suggestions && suggestions.length > 0 ? (
                   <div className="md:hidden h-dvh w-12 shrink-0" />
                 ) : null}
 

diff --git a/components/data-stream-handler.tsx b/components/data-stream-handler.tsx
@@ -6,13 +6,13 @@ import { BlockKind } from './block';
 import { Suggestion } from '@/lib/db/schema';
 import { initialBlockData, useBlock } from '@/hooks/use-block';
 import { useUserMessageId } from '@/hooks/use-user-message-id';
-import { cx } from 'class-variance-authority';
 import { useSWRConfig } from 'swr';
 
 type DataStreamDelta = {
   type:
     | 'text-delta'
     | 'code-delta'
+    | 'image-delta'
     | 'title'
     | 'id'
     | 'suggestion'
@@ -107,6 +107,14 @@ export function DataStreamHandler({ id }: { id: string }) {
               status: 'streaming',
             };
 
+          case 'image-delta':
+            return {
+              ...draftBlock,
+              content: delta.content as string,
+              isVisible: true,
+              status: 'streaming',
+            };
+
           case 'suggestion':
             setTimeout(() => {
               setOptimisticSuggestions((currentSuggestions) => [

diff --git a/components/document-preview.tsx b/components/document-preview.tsx
@@ -8,8 +8,8 @@ import {
   useMemo,
   useRef,
 } from 'react';
-import { UIBlock } from './block';
-import { FileIcon, FullscreenIcon, LoaderIcon } from './icons';
+import { BlockKind, UIBlock } from './block';
+import { FileIcon, FullscreenIcon, ImageIcon, LoaderIcon } from './icons';
 import { cn, fetcher } from '@/lib/utils';
 import { Document } from '@/lib/db/schema';
 import { InlineDocumentSkeleton } from './document-skeleton';
@@ -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 { ImageEditor } from './image-editor';
 
 interface DocumentPreviewProps {
   isReadonly: boolean;
@@ -78,7 +79,7 @@ export function DocumentPreview({
   }
 
   if (isDocumentsFetching) {
-    return <LoadingSkeleton />;
+    return <LoadingSkeleton blockKind={result.kind ?? args.kind} />;
   }
 
   const document: Document | null = previewDocument
@@ -94,21 +95,22 @@ export function DocumentPreview({
         }
       : null;
 
-  if (!document) return <LoadingSkeleton />;
+  if (!document) return <LoadingSkeleton blockKind={block.kind} />;
 
   return (
     <div className="relative w-full cursor-pointer">
       <HitboxLayer hitboxRef={hitboxRef} result={result} setBlock={setBlock} />
       <DocumentHeader
         title={document.title}
+        kind={document.kind}
         isStreaming={block.status === 'streaming'}
       />
       <DocumentContent document={document} />
     </div>
   );
 }
 
-const LoadingSkeleton = () => (
+const LoadingSkeleton = ({ blockKind }: { blockKind: BlockKind }) => (
   <div className="w-full">
     <div className="p-4 border rounded-t-2xl flex flex-row gap-2 items-center justify-between dark:bg-muted h-[57px] dark:border-zinc-700 border-b-0">
       <div className="flex flex-row items-center gap-3">
@@ -121,9 +123,15 @@ const LoadingSkeleton = () => (
         <FullscreenIcon />
       </div>
     </div>
-    <div className="overflow-y-scroll border rounded-b-2xl p-8 pt-4 bg-muted border-t-0 dark:border-zinc-700">
-      <InlineDocumentSkeleton />
-    </div>
+    {blockKind === 'image' ? (
+      <div className="overflow-y-scroll border rounded-b-2xl bg-muted border-t-0 dark:border-zinc-700">
+        <div className="animate-pulse h-[257px] bg-muted-foreground/20 w-full" />
+      </div>
+    ) : (
+      <div className="overflow-y-scroll border rounded-b-2xl p-8 pt-4 bg-muted border-t-0 dark:border-zinc-700">
+        <InlineDocumentSkeleton />
+      </div>
+    )}
   </div>
 );
 
@@ -185,9 +193,11 @@ const HitboxLayer = memo(PureHitboxLayer, (prevProps, nextProps) => {
 
 const PureDocumentHeader = ({
   title,
+  kind,
   isStreaming,
 }: {
   title: string;
+  kind: BlockKind;
   isStreaming: boolean;
 }) => (
   <div className="p-4 border rounded-t-2xl flex flex-row gap-2 items-start sm:items-center justify-between dark:bg-muted border-b-0 dark:border-zinc-700">
@@ -197,6 +207,8 @@ const PureDocumentHeader = ({
           <div className="animate-spin">
             <LoaderIcon />
           </div>
+        ) : kind === 'image' ? (
+          <ImageIcon />
         ) : (
           <FileIcon />
         )}
@@ -244,6 +256,15 @@ const DocumentContent = ({ document }: { document: Document }) => {
             <CodeEditor {...commonProps} />
           </div>
         </div>
+      ) : document.kind === 'image' ? (
+        <ImageEditor
+          title={document.title}
+          content={document.content ?? ''}
+          isCurrentVersion={true}
+          currentVersionIndex={0}
+          status={block.status}
+          isInline={true}
+        />
       ) : null}
     </div>
   );

diff --git a/components/document-skeleton.tsx b/components/document-skeleton.tsx
@@ -1,7 +1,13 @@
 'use client';
 
-export const DocumentSkeleton = () => {
-  return (
+import { BlockKind } from './block';
+
+export const DocumentSkeleton = ({ blockKind }: { blockKind: BlockKind }) => {
+  return blockKind === 'image' ? (
+    <div className="flex flex-col gap-4 w-full justify-center items-center h-[calc(100dvh-60px)]">
+      <div className="animate-pulse rounded-lg bg-muted-foreground/20 size-96" />
+    </div>
+  ) : (
     <div className="flex flex-col gap-4 w-full">
       <div className="animate-pulse rounded-lg h-12 bg-muted-foreground/20 w-1/2" />
       <div className="animate-pulse rounded-lg h-5 bg-muted-foreground/20 w-full" />

diff --git a/components/image-editor.tsx b/components/image-editor.tsx
@@ -0,0 +1,50 @@
+import { LoaderIcon } from './icons';
+import cn from 'classnames';
+
+interface ImageEditorProps {
+  title: string;
+  content: string;
+  isCurrentVersion: boolean;
+  currentVersionIndex: number;
+  status: string;
+  isInline: boolean;
+}
+
+export function ImageEditor({
+  title,
+  content,
+  isCurrentVersion,
+  currentVersionIndex,
+  status,
+  isInline,
+}: ImageEditorProps) {
+  return (
+    <div
+      className={cn('flex flex-row items-center justify-center w-full', {
+        'h-[calc(100dvh-60px)]': !isInline,
+        'h-[200px]': isInline,
+      })}
+    >
+      {status === 'streaming' ? (
+        <div className="flex flex-row gap-4 items-center">
+          {!isInline && (
+            <div className="animate-spin">
+              <LoaderIcon />
+            </div>
+          )}
+          <div>Generating Image...</div>
+        </div>
+      ) : (
+        <picture>
+          <img
+            className={cn('w-full h-fit max-w-[800px]', {
+              'p-0 md:p-20': !isInline,
+            })}
+            src={`data:image/png;base64,${content}`}
+            alt={title}
+          />
+        </picture>
+      )}
+    </div>
+  );
+}