feat: add image block type (#709)

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'}

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}
 

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) => [

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,13 +95,14 @@ 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} />
@@ -108,7 +110,7 @@ export function DocumentPreview({
   );
 }
 
-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>
   );

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" />

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>
+  );
+}

components/toolbar.tsx

@@ -29,13 +29,12 @@ import { sanitizeUIMessages } from '@/lib/utils';
 import {
   ArrowUpIcon,
   CodeIcon,
-  FileIcon,
   LogsIcon,
   MessageIcon,
   PenIcon,
+  SparklesIcon,
   StopIcon,
   SummarizeIcon,
-  TerminalIcon,
 } from './icons';
 import { BlockKind } from './block';
 
@@ -327,6 +326,7 @@ const toolsByBlockKind: Record<
       icon: <LogsIcon />,
     },
   ],
+  image: [],
 };
 
 export const Tools = ({
@@ -347,7 +347,7 @@ export const Tools = ({
   ) => Promise<string | null | undefined>;
   isAnimating: boolean;
   setIsToolbarVisible: Dispatch<SetStateAction<boolean>>;
-  blockKind: 'text' | 'code';
+  blockKind: BlockKind;
 }) => {
   const [primaryTool, ...secondaryTools] = toolsByBlockKind[blockKind];
 
@@ -407,7 +407,7 @@ const PureToolbar = ({
   ) => Promise<string | null | undefined>;
   stop: () => void;
   setMessages: Dispatch<SetStateAction<Message[]>>;
-  blockKind: 'text' | 'code';
+  blockKind: BlockKind;
 }) => {
   const toolbarRef = useRef<HTMLDivElement>(null);
   const timeoutRef = useRef<ReturnType<typeof setTimeout>>();
@@ -451,6 +451,10 @@ const PureToolbar = ({
     }
   }, [isLoading, setIsToolbarVisible]);
 
+  if (toolsByBlockKind[blockKind].length === 0) {
+    return null;
+  }
+
   return (
     <TooltipProvider delayDuration={0}>
       <motion.div