feat: image generation (#62)

* ui for image generation

* added intl

* updated prisma schema

* added image generation route

* updated UI

* shared fixes

* added fixed for input

Author: mikita-kandratsyeu

actions/chat/get-chat-conversations.ts

@@ -23,7 +23,10 @@ export const getChatConversations = async ({ sharedConversationId }: GetChatConv
         include: {
           conversation: {
             include: {
-              messages: true,
+              messages: {
+                orderBy: { createdAt: 'asc' },
+                include: { imageGeneration: true },
+              },
               shared: true,
               user: true,
             },
@@ -66,7 +69,7 @@ export const getChatConversations = async ({ sharedConversationId }: GetChatConv
         user: true,
         messages: {
           orderBy: { createdAt: 'asc' },
-          include: { feedback: true },
+          include: { feedback: true, imageGeneration: true },
         },
       },
     });

actions/uploadthing/upload-files.ts

@@ -0,0 +1,20 @@
+'use server';
+
+import { UTFile } from 'uploadthing/server';
+
+import { base64ToBlob } from '@/lib/utils';
+import { utapi } from '@/server/uploadthing';
+
+export const uploadFiles = async (
+  files: { name: string; contentType: string; base64: string }[],
+) => {
+  const UTFiles = files.map(({ name, contentType, base64 }) => {
+    const file = new UTFile([base64ToBlob(base64, contentType)], name);
+
+    return file;
+  });
+
+  const response = await utapi.uploadFiles(UTFiles);
+
+  return response;
+};

app/(chat)/(routes)/chat/[[...slug]]/_components/chat-main/chat-body.tsx

@@ -1,6 +1,6 @@
 'use client';
 
-import { ArrowDown } from 'lucide-react';
+import { ArrowDown, LoaderPinwheel } from 'lucide-react';
 import { useTranslations } from 'next-intl';
 import React, {
   createContext,
@@ -53,6 +53,7 @@ const Content = ({ children }: ContentProps) => {
 };
 
 type ChatBodyProps = {
+  assistantImage?: string;
   assistantMessage?: string;
   introMessages: string[];
   isShared?: boolean;
@@ -67,6 +68,7 @@ type ChatBodyProps = {
 };
 
 export const ChatBody = ({
+  assistantImage,
   assistantMessage,
   introMessages,
   isShared,
@@ -77,9 +79,10 @@ export const ChatBody = ({
   const t = useTranslations('chat.body');
 
   const { user } = useCurrentUser();
-  const { chatMessages, conversationId } = useChatStore((state) => ({
+  const { chatMessages, conversationId, isImageGeneration } = useChatStore((state) => ({
     chatMessages: state.chatMessages,
     conversationId: state.conversationId,
+    isImageGeneration: state.isImageGeneration,
   }));
 
   const [sticky, setSticky] = useState(false);
@@ -138,16 +141,24 @@ export const ChatBody = ({
                   </div>
                 );
               })}
-              {assistantMessage && (
-                <div className="flex flex-1 text-base md:px-5 lg:px-1 xl:px-5 mx-auto gap-3 md:max-w-3xl lg:max-w-[40rem] xl:max-w-4xl px-4 first:mt-4 last:mb-6">
+
+              <div className="flex flex-1 text-base md:px-5 lg:px-1 xl:px-5 mx-auto gap-3 md:max-w-3xl lg:max-w-[40rem] xl:max-w-4xl px-4 first:mt-4 last:mb-6">
+                {!assistantMessage && isSubmitting && isImageGeneration && (
+                  <div className="flex gap-x-2 items-center justify-center w-full mt-4">
+                    <LoaderPinwheel className="h-4 w-4 animate-spin text-muted-foreground" />
+                    <p className="text-sm text-muted-foreground">{t('image-loading')}</p>
+                  </div>
+                )}
+                {assistantMessage && (
                   <ChatBubble
+                    streamImage={assistantImage}
                     isSubmitting={isSubmitting}
                     message={{ role: ChatCompletionRole.ASSISTANT, content: '' }}
                     name="Nova Copilot"
                     streamMessage={assistantMessage}
                   />
-                </div>
-              )}
+                )}
+              </div>
             </Content>
           </ScrollToBottom>
         )}

app/(chat)/(routes)/chat/[[...slug]]/_components/chat-main/chat-bubble.tsx

@@ -1,6 +1,8 @@
 'use client';
 
-import { MoreHorizontal } from 'lucide-react';
+import { Download, MoreHorizontal } from 'lucide-react';
+import Image from 'next/image';
+import Link from 'next/link';
 
 import { CopyClipboard } from '@/components/common/copy-clipboard';
 import { MarkdownText } from '@/components/common/markdown-text';
@@ -19,9 +21,15 @@ type ChatBubbleProps = {
     model?: string;
     id?: string;
     feedback?: { feedback: string } | null;
+    imageGeneration?: {
+      model: string;
+      revisedPrompt: string;
+      url: string;
+    } | null;
   };
   name: string;
   picture?: string | null;
+  streamImage?: string;
   streamMessage?: string;
 };
 
@@ -31,10 +39,14 @@ export const ChatBubble = ({
   message,
   name,
   picture,
+  streamImage,
   streamMessage,
 }: ChatBubbleProps) => {
   const isAssistant = message.role === ChatCompletionRole.ASSISTANT;
+
   const text = streamMessage ?? message.content;
+  const image = streamImage ?? message?.imageGeneration?.url;
+  const model = message?.imageGeneration ? message.imageGeneration.model : message.model;
 
   return (
     <div className="pb-4 pt-2">
@@ -49,18 +61,32 @@ export const ChatBubble = ({
         <div className="flex flex-col">
           <div className="flex items-center space-x-2">
             <span className="text-medium font-bold">{name}</span>
-            {Boolean(isAssistant && streamMessage && isSubmitting) && (
+            {Boolean(isAssistant && isSubmitting) && (
               <MoreHorizontal className="w-6 h-6 animate-pulse" />
             )}
             {isAssistant && isShared && (
-              <div className="text-xs text-muted-foreground">{message.model}</div>
+              <div className="text-xs text-muted-foreground">{model}</div>
             )}
           </div>
+          {image && (
+            <div className="relative aspect-w-16 aspect-h-14 border my-4">
+              <Image alt="Image" fill src={image} className="rounded-sm" />
+            </div>
+          )}
           <MarkdownText text={text} />
           {isAssistant && !isSubmitting && (
-            <div className="flex gap-x-3 mt-4">
+            <div className="flex gap-x-3 mt-4 items-center">
               <CopyClipboard textToCopy={text} />
-              <ChatFeedback messageId={message.id} state={message.feedback?.feedback} />
+              {image && (
+                <Link href={image} target="_blank">
+                  <button className="flex items-center">
+                    <Download className="h-4 w-4" />
+                  </button>
+                </Link>
+              )}
+              {!isShared && (
+                <ChatFeedback messageId={message.id} state={message.feedback?.feedback} />
+              )}
             </div>
           )}
         </div>

app/(chat)/(routes)/chat/[[...slug]]/_components/chat-main/chat-feedback.tsx

@@ -82,7 +82,7 @@ export const ChatFeedback = ({ className, disabled, messageId, state }: ChatFeed
       <button onClick={handleClick} disabled={disabled || isFetching} name={Feedback.POSITIVE}>
         <ThumbsUp
           className={cn(
-            'h-4 w-4 hover:text-green-500 transition-all duration-300',
+            'h-4 w-4 hover:text-green-500 transition-colors duration-300',
             state === Feedback.POSITIVE && 'text-green-500',
             clicked[Feedback.POSITIVE] ? 'animate-spin-once' : '',
             className,
@@ -92,7 +92,7 @@ export const ChatFeedback = ({ className, disabled, messageId, state }: ChatFeed
       <button onClick={handleClick} disabled={disabled || isFetching} name={Feedback.NEGATIVE}>
         <ThumbsDown
           className={cn(
-            'h-4 w-4 hover:text-red-500 transition-all duration-300',
+            'h-4 w-4 hover:text-red-500 transition-color duration-300',
             state === Feedback.NEGATIVE && 'text-red-500',
             clicked[Feedback.NEGATIVE] ? 'animate-spin-once' : '',
             className,

app/(chat)/(routes)/chat/[[...slug]]/_components/chat-main/chat-input-footer.tsx

@@ -0,0 +1,72 @@
+'use client';
+
+import { ImageIcon, Paperclip, SendHorizonal, StopCircle } from 'lucide-react';
+import { useTranslations } from 'next-intl';
+
+import { Badge, Button, Separator } from '@/components/ui';
+import { OPEN_AI_IMAGE_MODELS } from '@/constants/open-ai';
+import { useChatStore } from '@/hooks/store/use-chat-store';
+import { cn } from '@/lib/utils';
+
+type ChatInputFooterProps = {
+  isDisabled?: boolean;
+  isSubmitting?: boolean;
+  onSendMessage: () => void;
+};
+
+export const ChatInputFooter = ({
+  isDisabled,
+  isSubmitting,
+  onSendMessage,
+}: ChatInputFooterProps) => {
+  const t = useTranslations('chat.input');
+
+  const { isImageGeneration, setIsImageGeneration } = useChatStore((state) => ({
+    isImageGeneration: state.isImageGeneration,
+    setIsImageGeneration: state.setIsImageGeneration,
+  }));
+
+  return (
+    <div className="flex bg-background justify-between px-2 py-2 items-center">
+      <div className="text-xs text-muted-foreground flex items-center gap-x-2 pr-2">
+        {isImageGeneration && (
+          <Badge variant="secondary" className="rounded-sm px-1 font-normal line-clamp-2">
+            {t('image-generation-mode', { model: OPEN_AI_IMAGE_MODELS[0].label })}
+          </Badge>
+        )}
+      </div>
+      <div className="flex items-center">
+        <button
+          type="button"
+          className="mr-3"
+          disabled={isSubmitting}
+          onClick={() => setIsImageGeneration(!isImageGeneration)}
+        >
+          <ImageIcon
+            className={cn(
+              'w-4 h-4 text-muted-foreground transition-colors duration-300',
+              isImageGeneration && 'text-purple-500',
+            )}
+          />
+        </button>
+        <button type="button" disabled={isSubmitting}>
+          <Paperclip className="w-4 h-4 text-muted-foreground" />
+        </button>
+        <Separator orientation="vertical" className="mr-4 ml-2 h-6" />
+        <Button
+          className={cn(
+            'bg-gradient-to-r from-indigo-500 via-purple-500 to-pink-500 text-white hover:text-white font-medium z-10 px-2 text-sm',
+            isSubmitting && 'w-12',
+          )}
+          disabled={isDisabled || (isSubmitting && isImageGeneration)}
+          type={isSubmitting ? 'button' : 'submit'}
+          variant="outline"
+          onClick={onSendMessage}
+        >
+          {isSubmitting && <StopCircle className="w-4 h-4 mx-2" />}
+          {!isSubmitting && <SendHorizonal className="w-4 h-4 mx-2" />}
+        </Button>
+      </div>
+    </div>
+  );
+};

app/(chat)/(routes)/chat/[[...slug]]/_components/chat-main/chat-input.tsx

@@ -1,13 +1,15 @@
 'use client';
 
-import { SendHorizonal, StopCircle } from 'lucide-react';
 import { useTranslations } from 'next-intl';
 import { SyntheticEvent, useState } from 'react';
 
-import { Button, Textarea } from '@/components/ui';
+import { Textarea } from '@/components/ui';
 import { LIMIT_CHAT_INPUT } from '@/constants/chat';
+import { useChatStore } from '@/hooks/store/use-chat-store';
 import { cn } from '@/lib/utils';
 
+import { ChatInputFooter } from './chat-input-footer';
+
 type ChatInputProps = {
   currenMessage: string;
   isSubmitting?: boolean;
@@ -25,6 +27,10 @@ export const ChatInput = ({
 }: ChatInputProps) => {
   const t = useTranslations('chat.input');
 
+  const { isImageGeneration } = useChatStore((state) => ({
+    isImageGeneration: state.isImageGeneration,
+  }));
+
   const [inputLength, setInputLength] = useState(0);
 
   return (
@@ -35,14 +41,16 @@ export const ChatInput = ({
             className={cn(
               'mx-auto flex flex-col lg:max-w-2xl xl:max-w-4xl w-full h-full border rounded-sm z-10 focus-within:border-b-indigo-500 focus-within:border-b-2 transition-colors duration-200 ease-in-out',
               inputLength >= LIMIT_CHAT_INPUT && 'focus-within:border-b-red-600',
+              isImageGeneration &&
+                'border-b-purple-500 border-b-2 focus-within:border-b-purple-500 ',
             )}
             onSubmit={onSubmit}
           >
             <Textarea
               className="resize-none overflow-auto z-10 border-none"
               disabled={isSubmitting}
               maxLength={LIMIT_CHAT_INPUT}
-              placeholder={t('enterMessage')}
+              placeholder={t(isImageGeneration ? 'enterImageMessage' : 'enterMessage')}
               value={currenMessage}
               onChange={(event) => {
                 setCurrentMessage(event.target.value);
@@ -55,24 +63,11 @@ export const ChatInput = ({
                 }
               }}
             />
-            <div className="flex bg-background justify-between px-2 pb-4 pt-2 items-end">
-              <div className="text-xs text-muted-foreground">
-                {inputLength} / {LIMIT_CHAT_INPUT}
-              </div>
-              <Button
-                className={cn(
-                  'bg-gradient-to-r from-indigo-500 via-purple-500 to-pink-500 text-white hover:text-white font-medium z-10 px-2 text-sm',
-                  isSubmitting && 'w-12',
-                )}
-                disabled={!currenMessage && !isSubmitting}
-                type={isSubmitting ? 'button' : 'submit'}
-                variant="outline"
-                onClick={isSubmitting ? onAbortGenerating : () => {}}
-              >
-                {isSubmitting && <StopCircle className="w-4 h-4 mx-2" />}
-                {!isSubmitting && <SendHorizonal className="w-4 h-4 mx-2" />}
-              </Button>
-            </div>
+            <ChatInputFooter
+              isDisabled={!currenMessage && !isSubmitting}
+              isSubmitting={isSubmitting}
+              onSendMessage={isSubmitting ? onAbortGenerating : () => {}}
+            />
           </form>
           <div className="p-2 text-center text-xs text-muted-foreground z-10">{t('footer')}</div>
         </div>