feat: DB change - combine user and assistant message (#70)

app/api/chat/route.ts

@@ -6,14 +6,14 @@ import {
   ChatCompletionMessageParam,
   ChatCompletionContentPart,
 } from 'openai/resources';
-import { MessageBase } from '../../../lib/types';
+import { MessageUI } from '@/lib/types';
 
 export const runtime = 'edge';
 
 export const POST = async (req: Request) => {
   const json = await req.json();
   const { messages } = json as {
-    messages: MessageBase[];
+    messages: MessageUI[];
     id: string;
     url: string;
   };

app/api/vision-agent/route.ts

@@ -1,10 +1,10 @@
 import { StreamingTextResponse } from 'ai';
 
 // import { auth } from '@/auth';
-import { MessageBase, SignedPayload } from '../../../lib/types';
+import { MessageUI, SignedPayload } from '@/lib/types';
 import { logger, withLogging } from '@/lib/logger';
 import { CLEANED_SEPARATOR } from '@/lib/constants';
-import { cleanAnswerMessage, cleanInputMessage } from '@/lib/messageUtils';
+import { cleanAnswerMessage, cleanInputMessage } from '@/lib/utils/content';
 import { fetcher } from '@/lib/utils';
 
 // export const runtime = 'edge';
@@ -53,7 +53,7 @@ export const POST = withLogging(
   async (
     session,
     json: {
-      messages: MessageBase[];
+      messages: MessageUI[];
       id: string;
       mediaUrl: string;
     },

app/chat/layout.tsx

@@ -1,38 +1,7 @@
-// import { sessionUser } from '@/auth';
-// import ChatSidebarList from '@/components/chat-sidebar/ChatListSidebar';
-// import Loading from '@/components/ui/Loading';
-// import { dbGetMyChatListWithMessages } from '@/lib/db/functions';
-// import { Suspense } from 'react';
-
 interface ChatLayoutProps {
   children: React.ReactNode;
 }
 
-// export default async function Layout({ children }: ChatLayoutProps) {
-//   const { email, user, id } = await sessionUser();
-//   const chats = await dbGetMyChatListWithMessages();
-
-//   return (
-//     <div className="relative flex h-[calc(100vh_-_theme(spacing.16))] overflow-hidden">
-//       {user && (
-//         <div
-//           data-state={email ? 'open' : 'closed'}
-//           className="peer absolute inset-y-0 z-30 hidden border-r bg-muted duration-300 ease-in-out -translate-x-full data-[state=open]:translate-x-0 lg:flex lg:w-[250px] h-full flex-col overflow-auto py-2"
-//         >
-//           <Suspense fallback={<Loading />}>
-//             <ChatSidebarList chats={chats} />
-//           </Suspense>
-//         </div>
-//       )}
-//       <Suspense fallback={<Loading />}>
-//         <div className="group w-full overflow-auto pl-0 animate-in duration-300 ease-in-out peer-[[data-state=open]]:lg:pl-[250px]">
-//           {children}
-//         </div>
-//       </Suspense>
-//     </div>
-//   );
-// }
-
 export default async function Layout({ children }: ChatLayoutProps) {
   // return <Suspense fallback={<Loading />}>{children}</Suspense>;
   return children;

app/page.tsx

@@ -1,15 +1,13 @@
 'use client';
 
-import { generateInputImageMarkdown } from '@/lib/messageUtils';
 import { useRouter } from 'next/navigation';
 
-import { MessageRaw } from '@/lib/db/types';
 import { useRef, useState } from 'react';
 import Composer, { ComposerRef } from '@/components/chat/Composer';
 import { dbPostCreateChat } from '@/lib/db/functions';
 import { nanoid } from '@/lib/utils';
 import Chip from '@/components/ui/Chip';
-import { IconArrowUpRight, IconImage } from '@/components/ui/Icons';
+import { IconArrowUpRight } from '@/components/ui/Icons';
 
 const EXAMPLES = [
   {
@@ -18,12 +16,6 @@ const EXAMPLES = [
       'https://vision-agent-dev.s3.us-east-2.amazonaws.com/examples/flower.png',
     prompt: 'Count the number of flowers in this image.',
   },
-  // {
-  //   heading: 'Detecting',
-  //   url: 'https://landing-lens-support.s3.us-east-2.amazonaws.com/vision-agent-examples/cereal-example.jpg',
-  //   subheading: 'number of cereals in an image',
-  //   message: `How many cereals are there in the image?`,
-  // },
 ];
 
 export default function Page() {
@@ -46,19 +38,12 @@ export default function Page() {
               const newId = nanoid();
               const resp = await dbPostCreateChat({
                 id: newId,
-                mediaUrl: mediaUrl,
                 title: `conversation-${newId}`,
-                initMessages: [
-                  {
-                    role: 'user',
-                    content:
-                      input +
-                      (mediaUrl
-                        ? '\n\n' + generateInputImageMarkdown(mediaUrl)
-                        : ''),
-                    result: null,
-                  },
-                ],
+                mediaUrl,
+                message: {
+                  prompt: input,
+                  mediaUrl,
+                },
               });
               if (resp) {
                 router.push(`/chat/${newId}`);
@@ -66,25 +51,23 @@ export default function Page() {
             }}
           />
         </div>
-        <>
-          {EXAMPLES.map((example, index) => {
-            return (
-              <Chip
-                key={index}
-                className="bg-transparent border border-zinc-500 cursor-pointer px-4"
-                onClick={() => {
-                  composerRef.current?.setInput(example.prompt);
-                  composerRef.current?.setMediaUrl(example.mediaUrl);
-                }}
-              >
-                <div className="flex flex-row items-center space-x-2">
-                  <p className="text-primary text-sm">{example.title}</p>
-                  <IconArrowUpRight className="text-primary" />
-                </div>
-              </Chip>
-            );
-          })}
-        </>
+        {EXAMPLES.map((example, index) => {
+          return (
+            <Chip
+              key={index}
+              className="bg-transparent border border-zinc-500 cursor-pointer px-4"
+              onClick={() => {
+                composerRef.current?.setInput(example.prompt);
+                composerRef.current?.setMediaUrl(example.mediaUrl);
+              }}
+            >
+              <div className="flex flex-row items-center space-x-2">
+                <p className="text-primary text-sm">{example.title}</p>
+                <IconArrowUpRight className="text-primary" />
+              </div>
+            </Chip>
+          );
+        })}
       </div>
     </div>
   );

auth.ts

@@ -46,19 +46,21 @@ export const {
       return false;
     },
     async jwt({ token, profile, user }) {
-      // console.log('[Ming] ~ jwt ~ user:', user, token);
-      // const dbUser = await dbFindOrCreateUser(email, name);
-      // console.log('[Ming] ~ signIn ~ dbUser:', dbUser);
       if (profile) {
         token.id = profile.id || profile.sub;
         token.image = profile.avatar_url || profile.picture;
       }
       return token;
     },
-    session: ({ session, token }) => {
-      if (token) {
+    async session({ session, token }) {
+      // TODO: this is temporary between we switch DB and make migration
+      // so also UI might still have session, DB might already have cleaned up
+      const email = session?.user?.email;
+      const name = session?.user?.name;
+      if (email && name) {
+        const dbUser = await dbFindOrCreateUser(email, name);
         // put db user id into session
-        session.user.id = token.sub ?? '';
+        session.user.id = dbUser.id;
       }
       return session;
     },

components/chat-sidebar/ChatAdminToggle.tsx

@@ -1,14 +0,0 @@
-'use client';
-
-import { chatViewMode } from '@/state/chat';
-import { useAtom } from 'jotai';
-import React from 'react';
-
-export interface ChatAdminToggleProps {}
-
-const ChatAdminToggle: React.FC<ChatAdminToggleProps> = () => {
-  const modeAtom = useAtom(chatViewMode);
-  return null;
-};
-
-export default ChatAdminToggle;

components/chat-sidebar/ChatCard.tsx

@@ -1,75 +0,0 @@
-'use client';
-
-import { PropsWithChildren } from 'react';
-import Link from 'next/link';
-import { useParams, usePathname, useRouter } from 'next/navigation';
-import { cn } from '@/lib/utils';
-import Image from 'next/image';
-import clsx from 'clsx';
-import Img from '../ui/Img';
-import { format } from 'date-fns';
-import { cleanInputMessage } from '@/lib/messageUtils';
-import { IconClose } from '../ui/Icons';
-import { ChatWithMessages } from '@/lib/db/types';
-import { dbDeleteChat } from '@/lib/db/functions';
-
-type ChatCardProps = PropsWithChildren<{
-  chat: ChatWithMessages;
-  isAdminView?: boolean;
-}>;
-
-export const ChatCardLayout: React.FC<
-  PropsWithChildren<{ link: string; classNames?: clsx.ClassValue }>
-> = ({ link, children, classNames }) => {
-  return (
-    <Link
-      className={cn(
-        'p-2 bg-background max-h-[100px] rounded-xl shadow-md flex items-center border border-transparent hover:border-gray-500 transition-all cursor-pointer w-full',
-        classNames,
-      )}
-      href={link}
-    >
-      {children}
-    </Link>
-  );
-};
-
-const ChatCard: React.FC<ChatCardProps> = ({ chat, isAdminView }) => {
-  const { id: chatIdFromParam } = useParams();
-  const { id, mediaUrl, messages, userId, updatedAt } = chat;
-  if (!id) {
-    return null;
-  }
-  const firstMessage = cleanInputMessage(messages?.[0]?.content ?? '');
-  const title = firstMessage
-    ? firstMessage.length > 50
-      ? firstMessage.slice(0, 50) + '...'
-      : firstMessage
-    : '(No messages yet)';
-  return (
-    <ChatCardLayout
-      link={isAdminView ? `/all/chat/${id}` : `/chat/${id}`}
-      classNames={chatIdFromParam === id && 'border-gray-500'}
-    >
-      <div className="overflow-hidden flex items-center size-full group">
-        <Img src={mediaUrl} alt={`chat-${id}-card-image`} className="w-1/4" />
-        <div className="flex items-start flex-col h-full ml-3 w-3/4">
-          <p className="text-sm mb-1">{title}</p>
-          <p className="text-xs text-gray-500">
-            {updatedAt ? format(Number(updatedAt), 'yyyy-MM-dd') : '-'}
-          </p>
-          {isAdminView && <p className="text-xs text-gray-500">{userId}</p>}
-          <IconClose
-            onClick={async e => {
-              e.stopPropagation();
-              await dbDeleteChat(id);
-            }}
-            className="absolute right-4 opacity-0 group-hover:opacity-100 top-1/2 -translate-y-1/2"
-          />
-        </div>
-      </div>
-    </ChatCardLayout>
-  );
-};
-
-export default ChatCard;

components/chat-sidebar/ChatListSidebar.tsx

@@ -1,69 +0,0 @@
-'use client';
-
-import ChatCard, { ChatCardLayout } from './ChatCard';
-import { IconPlus } from '../ui/Icons';
-import { auth } from '@/auth';
-import { VariableSizeList as List } from 'react-window';
-import { cleanInputMessage } from '@/lib/messageUtils';
-import AutoSizer from 'react-virtualized-auto-sizer';
-import { ChatWithMessages } from '@/lib/db/types';
-
-export interface ChatSidebarListProps {
-  chats: ChatWithMessages[];
-  isAdminView?: boolean;
-}
-
-const getItemSize = (message: string, isAdminView?: boolean) => {
-  if (message.length >= 45) return 116;
-  else if (message.length >= 20) return 104;
-  else return 88;
-};
-
-export default function ChatSidebarList({
-  chats,
-  isAdminView,
-}: ChatSidebarListProps) {
-  return (
-    <>
-      {!isAdminView && (
-        <div className="p-2">
-          <ChatCardLayout link="/chat">
-            <div className="overflow-hidden flex items-center size-full">
-              <IconPlus className="w-1/4 font-bold" />
-              <p className="text-sm w-3/4 ml-3 font-bold">New chat</p>
-            </div>
-          </ChatCardLayout>
-        </div>
-      )}
-      <AutoSizer>
-        {({ height, width }) => (
-          <List
-            itemData={chats}
-            height={height}
-            itemCount={chats.length}
-            itemSize={index =>
-              getItemSize(
-                cleanInputMessage(chats[index].messages?.[0]?.content ?? ''),
-                isAdminView,
-              )
-            }
-            width={width}
-          >
-            {({ style, index, data }) => (
-              <div
-                style={style}
-                className="px-2 flex items-center overflow-hidden"
-              >
-                <ChatCard
-                  key={data[index].id}
-                  chat={data[index]}
-                  isAdminView={isAdminView}
-                />
-              </div>
-            )}
-          </List>
-        )}
-      </AutoSizer>
-    </>
-  );
-}

components/chat/ChatClient.tsx

@@ -6,12 +6,12 @@ import useVisionAgent from '@/lib/hooks/useVisionAgent';
 import { useScrollAnchor } from '@/lib/hooks/useScrollAnchor';
 import { Session } from 'next-auth';
 import { useEffect } from 'react';
-import { ChatWithMessages } from '@/lib/db/types';
+import { ChatWithMessages, MessageUserInput } from '@/lib/types';
 import { ChatMessage } from './ChatMessage';
 import { Button } from '../ui/Button';
 import { cn } from '@/lib/utils';
 import { IconArrowDown } from '../ui/Icons';
-import { generateInputImageMarkdown } from '@/lib/messageUtils';
+import { dbPostCreateMessage } from '@/lib/db/functions';
 
 export interface ChatClientProps {
   chat: ChatWithMessages;
@@ -20,8 +20,8 @@ export interface ChatClientProps {
 export const SCROLL_BOTTOM = 120;
 
 const ChatClient: React.FC<ChatClientProps> = ({ chat }) => {
-  const { mediaUrl, id } = chat;
-  const { messages, append, isLoading, reload } = useVisionAgent(chat);
+  const { id, messages: dbMessages } = chat;
+  const { messages, append, isLoading } = useVisionAgent(chat);
 
   const { messagesRef, scrollRef, visibilityRef, isVisible, scrollToBottom } =
     useScrollAnchor(SCROLL_BOTTOM);
@@ -56,17 +56,13 @@ const ChatClient: React.FC<ChatClientProps> = ({ chat }) => {
       </div>
       <div className="absolute bottom-4 w-full">
         <Composer
-          initMediaUrl={mediaUrl}
+          // Use the last message mediaUrl as the initial mediaUrl
+          initMediaUrl={dbMessages[dbMessages.length - 1]?.mediaUrl}
           isLoading={isLoading}
           onSubmit={async ({ input, mediaUrl: newMediaUrl }) => {
             append({
-              id,
-              content:
-                input +
-                (newMediaUrl
-                  ? '\n\n' + generateInputImageMarkdown(newMediaUrl)
-                  : ''),
-              role: 'user',
+              prompt: input,
+              mediaUrl: newMediaUrl,
             });
           }}
         />

components/chat/ChatList.tsx

@@ -2,13 +2,13 @@
 
 import { Separator } from '@/components/ui/Separator';
 import { ChatMessage } from '@/components/chat/ChatMessage';
-import { MessageBase } from '../../lib/types';
+import { MessageUI } from '@/lib/types';
 import { Session } from 'next-auth';
 import { IconExclamationTriangle } from '../ui/Icons';
 import Link from 'next/link';
 
 export interface ChatList {
-  messages: MessageBase[];
+  messages: MessageUI[];
   session: Session | null;
   isLoading: boolean;
 }

components/chat/ChatMessage.tsx

@@ -21,9 +21,9 @@ import {
   IconOutput,
   IconLog,
 } from '@/components/ui/Icons';
-import { MessageBase } from '../../lib/types';
+import { MessageUI } from '@/lib/types';
 import Img from '../ui/Img';
-import { ChunkBody, CodeResult, formatStreamLogs } from '@/lib/messageUtils';
+import { ChunkBody, CodeResult, formatStreamLogs } from '@/lib/utils/content';
 import {
   Table,
   TableBody,
@@ -43,7 +43,7 @@ import {
 import { Dialog, DialogContent, DialogTrigger } from '../ui/Dialog';
 
 export interface ChatMessageProps {
-  message: MessageBase;
+  message: MessageUI;
   isLoading: boolean;
 }
 
@@ -194,7 +194,7 @@ const ChunkPayloadAction: React.FC<{
           </Button>
         </DialogTrigger>
         <DialogContent className="max-w-5xl">
-          <Table className="border rounded-lg bg-zinc-700 overflow-hidden">
+          <Table>
             <TableHeader>
               <TableRow className="border-primary/50">
                 {keyArray.map(header => (
@@ -331,7 +331,7 @@ const AssistantChatMessage: React.FC<{
         <IconLandingAI />
       </div>
       <div className="flex-1 px-1 space-y-4 ml-4 overflow-hidden">
-        <Table className="border rounded-lg bg-zinc-700 overflow-hidden w-[400px]">
+        <Table className="w-[400px]">
           <TableBody>
             {formattedSections.map(section => (
               <TableRow className="border-primary/50" key={section.type}>

components/chat/ChatMessageActions.tsx

@@ -6,10 +6,10 @@ import { Button } from '@/components/ui/Button';
 import { IconCheck, IconCopy } from '@/components/ui/Icons';
 import { useCopyToClipboard } from '@/lib/hooks/useCopyToClipboard';
 import { cn } from '@/lib/utils';
-import { MessageBase } from '../../lib/types';
+import { MessageUI } from '../../lib/types';
 
 interface ChatMessageActionsProps extends React.ComponentProps<'div'> {
-  message: MessageBase;
+  message: MessageUI;
 }
 
 export function ChatMessageActions({

components/ui/Table.tsx

@@ -9,7 +9,10 @@ const Table = React.forwardRef<
   <div className="relative w-full overflow-auto">
     <table
       ref={ref}
-      className={cn('w-full caption-bottom text-sm', className)}
+      className={cn(
+        'w-full caption-bottom text-sm border rounded-lg bg-zinc-700 overflow-hidden',
+        className,
+      )}
       {...props}
     />
   </div>
@@ -87,7 +90,7 @@ const TableCell = React.forwardRef<
 >(({ className, ...props }, ref) => (
   <td
     ref={ref}
-    className={cn('p-4 align-middle [&:has([role=checkbox])]:pr-0', className)}
+    className={cn('p-2 align-middle [&:has([role=checkbox])]:pr-0', className)}
     {...props}
   />
 ));

lib/db/functions.ts

@@ -2,10 +2,25 @@
 
 import { sessionUser } from '@/auth';
 import prisma from './prisma';
-import { ChatWithMessages, MessageRaw } from './types';
+import {
+  ChatWithMessages,
+  MessageAssistantResponse,
+  MessageUserInput,
+} from '../types';
 import { revalidatePath } from 'next/cache';
 import { Chat } from '@prisma/client';
 
+async function getUserConnect() {
+  const { id: userId } = await sessionUser();
+  return userId
+    ? {
+        user: {
+          connect: { id: userId }, // Connect the chat to an existing user
+        },
+      }
+    : {};
+}
+
 /**
  * Finds or creates a user in the database based on the provided email and name.
  * If the user already exists, it returns the existing user.
@@ -48,25 +63,6 @@ export async function dbGetMyChatList(): Promise<Chat[]> {
   });
 }
 
-/**
- * Retrieves all chats with their associated messages for the current user.
- * @returns A promise that resolves to an array of `ChatWithMessages` objects.
- */
-export async function dbGetMyChatListWithMessages(): Promise<
-  ChatWithMessages[]
-> {
-  const { id: userId } = await sessionUser();
-
-  if (!userId) return [];
-
-  return prisma.chat.findMany({
-    where: { userId },
-    include: {
-      messages: true,
-    },
-  });
-}
-
 /**
  * Retrieves a chat with messages from the database based on the provided ID.
  * @param id - The ID of the chat to retrieve.
@@ -82,45 +78,34 @@ export async function dbGetChat(id: string): Promise<ChatWithMessages | null> {
 }
 
 /**
- * Creates a new chat in the database.
- *
+ * Creates a new chat in the database with the provided information.
+ * @param {Object} options - The options for creating the chat.
  * @param {string} options.id - The ID of the chat (optional).
- * @param {string} options.mediaUrl - The media URL for the chat.
- * @param {MessageRaw[]} options.initMessages - The initial messages for the chat (optional).
- * @returns {Promise<Chat>} The created chat object.
+ * @param {string} options.title - The title of the chat (optional).
+ * @param {MessageUserInput} options.message - The message to be added to the chat.
+ * @returns {Promise<Chat>} - A promise that resolves to the created chat.
  */
 export async function dbPostCreateChat({
   id,
-  mediaUrl,
   title,
-  initMessages = [],
+  mediaUrl,
+  message,
 }: {
   id?: string;
-  mediaUrl: string;
   title?: string;
-  initMessages?: MessageRaw[];
+  mediaUrl: string;
+  message: MessageUserInput;
 }) {
-  const { id: userId } = await sessionUser();
-  const userConnect = userId
-    ? {
-        user: {
-          connect: { id: userId }, // Connect the chat to an existing user
-        },
-      }
-    : {};
+  const userConnect = await getUserConnect();
   try {
     const response = await prisma.chat.create({
       data: {
         id,
-        mediaUrl: mediaUrl,
         ...userConnect,
+        mediaUrl,
         title,
         messages: {
-          create: initMessages.map(message => ({
-            ...message,
-            ...userConnect,
-            result: undefined,
-          })),
+          create: [{ ...message, ...userConnect }],
         },
       },
       include: {
@@ -136,29 +121,46 @@ export async function dbPostCreateChat({
 }
 
 /**
- * Creates a new message in the database.
- * @param chatId - The ID of the chat where the message belongs.
- * @param message - The message object to be created.
- * @returns A promise that resolves to the created message.
+ * Creates a new message in the database for a given chat.
+ * @param chatId - The ID of the chat where the message will be created.
+ * @param message - The message to be created.
  */
-export async function dbPostCreateMessage(chatId: string, message: MessageRaw) {
-  const { id: userId } = await sessionUser();
-  const userConnect = userId
-    ? {
-        user: {
-          connect: { id: userId }, // Connect the chat to an existing user
-        },
-      }
-    : {};
+export async function dbPostCreateMessage(
+  chatId: string,
+  message: MessageUserInput,
+) {
+  const userConnect = await getUserConnect();
 
-  await prisma.message.create({
+  const resp = await prisma.message.create({
     data: {
-      content: message.content,
-      role: message.role,
+      ...message,
+      ...userConnect,
       chat: {
         connect: { id: chatId },
       },
-      ...userConnect,
+    },
+  });
+
+  revalidatePath('/chat');
+  return resp;
+}
+
+/**
+ * Updates a message in the database with the provided response.
+ * @param messageId - The ID of the message to update.
+ * @param messageResponse - The response to update the message with.
+ */
+export async function dbPostUpdateMessageResponse(
+  messageId: string,
+  messageResponse: MessageAssistantResponse,
+) {
+  await prisma.message.update({
+    data: {
+      response: messageResponse.response,
+      result: messageResponse.result ?? undefined,
+    },
+    where: {
+      id: messageId,
     },
   });
 

lib/db/types.ts

@@ -1,10 +0,0 @@
-import { Chat, Message } from '@prisma/client';
-
-export type ChatWithMessages = Chat & { messages: Message[] };
-export type ChatMessage = Message;
-
-export type MessageRaw = {
-  role: Message['role'];
-  result: Message['result'];
-  content: Message['content'];
-};

lib/hooks/useVisionAgent.ts

@@ -1,24 +1,26 @@
-import { useChat, UseChatHelpers } from 'ai/react';
+import { useChat } from 'ai/react';
 import { toast } from 'react-hot-toast';
-import { useEffect, useRef } from 'react';
-import { ChatWithMessages } from '../db/types';
-import { dbPostCreateMessage } from '../db/functions';
+import { useEffect, useRef, useState } from 'react';
+import { ChatWithMessages, MessageUserInput } from '../types';
 import {
-  convertDbMessageToMessage,
-  convertMessageToDbMessage,
-} from '../messageUtils';
+  dbPostCreateMessage,
+  dbPostUpdateMessageResponse,
+} from '../db/functions';
+import {
+  convertAssistantUIMessageToDBMessageResponse,
+  convertDBMessageToUIMessage,
+} from '../utils/message';
 
 const useVisionAgent = (chat: ChatWithMessages) => {
-  const { messages: initialMessages, id, mediaUrl } = chat;
+  const { messages: dbMessages, id, mediaUrl } = chat;
+  const latestDbMessage = dbMessages[dbMessages.length - 1];
 
-  const {
-    messages,
-    append: appendRaw,
-    isLoading,
-    reload,
-  } = useChat({
+  // Temporary solution for now while single we have to pass mediaUrl separately outside of the messages
+  const currMediaUrl = useRef<string>(mediaUrl);
+  const currMessageId = useRef<string>(latestDbMessage?.id);
+
+  const { messages, append, isLoading, reload } = useChat({
     api: '/api/vision-agent',
-    // @ts-ignore https://sdk.vercel.ai/docs/troubleshooting/common-issues/use-chat-failed-to-parse-stream
     streamMode: 'text',
     onResponse(response) {
       if (response.status !== 200) {
@@ -26,11 +28,15 @@ const useVisionAgent = (chat: ChatWithMessages) => {
       }
     },
     onFinish: async message => {
-      await dbPostCreateMessage(id, convertMessageToDbMessage(message));
+      await dbPostUpdateMessageResponse(
+        currMessageId.current,
+        convertAssistantUIMessageToDBMessageResponse(message),
+      );
     },
-    initialMessages: initialMessages.map(convertDbMessageToMessage),
+    sendExtraMessageFields: true,
+    initialMessages: convertDBMessageToUIMessage(dbMessages),
     body: {
-      mediaUrl,
+      mediaUrl: currMediaUrl.current,
       id,
     },
     onError: err => {
@@ -54,18 +60,18 @@ const useVisionAgent = (chat: ChatWithMessages) => {
     }
   }, [isLoading, messages, reload]);
 
-  const append: UseChatHelpers['append'] = async message => {
-    dbPostCreateMessage(id, {
-      role: message.role as 'user' | 'assistant',
-      content: message.content,
-      result: null,
-    });
-    return appendRaw(message);
-  };
-
   return {
     messages,
-    append,
+    append: async (messageInput: MessageUserInput) => {
+      currMediaUrl.current = messageInput.mediaUrl;
+      append({
+        id,
+        role: 'user',
+        content: messageInput.prompt,
+      });
+      const resp = await dbPostCreateMessage(id, messageInput);
+      currMessageId.current = resp.id;
+    },
     reload,
     isLoading,
   };

lib/types.ts

@@ -1,10 +1,14 @@
-import { type Message } from 'ai';
+import { Chat, Message } from '@prisma/client';
+import { type Message as MessageAI } from 'ai';
 
-export type MessageBase = {
-  role: Message['role'];
-  content: string;
-  id: string;
-};
+export type ChatWithMessages = Chat & { messages: Message[] };
+
+export type MessageUserInput = Pick<Message, 'prompt' | 'mediaUrl'>;
+export type MessageAssistantResponse = Partial<
+  Pick<Message, 'response' | 'result'>
+>;
+
+export type MessageUI = Pick<MessageAI, 'role' | 'content' | 'id'>;
 
 export interface SignedPayload {
   id: string;

lib/{messageUtils.ts → utils/content.ts}

@@ -1,5 +1,4 @@
 import toast from 'react-hot-toast';
-import type { ChatMessage } from '@/lib/db/types';
 import { Message } from 'ai';
 
 const PAIRS: Record<string, string> = {
@@ -13,22 +12,11 @@ const MIDDLE_STARTER = '┝';
 const MIDDLE_SEPARATOR = '┿';
 
 const ANSWERS_PREFIX = 'answers';
-const INPUT_PREFIX = 'input';
 
 export const generateAnswersImageMarkdown = (index: number, url: string) => {
   return `![${ANSWERS_PREFIX}-${index}](${url})`;
 };
 
-export const generateInputImageMarkdown = (url: string, index = 0) => {
-  if (url.toLowerCase().endsWith('.mp4')) {
-    const prefix = 'input-video';
-    return `![${INPUT_PREFIX}-${index}](<${url}>)`;
-  } else {
-    const prefix = 'input';
-    return `![${INPUT_PREFIX}-${index}](<${url}>)`;
-  }
-};
-
 export const cleanInputMessage = (content: string) => {
   return content
     .replace(/!\[input-.*?\)/g, '')
@@ -244,100 +232,21 @@ const parseLine = (json: MessageBody) => {
   }
 };
 
-export const getFormattedMessage = ({
-  content,
-  role,
-}: Pick<ChatMessage, 'role' | 'content'>) => {
-  if (role === 'user') {
-    return {
-      content,
-    };
-  }
-  const lines = content.split('\n');
-  let formattedLogs = '';
-  let finalResult = null;
-  const jsons: MessageBody[] = [];
-  for (let line of lines) {
-    console.log(line);
-    if (!line.trim()) {
-      continue;
-    }
-    try {
-      const json = JSON.parse(line) as MessageBody;
-      if (json.type === 'final_code') {
-        const result = JSON.parse(
-          json.payload.result as unknown as string,
-        ) as PrismaJson.FinalChatResult['payload']['result'];
-        finalResult = generateFinalCodeMarkdown(
-          json.payload.code,
-          json.payload.test,
-          result,
-        );
-      } else if (
-        jsons.length > 0 &&
-        json.type === jsons[jsons.length - 1].type &&
-        json.status !== 'started'
-      ) {
-        jsons[jsons.length - 1] = json;
-      } else {
-        jsons.push(json);
-      }
-    } catch (e) {
-      console.error((e as Error).message);
-      console.error(line);
-    }
-  }
-  jsons.forEach(json => (formattedLogs += parseLine(json)));
-  return {
-    content: formattedLogs + (finalResult ?? ''),
-  };
-};
-
-export const convertDbMessageToMessage = (message: ChatMessage): Message => {
-  return {
-    id: message.id,
-    role: message.role,
-    createdAt: message.createdAt,
-    content: message.result
-      ? message.content + '\n' + JSON.stringify(message.result)
-      : message.content,
-  };
-};
-
-export const convertMessageToDbMessage = (
-  message: Message,
-): Pick<ChatMessage, 'content' | 'result' | 'role'> => {
-  let result = null;
-  const lines = message.content.split('\n');
-  for (let line of lines) {
-    try {
-      const json = JSON.parse(line) as MessageBody;
-      if (json.type == 'final_code') {
-        result = json as PrismaJson.FinalChatResult;
-        break;
-      }
-    } catch (e) {
-      console.error((e as Error).message);
-    }
-  }
-  return {
-    role: message.role as ChatMessage['role'],
-    content: message.content,
-    result,
-  };
-};
-
 export type CodeResult = {
   code: string;
   test: string;
   result: string;
 };
 
-export type ChunkBody = {
-  type: 'plans' | 'tools' | 'code' | 'final_code';
-  status: 'started' | 'completed' | 'failed' | 'running';
-  payload: Array<Record<string, string>> | CodeResult;
-};
+export type ChunkBody =
+  | {
+      type: 'plans' | 'tools' | 'code' | 'final_code';
+      status: 'started' | 'completed' | 'failed' | 'running';
+      payload:
+        | Array<Record<string, string>> // PlansBody | ToolsBody
+        | CodeResult; // CodeBody
+    }
+  | PrismaJson.FinalChatResult;
 
 /**
  * Formats the stream logs and returns an array of grouped sections.

lib/utils/message.ts

@@ -0,0 +1,63 @@
+import { Message } from '@prisma/client';
+import { MessageAssistantResponse, MessageUI } from '../types';
+import { ChunkBody } from './content';
+
+const INPUT_PREFIX = 'input';
+
+const generateInputImageMarkdown = (url: string) => {
+  if (url.toLowerCase().endsWith('.mp4')) {
+    return `![${INPUT_PREFIX}](<${url}>)`;
+  } else {
+    return `![${INPUT_PREFIX}](<${url}>)`;
+  }
+};
+
+/**
+ * The Message we saved to database consists of a prompt and a response
+ * for the UI to use, we need to break them to 2 messages, User and Assistant(if responded)
+ */
+export const convertDBMessageToUIMessage = (
+  messages: Message[],
+): MessageUI[] => {
+  return messages.reduce((acc, message) => {
+    const { id, mediaUrl, prompt, response, result } = message;
+    if (mediaUrl && prompt) {
+      acc.push({
+        id: id + '-user',
+        role: 'user',
+        // add mediaUrl to the prompt
+        content: prompt + '\n\n' + generateInputImageMarkdown(mediaUrl),
+      });
+    }
+    if (response) {
+      acc.push({
+        id: id + '-assistant',
+        role: 'assistant',
+        content: response,
+      });
+    }
+    return acc;
+  }, [] as MessageUI[]);
+};
+
+export const convertAssistantUIMessageToDBMessageResponse = (
+  message: MessageUI,
+): MessageAssistantResponse => {
+  let result = null;
+  const lines = message.content.split('\n');
+  for (let line of lines) {
+    try {
+      const json = JSON.parse(line) as ChunkBody;
+      if (json.type == 'final_code') {
+        result = json as PrismaJson.FinalChatResult;
+        break;
+      }
+    } catch (e) {
+      console.error((e as Error).message);
+    }
+  }
+  return {
+    response: message.content,
+    result,
+  };
+};

prisma/schema.prisma

@@ -29,8 +29,8 @@ model Chat {
   createdAt DateTime  @default(now()) @map("created_at")
   updatedAt DateTime  @updatedAt @map("updated_at")
   title     String    @default("(no title)")
-  userId    String?
   mediaUrl  String
+  userId    String?
   user      User?     @relation(fields: [userId], references: [id])
   messages  Message[]
 
@@ -38,23 +38,18 @@ model Chat {
 }
 
 model Message {
-  id        String      @id @default(cuid())
-  createdAt DateTime    @default(now()) @map("created_at")
-  updatedAt DateTime    @updatedAt @map("updated_at")
+  id        String   @id @default(cuid())
+  createdAt DateTime @default(now()) @map("created_at")
+  updatedAt DateTime @updatedAt @map("updated_at")
   userId    String?
   chatId    String
-  content   String
+  mediaUrl  String
+  prompt    String
+  response  String?
   /// [FinalChatResult]
   result    Json?
-  role      MessageRole
-  chat      Chat        @relation(fields: [chatId], references: [id], onDelete: Cascade)
-  user      User?       @relation(fields: [userId], references: [id])
+  chat      Chat     @relation(fields: [chatId], references: [id], onDelete: Cascade)
+  user      User?    @relation(fields: [userId], references: [id])
 
   @@map("message")
 }
-
-enum MessageRole {
-  system
-  user
-  assistant
-}