Merge branch 'main' of github.com:landing-ai/vision-agent-ui

app/api/upload/route.ts

@@ -1,8 +1,9 @@
 import { auth } from '@/auth';
 import { upload } from '@/lib/aws';
+import { createKVChat } from '@/lib/kv/chat';
 import { ChatEntity, MessageBase } from '@/lib/types';
 import { nanoid } from '@/lib/utils';
-import { kv } from '@vercel/kv';
+import { revalidatePath } from 'next/cache';
 
 /**
  * TODO: this should be replaced with actual upload to S3
@@ -11,7 +12,7 @@ import { kv } from '@vercel/kv';
  */
 export async function POST(req: Request): Promise<Response> {
   const session = await auth();
-  const email = session?.user?.email;
+  const user = session?.user?.email ?? 'anonymous';
   // if (!email) {
   //   return new Response('Unauthorized', {
   //     status: 401,
@@ -37,7 +38,7 @@ export async function POST(req: Request): Promise<Response> {
 
     let urlToSave = url;
     if (base64) {
-      const fileName = `${email}/${id}/${Date.now()}-image.jpg`;
+      const fileName = `${user}/${id}/${Date.now()}-image.jpg`;
       const res = await upload(base64, fileName, fileType ?? 'image/png');
       if (res.ok) {
         console.log('Image uploaded successfully');
@@ -50,21 +51,13 @@ export async function POST(req: Request): Promise<Response> {
     const payload: ChatEntity = {
       url: urlToSave!, // TODO can be uploaded as well
       id,
-      user: email || 'anonymous',
+      user,
       messages: initMessages ?? [],
     };
 
-    await kv.hmset(`chat:${id}`, payload);
-    if (email) {
-      await kv.zadd(`user:chat:${email}`, {
-        score: Date.now(),
-        member: `chat:${id}`,
-      });
-    }
-    await kv.zadd('user:chat:all', {
-      score: Date.now(),
-      member: `chat:${id}`,
-    });
+    await createKVChat(payload);
+
+    revalidatePath('/chat', 'layout');
 
     return Response.json(payload);
   } catch (error) {

app/chat/[id]/page.tsx

@@ -1,6 +1,6 @@
-import { nanoid } from '@/lib/utils';
-import { Chat } from '@/components/chat';
-import { getKVChat } from '@/lib/kv/chat';
+import ChatDataLoad from '@/components/chat/ChatDataLoad';
+import { Suspense } from 'react';
+import Loading from '@/components/ui/Loading';
 
 interface PageProps {
   params: {
@@ -11,7 +11,9 @@ interface PageProps {
 export default async function Page({ params }: PageProps) {
   const { id: chatId } = params;
 
-  const chat = await getKVChat(chatId);
-
-  return <Chat chat={chat} />;
+  return (
+    <Suspense fallback={<Loading />}>
+      <ChatDataLoad chatId={chatId} />
+    </Suspense>
+  );
 }

app/globals.css

@@ -110,3 +110,54 @@
     pointer-events: none;
   }
 }
+
+/* Light theme. */
+:root {
+  --color-canvas-default: #ffffff;
+  --color-canvas-subtle: #f6f8fa;
+  --color-border-default: #d0d7de;
+  --color-border-muted: hsla(210, 18%, 87%, 1);
+}
+
+/* Dark theme. */
+@media (prefers-color-scheme: dark) {
+  :root {
+    --color-canvas-default: #0d1117;
+    --color-canvas-subtle: #161b22;
+    --color-border-default: #30363d;
+    --color-border-muted: #21262d;
+  }
+}
+
+table {
+  border-spacing: 0;
+  border-collapse: collapse;
+  display: block;
+  margin-top: 0;
+  margin-bottom: 16px;
+  width: max-content;
+  max-width: 100%;
+  overflow: auto;
+}
+
+tr {
+  border-top: 1px solid var(--color-border-muted);
+}
+
+tr:nth-child(2n) {
+  background-color: var(--color-canvas-subtle);
+}
+
+td,
+th {
+  padding: 6px 13px;
+  border: 1px solid var(--color-border-default);
+}
+
+th {
+  font-weight: 600;
+}
+
+table img {
+  background-color: transparent;
+}

components/chat-sidebar/ChatCard.tsx

@@ -18,7 +18,7 @@ export const ChatCardLayout: React.FC<
   return (
     <Link
       className={cn(
-        'p-2 m-2 bg-background l:h-[250px] rounded-xl shadow-md flex items-center border border-transparent hover:border-gray-500 transition-all cursor-pointer',
+        'p-2 m-2 bg-background max-h-[100px] rounded-xl shadow-md flex items-center border border-transparent hover:border-gray-500 transition-all cursor-pointer',
         classNames,
       )}
       href={link}
@@ -31,7 +31,12 @@ export const ChatCardLayout: React.FC<
 const ChatCard: React.FC<ChatCardProps> = ({ chat }) => {
   const { id: chatIdFromParam } = useParams();
   const { id, url, messages, user } = chat;
-  const firstMessage = messages?.[0]?.content.slice(0, 50);
+  const firstMessage = messages?.[0]?.content;
+  const title = firstMessage
+    ? firstMessage.length > 50
+      ? firstMessage.slice(0, 50) + '...'
+      : firstMessage
+    : '(No messages yet)';
   return (
     <ChatCardLayout
       link={`/chat/${id}`}
@@ -45,9 +50,9 @@ const ChatCard: React.FC<ChatCardProps> = ({ chat }) => {
           height={50}
           className="rounded w-1/4 "
         />
-        <p className="text-sm w-3/4 ml-3">
-          {firstMessage ? firstMessage + ' ...' : '(No messages yet)'}
-        </p>
+        <div className="flex items-start h-full">
+          <p className="text-sm w-3/4 ml-3">{title}</p>
+        </div>
       </div>
     </ChatCardLayout>
   );

components/chat/ChatDataLoad.tsx

@@ -0,0 +1,11 @@
+import { getKVChat } from '@/lib/kv/chat';
+import { Chat } from '.';
+
+export interface ChatDataLoadProps {
+  chatId: string;
+}
+
+export default async function ChatDataLoad({ chatId }: ChatDataLoadProps) {
+  const chat = await getKVChat(chatId);
+  return <Chat chat={chat} />;
+}

components/chat/ChatMessage.tsx

@@ -15,7 +15,67 @@ export interface ChatMessageProps {
   message: MessageBase;
 }
 
+const PAIRS: Record<string, string> = {
+  '┍': '┑',
+  '┝': '┥',
+  '├': '┤',
+  '┕': '┙',
+};
+
+const MIDDLE_STARTER = '┝';
+const MIDDLE_SEPARATOR = '┿';
+
 export function ChatMessage({ message, ...props }: ChatMessageProps) {
+  const cleanupMessage = ({ content, role }: MessageBase) => {
+    if (role === 'user') {
+      return {
+        content,
+      };
+    }
+    const [logs = '', answer = ''] = content.split('<ANSWER>');
+    const cleanedLogs = [];
+    let left = 0;
+    let right = 0;
+    while (right < logs.length) {
+      if (Object.keys(PAIRS).includes(content[right])) {
+        cleanedLogs.push(content.substring(left, right));
+        left = right++;
+        while (
+          right < content.length &&
+          PAIRS[content[left]] !== content[right]
+        ) {
+          right++;
+        }
+        if (content[left] === MIDDLE_STARTER) {
+          // add the text alignment so it can be shown as a table
+          const separators = logs
+            .substring(left, right)
+            .split(MIDDLE_SEPARATOR).length;
+          if (separators > 0) {
+            cleanedLogs.push(
+              Array(separators + 1)
+                .fill('|')
+                .join(' :- '),
+            );
+          }
+        }
+        left = ++right;
+      } else {
+        right++;
+      }
+    }
+    cleanedLogs.push(content.substring(left, right));
+    return {
+      logs: cleanedLogs
+        .join('')
+        .replace(/│/g, '|')
+        .split('|\n\n|')
+        .join('|\n|'),
+      content: answer.replace('</</ANSWER>', '').replace('</ANSWER>', ''),
+    };
+  };
+
+  const { logs, content } = cleanupMessage(message);
   return (
     <div className={cn('group relative mb-4 flex items-start')} {...props}>
       <div
@@ -29,12 +89,39 @@ export function ChatMessage({ message, ...props }: ChatMessageProps) {
         {message.role === 'user' ? <IconUser /> : <IconOpenAI />}
       </div>
       <div className="flex-1 px-1 ml-4 space-y-2 overflow-hidden">
+        {logs && message.role !== 'user' && (
+          <div className="bg-slate-100 dark:bg-slate-900 mb-4 p-4 max-h-[400px] overflow-auto">
+            <div className="text-xl font-bold">Thinking Process</div>
+            <MemoizedReactMarkdown
+              className="break-words text-sm"
+              remarkPlugins={[remarkGfm, remarkMath]}
+              components={{
+                p({ children }) {
+                  return (
+                    <p className="my-2 last:mb-0 whitespace-pre-line">
+                      {children}
+                    </p>
+                  );
+                },
+                code({ children, ...props }) {
+                  return (
+                    <code className="whitespace-pre-line">{children}</code>
+                  );
+                },
+              }}
+            >
+              {logs}
+            </MemoizedReactMarkdown>
+          </div>
+        )}
         <MemoizedReactMarkdown
           className="break-words"
           remarkPlugins={[remarkGfm, remarkMath]}
           components={{
             p({ children }) {
-              return <p className="mb-2 last:mb-0">{children}</p>;
+              return (
+                <p className="my-2 last:mb-0 whitespace-pre-line">{children}</p>
+              );
             },
             code({ node, inline, className, children, ...props }) {
               if (children.length) {
@@ -48,7 +135,6 @@ export function ChatMessage({ message, ...props }: ChatMessageProps) {
               }
 
               const match = /language-(\w+)/.exec(className || '');
-
               if (inline) {
                 return (
                   <code className={className} {...props}>
@@ -68,7 +154,7 @@ export function ChatMessage({ message, ...props }: ChatMessageProps) {
             },
           }}
         >
-          {message.content}
+          {content}
         </MemoizedReactMarkdown>
         <ChatMessageActions message={message} />
       </div>

components/chat/ImageSelector.tsx

@@ -51,13 +51,13 @@ const ImageSelector: React.FC<ImageSelectorProps> = () => {
       )}
     >
       <input {...getInputProps()} />
-      <p className="text-gray-400 text-lg">
+      <div className="text-gray-400 text-lg">
         {isUploading ? (
           <Loading />
         ) : (
           'Drag or drop image here, or click to select images'
         )}
-      </p>
+      </div>
     </div>
   );
 };

components/chat/PromptForm.tsx

@@ -54,8 +54,8 @@ export function PromptForm({
             <TooltipTrigger asChild>
               <Image
                 src={url}
-                width={60}
-                height={60}
+                width={250}
+                height={250}
                 alt="chosen image"
                 className="w-1/5 my-4 mx-2 rounded-md"
               />

components/ui/ImageLoader.tsx

@@ -0,0 +1,11 @@
+// import Image from 'next/image';
+// import React from 'react';
+
+// type ImageLoaderProps = typeof Image;
+
+// const ImageLoader: React.FC<ImageLoaderProps> = props => {
+//   const { alt, width, height } = props;
+//   return <Image alt={alt} />;
+// };
+
+// export default ImageLoader;

lib/hooks/useVisionAgent.tsx

@@ -1,13 +1,14 @@
-import { useChat, type Message } from 'ai/react';
+import { useChat, type Message, UseChatHelpers } from 'ai/react';
 import { toast } from 'react-hot-toast';
 import { useEffect, useState } from 'react';
 import { ChatEntity, MessageBase } from '../types';
+import { saveKVChatMessage } from '../kv/chat';
 
 const useVisionAgent = (chat: ChatEntity) => {
   const { messages: initialMessages, id, url } = chat;
   const {
     messages,
-    append,
+    append: appendRaw,
     reload,
     stop,
     isLoading,
@@ -22,6 +23,9 @@ const useVisionAgent = (chat: ChatEntity) => {
         toast.error(response.statusText);
       }
     },
+    onFinish(message) {
+      saveKVChatMessage(id, message);
+    },
     initialMessages: initialMessages,
     body: {
       url,
@@ -58,6 +62,16 @@ const useVisionAgent = (chat: ChatEntity) => {
     };
   }, [isLoading]);
 
+  useEffect(() => {
+    if (
+      !isLoading &&
+      messages.length &&
+      messages[messages.length - 1].role === 'user'
+    ) {
+      reload();
+    }
+  }, [isLoading, messages, reload]);
+
   const assistantLoadingMessage = {
     id: 'loading',
     content: loadingDots,
@@ -71,6 +85,11 @@ const useVisionAgent = (chat: ChatEntity) => {
       ? [...messages, assistantLoadingMessage]
       : messages;
 
+  const append: UseChatHelpers['append'] = async message => {
+    await saveKVChatMessage(id, message as MessageBase);
+    return appendRaw(message);
+  };
+
   return {
     messages: messageWithLoading as MessageBase[],
     append,

lib/kv/chat.ts

@@ -4,16 +4,17 @@ import { revalidatePath } from 'next/cache';
 import { kv } from '@vercel/kv';
 
 import { auth } from '@/auth';
-import { ChatEntity } from '@/lib/types';
-import { redirect } from 'next/navigation';
+import { ChatEntity, MessageBase } from '@/lib/types';
+import { notFound, redirect } from 'next/navigation';
+import { nanoid } from '../utils';
 
 async function authCheck() {
   const session = await auth();
   const email = session?.user?.email;
-  if (!email) {
-    redirect('/');
-  }
-  return { email, isAdmin: email.endsWith('landing.ai') };
+  // if (!email) {
+  //   redirect('/');
+  // }
+  return { email, isAdmin: !!email?.endsWith('landing.ai') };
 }
 
 export async function getKVChats() {
@@ -48,6 +49,37 @@ export async function getKVChat(id: string) {
   return chat;
 }
 
+export async function createKVChat(chat: ChatEntity) {
+  // const { email, isAdmin } = await authCheck();
+  const { email } = await authCheck();
+
+  await kv.hmset(`chat:${chat.id}`, chat);
+  if (email) {
+    await kv.zadd(`user:chat:${email}`, {
+      score: Date.now(),
+      member: `chat:${chat.id}`,
+    });
+  }
+  await kv.zadd('user:chat:all', {
+    score: Date.now(),
+    member: `chat:${chat.id}`,
+  });
+  revalidatePath('/chat', 'layout');
+}
+
+export async function saveKVChatMessage(id: string, message: MessageBase) {
+  const chat = await kv.hgetall<ChatEntity>(`chat:${id}`);
+  if (!chat) {
+    notFound();
+  }
+  const { messages } = chat;
+  await kv.hmset(`chat:${id}`, {
+    ...chat,
+    messages: [...messages, message],
+  });
+  revalidatePath('/chat', 'layout');
+}
+
 export async function removeKVChat({ id, path }: { id: string; path: string }) {
   const session = await auth();
 
@@ -69,6 +101,6 @@ export async function removeKVChat({ id, path }: { id: string; path: string }) {
   await kv.del(`chat:${id}`);
   await kv.zrem(`user:chat:${session.user.id}`, `chat:${id}`);
 
-  revalidatePath('/');
+  revalidatePath('/chat/layout', 'layout');
   return revalidatePath(path);
 }