Save chat message to KV store and UI revalidation (#14)

app/api/upload/route.ts

@@ -1,8 +1,8 @@
 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';
 
 /**
  * TODO: this should be replaced with actual upload to S3
@@ -11,7 +11,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 +37,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 +50,11 @@ 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);
 
     return Response.json(payload);
   } catch (error) {

components/chat-sidebar/ChatCard.tsx

@@ -46,7 +46,7 @@ const ChatCard: React.FC<ChatCardProps> = ({ chat }) => {
           className="rounded w-1/4 "
         />
         <p className="text-sm w-3/4 ml-3">
-          {firstMessage ? firstMessage + ' ...' : '(No messages yet)'}
+          {firstMessage ?? '(No messages yet)'}
         </p>
       </div>
     </ChatCardLayout>

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', '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', '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);
 }