Merge pull request #122 from sebiweise/feature/image_generation

Add image generation
jorge-menjivar · Dec 13, 2023 · f89ffb0 · f89ffb0
2 parents 83267a1 + ca7ada6
commit f89ffb0
Show file tree

Hide file tree

Showing 30 changed files with 796 additions and 328 deletions.
diff --git a/apps/docs/package.json b/apps/docs/package.json
@@ -31,7 +31,7 @@
     "@docusaurus/tsconfig": "3.0.1",
     "@docusaurus/types": "3.0.1",
     "autoprefixer": "^10.4.16",
-    "postcss": "^8.4.31",
+    "postcss": "^8.4.32",
     "tailwindcss": "^3.3.5",
     "typescript": "~5.3.3"
   },

diff --git a/apps/unsaged/app/api/chat/route.ts b/apps/unsaged/app/api/chat/route.ts
@@ -2,6 +2,7 @@ import { getStream } from '@/utils/server/ai_vendors/stream';
 import { getTokenCount } from '@/utils/server/ai_vendors/token-count';
 
 import { ChatBody } from '@/types/chat';
+import { StreamingTextResponse } from 'ai';
 
 export const runtime = 'edge';
 
@@ -47,9 +48,7 @@ const handler = async (req: Request): Promise<Response> => {
     });
   }
 
-  return new Response(stream, {
-    headers: { 'Content-Type': 'text/event-stream; charset=utf-8' },
-  });
+  return new StreamingTextResponse(stream);
 };
 
 export { handler as GET, handler as POST };
diff --git a/apps/unsaged/app/api/image/route.ts b/apps/unsaged/app/api/image/route.ts
@@ -0,0 +1,36 @@
+import { getImage } from '@/utils/server/ai_vendors/image';
+
+import { ImageBody } from '@/types/chat';
+
+export const runtime = 'edge';
+
+const handler = async (req: Request): Promise<Response> => {
+  const { model, prompt, apiKey, params } =
+    (await req.json()) as ImageBody;
+
+  const { error: imageError, images } = await getImage(
+    model,
+    params,
+    apiKey,
+    prompt
+  );
+
+  if (imageError) {
+    let message = imageError;
+
+    if (message.message) {
+      message = message.message;
+    }
+
+    console.error(message);
+
+    return new Response('Error', {
+      status: 500,
+      statusText: message,
+    });
+  }
+
+  return new Response(JSON.stringify(images), { status: 200 });
+};
+
+export { handler as GET, handler as POST };
diff --git a/apps/unsaged/components/Home/components/ChatZone/Screens/Chat/ChatInput.tsx b/apps/unsaged/components/Home/components/ChatZone/Screens/Chat/ChatInput.tsx
@@ -70,15 +70,19 @@ export const ChatInput = ({
 
   const handleChange = (e: React.ChangeEvent<HTMLTextAreaElement>) => {
     const value = e.target.value;
-    const maxLength = selectedConversation?.model.maxLength;
-
-    if (maxLength && value.length > maxLength) {
-      alert(
-        t('messageLimit',
-          { maxLength, valueLength: value.length },
-        ),
-      );
-      return;
+
+    if (selectedConversation?.model.type == 'text') {
+      const maxLength = selectedConversation?.model.maxLength;
+
+      if (maxLength && value.length > maxLength) {
+        alert(
+          t(
+            `messageLimit`,
+            { maxLength, valueLength: value.length },
+          ),
+        );
+        return;
+      }
     }
 
     setContent(value);
@@ -326,8 +330,8 @@ export const ChatInput = ({
               bottom: `${textareaRef?.current?.scrollHeight}px`,
               maxHeight: '400px',
               overflow: `${textareaRef.current && textareaRef.current.scrollHeight > 400
-                ? 'auto'
-                : 'hidden'
+                  ? 'auto'
+                  : 'hidden'
                 }`,
             }}
             placeholder={

diff --git a/apps/unsaged/components/Home/components/ChatZone/Screens/Chat/ChatMessage.tsx b/apps/unsaged/components/Home/components/ChatZone/Screens/Chat/ChatMessage.tsx
@@ -26,6 +26,7 @@ import rehypeKatex from 'rehype-katex';
 import rehypeRaw from 'rehype-raw';
 import remarkGfm from 'remark-gfm';
 import remarkMath from 'remark-math';
+import { ImageChatMessage } from './ImageChatMessage';
 
 export interface Props {
   message: Message;
@@ -129,7 +130,11 @@ export const ChatMessage: FC<Props> = memo(
     const copyOnClick = () => {
       if (!navigator.clipboard) return;
 
-      navigator.clipboard.writeText(message.content).then(() => {
+      const copyText = selectedConversation?.model.type === 'text'
+        ? message.content
+        : JSON.parse(message.content)[0].url;
+
+      navigator.clipboard.writeText(copyText).then(() => {
         setMessageCopied(true);
         setTimeout(() => {
           setMessageCopied(false);
@@ -151,11 +156,10 @@ export const ChatMessage: FC<Props> = memo(
     return (
       <div
         className={`sm:px-4 lg:px-8
-        ${
-          message.role === 'assistant'
+        ${message.role === 'assistant'
             ? 'bg-gray-200 text-gray-800 dark:bg-[#444654] dark:text-gray-100'
             : 'bg-theme-light text-gray-800 dark:bg-theme-dark dark:text-gray-100'
-        }`}
+          }`}
         style={{ overflowWrap: 'anywhere' }}
       >
         <div className="group relative m-auto flex p-4 text-base md:max-w-2xl md:py-6 lg:max-w-full lg:px-0 2xl:max-w-5xl 3xl:max-w-6xl 4xl:max-w-7xl 5xl:max-w-[1920px] justify-center">
@@ -230,99 +234,94 @@ export const ChatMessage: FC<Props> = memo(
                 )}
               </div>
             )}
-            {message.role === 'assistant' && (
+            {message.role === 'assistant' && message.content && (
               <div className="flex flex-row">
-                <MemoizedReactMarkdown
-                  className="flex-1"
-                  remarkPlugins={[remarkGfm, remarkMath]}
-                  // @ts-ignore
-                  rehypePlugins={[rehypeRaw, rehypeKatex]}
-                  components={{
-                    code({ node, inline, className, children, ...props }) {
-                      if (children.length) {
-                        if (children[0] == '▍') {
+                {selectedConversation?.model.type === "text" ?
+                  <>
+                    <MemoizedReactMarkdown
+                      className="flex-1"
+                      remarkPlugins={[remarkGfm, remarkMath]}
+                      // @ts-ignore
+                      rehypePlugins={[rehypeRaw, rehypeKatex]}
+                      components={{
+                        code({ node, inline, className, children, ...props }) {
+                          if (children.length) {
+                            if (children[0] == '▍') {
+                              return (
+                                <span className="animate-pulse cursor-default mt-1">
+                                  ▍
+                                </span>
+                              );
+                            }
+
+                            children[0] = (children[0] as string).replace(
+                              '`▍`',
+                              '▍',
+                            );
+                          }
+
+                          const match = /language-(\w+)/.exec(className || '');
+
+                          return !inline ? (
+                            <CodeBlock
+                              key={Math.random()}
+                              language={(match && match[1]) || ''}
+                              value={String(children).replace(/\n$/, '')}
+                              {...props}
+                            />
+                          ) : (
+                            <code className={className} {...props}>
+                              {children}
+                            </code>
+                          );
+                        },
+                        table({ children }) {
+                          return (
+                            <table className="border-collapse border border-black px-3 py-1 dark:border-white">
+                              {children}
+                            </table>
+                          );
+                        },
+                        th({ children }) {
+                          return (
+                            <th className="break-words border border-black bg-gray-500 px-3 py-1 text-white dark:border-white">
+                              {children}
+                            </th>
+                          );
+                        },
+                        td({ children }) {
                           return (
-                            <span className="animate-pulse cursor-default mt-1">
-                              ▍
-                            </span>
+                            <td className="break-words border border-black px-3 py-1 dark:border-white">
+                              {children}
+                            </td>
                           );
-                        }
-
-                        children[0] = (children[0] as string).replace(
-                          '`▍`',
-                          '▍',
-                        );
-                      }
-
-                      const match = /language-(\w+)/.exec(className || '');
-
-                      return !inline ? (
-                        <CodeBlock
-                          key={Math.random()}
-                          language={(match && match[1]) || ''}
-                          value={String(children).replace(/\n$/, '')}
-                          {...props}
-                        />
-                      ) : (
-                        <code className={className} {...props}>
-                          {children}
-                        </code>
-                      );
-                    },
-                    table({ children }) {
-                      return (
-                        <table className="border-collapse border border-black px-3 py-1 dark:border-white">
-                          {children}
-                        </table>
-                      );
-                    },
-                    th({ children }) {
-                      return (
-                        <th className="break-words border border-black bg-gray-500 px-3 py-1 text-white dark:border-white">
-                          {children}
-                        </th>
-                      );
-                    },
-                    td({ children }) {
-                      return (
-                        <td className="break-words border border-black px-3 py-1 dark:border-white">
-                          {children}
-                        </td>
-                      );
-                    },
-                    a({ children, ...props }) {
-                      return (
-                        <a {...props} target="_blank">
-                          {children}
-                        </a>
-                      );
-                    },
-                    img({ src, alt, width, height }) {
-                      if (!width && !height) {
-                        width = '1024px';
-                        height = '1024px';
-                      }
-                      return (
-                        // eslint-disable-next-line @next/next/no-img-element
-                        <img
-                          src={src!}
-                          alt={alt!}
-                          width={parseInt(width as string)}
-                          height={parseInt(height as string)}
-                          className="m-1"
-                        />
-                      );
-                    },
-                  }}
-                >
-                  {`${message.content}${
-                    messageIsStreaming &&
-                    messageIndex ==
-                      (selectedConversationMessages.length ?? 0) - 1
-                      ? '`▍`'
-                      : ''
-                  }`}
-                </MemoizedReactMarkdown>
+                        },
+                        a({ children, ...props }) {
+                          return (
+                            <a {...props} target="_blank">
+                              {children}
+                            </a>
+                          );
+                        },
+                        img({ src, alt, width, height }) {
+                          return (
+                            <ImageChatMessage src={src} alt={alt} width={width} height={height} />
+                          );
+                        },
+                      }}
+                    >
+                      {`${message.content}${messageIsStreaming &&
+                        messageIndex ==
+                        (selectedConversationMessages.length ?? 0) - 1
+                        ? '`▍`'
+                        : ''
+                        }`}
+                    </MemoizedReactMarkdown>
+                  </> :
+                  <>
+                    <ImageChatMessage src={JSON.parse(message.content)[0].url} />
+                  </>
+                }
 
                 <div className="md:-mr-8 ml-1 md:ml-0 flex flex-col md:flex-row gap-4 md:gap-1 items-center md:items-start justify-end md:justify-start">
                   {messagedCopied ? (
@@ -343,7 +342,7 @@ export const ChatMessage: FC<Props> = memo(
             )}
           </div>
         </div>
-      </div>
+      </div >
     );
   },
 );

diff --git a/apps/unsaged/components/Home/components/ChatZone/Screens/Chat/ImageChatMessage.tsx b/apps/unsaged/components/Home/components/ChatZone/Screens/Chat/ImageChatMessage.tsx
@@ -0,0 +1,25 @@
+import { FC } from "react";
+
+interface Props {
+    src: string | undefined;
+    alt?: string | undefined;
+    width?: string | number | undefined;
+    height?: string | number | undefined;
+}
+
+export const ImageChatMessage: FC<Props> = ({ src, alt, width, height }) => {
+    if (!width && !height) {
+        width = '1024px';
+        height = '1024px';
+    }
+    return (
+        // eslint-disable-next-line @next/next/no-img-element
+        <img
+            src={src!}
+            alt={alt!}
+            width={parseInt(width as string)}
+            height={parseInt(height as string)}
+            className="m-1"
+        />
+    );
+};