feat: Directly us DBMessage as UIMessage + Code viewer (#74)


app/api/vision-agent/route.ts

@@ -1,7 +1,7 @@
 import { StreamingTextResponse, experimental_StreamData } from 'ai';
 
 // import { auth } from '@/auth';
-import { MessageUI, SignedPayload } from '@/lib/types';
+import { MessageUI, ResultPayload, SignedPayload } from '@/lib/types';
 
 import { logger, withLogging } from '@/lib/logger';
 import { CLEANED_SEPARATOR } from '@/lib/constants';
@@ -169,9 +169,7 @@ export const POST = withLogging(
                 if (msg.type !== 'final_code') {
                   return line;
                 }
-                const result = JSON.parse(
-                  msg.payload.result,
-                ) as PrismaJson.FinalChatResult['payload']['result'];
+                const result = JSON.parse(msg.payload.result) as ResultPayload;
                 for (let index = 0; index < result.results.length; index++) {
                   const png = result.results[index].png ?? '';
                   const mp4 = result.results[index].mp4 ?? '';

app/chat/[id]/page.tsx

@@ -1,6 +1,8 @@
 import { Suspense } from 'react';
-import ChatServer from '@/components/chat/ChatServer';
+import ChatServer from './server';
 import Loading from '@/components/ui/Loading';
+import { auth } from '@/auth';
+import { LoginPrompt } from '@/components/chat/LoginPrompt';
 
 interface PageProps {
   params: {
@@ -10,6 +12,7 @@ interface PageProps {
 
 export default async function Page({ params }: PageProps) {
   const { id: chatId } = params;
+  const session = await auth();
   return (
     <Suspense
       fallback={
@@ -18,7 +21,10 @@ export default async function Page({ params }: PageProps) {
         </div>
       }
     >
-      <ChatServer id={chatId} />
+      <div className="w-[1600px] max-w-full mx-auto flex flex-col space-y-4 items-center">
+        {!session && <LoginPrompt />}
+        <ChatServer id={chatId} />
+      </div>
     </Suspense>
   );
 }

components/chat/ChatServer.tsx → app/chat/[id]/server.tsx

@@ -1,4 +1,4 @@
-import ChatClient from './ChatClient';
+import ChatInterface from '../../../components/ChatInterface';
 import { auth } from '@/auth';
 import { dbGetChat } from '@/lib/db/functions';
 import { redirect } from 'next/navigation';
@@ -15,5 +15,5 @@ export default async function ChatServer({ id }: ChatServerProps) {
     revalidatePath('/');
     redirect('/');
   }
-  return <ChatClient chat={chat} />;
+  return <ChatInterface chat={chat} />;
 }

app/layout.tsx

@@ -53,7 +53,7 @@ export default function RootLayout(props: RootLayoutProps) {
         >
           <div className="flex flex-col min-h-screen">
             <Header />
-            <main className="flex py-8 h-[calc(100vh-64px)] bg-background overflow-hidden relative">
+            <main className="flex p-4 h-[calc(100vh-64px)] bg-background overflow-hidden relative w-screen">
               {children}
             </main>
           </div>

app/project/[projectId]/page.tsx

@@ -1,33 +0,0 @@
-import MediaGrid from '@/components/project/MediaGrid';
-import { fetchProjectClass, fetchProjectMedia } from '@/lib/fetch';
-import ProjectChat from '@/components/project/ProjectChat';
-import ClassBar from '@/components/project/ClassBar';
-
-interface PageProps {
-  params: {
-    projectId: string;
-  };
-}
-
-export default async function Page({ params }: PageProps) {
-  const { projectId } = params;
-
-  const [mediaList, classList] = await Promise.all([
-    fetchProjectMedia({ projectId: Number(projectId) }),
-    fetchProjectClass({ projectId: Number(projectId) }),
-  ]);
-
-  return (
-    <div className="pt-4 md:pt-10 h-full">
-      <div className="flex h-full">
-        <div className="w-1/2 relative border-r border-gray-300 overflow-auto">
-          <ClassBar classList={classList} />
-          <MediaGrid mediaList={mediaList} />
-        </div>
-        <div className="w-1/2 relative overflow-auto">
-          <ProjectChat mediaList={mediaList} />
-        </div>
-      </div>
-    </div>
-  );
-}

components/ChatInterface.tsx

@@ -0,0 +1,39 @@
+'use client';
+
+import { ChatWithMessages } from '@/lib/types';
+import React from 'react';
+import ChatList from './chat/ChatList';
+import { Card } from './ui/Card';
+import { useAtom, useAtomValue } from 'jotai';
+import { selectedMessageId } from '@/state/chat';
+import CodeResultDisplay from './CodeResultDisplay';
+
+export interface ChatInterfaceProps {
+  chat: ChatWithMessages;
+}
+
+const ChatInterface: React.FC<ChatInterfaceProps> = ({ chat }) => {
+  const messageId = useAtomValue(selectedMessageId);
+  const messageCodeResult = chat.messages.find(
+    message => message.id === messageId,
+  )?.result;
+  return (
+    <div className="relative flex overflow-hidden space-x-4 size-full">
+      <div
+        data-state={messageCodeResult?.payload ? 'open' : 'closed'}
+        className="pl-4 peer absolute right-0 inset-y-0 hidden translate-x-full data-[state=open]:translate-x-0 z-30 duration-300 ease-in-out xl:flex flex-col items-start xl:w-1/2 h-full dark:bg-zinc-950 overflow-auto"
+      >
+        {messageCodeResult?.payload && (
+          <Card className="w-full">
+            <CodeResultDisplay codeResult={messageCodeResult.payload} />
+          </Card>
+        )}
+      </div>
+      <div className="w-full flex justify-center overflow-auto pr-0 animate-in duration-300 ease-in-out peer-[[data-state=open]]:xl:pr-[50%]">
+        <ChatList chat={chat} />
+      </div>
+    </div>
+  );
+};
+
+export default ChatInterface;

components/CodeResultDisplay.tsx

@@ -0,0 +1,123 @@
+import React from 'react';
+
+import { ChunkBody, CodeResult, formatStreamLogs } from '@/lib/utils/content';
+import { CodeBlock } from './ui/CodeBlock';
+import {
+  Dialog,
+  DialogTrigger,
+  DialogContent,
+  DialogHeader,
+  DialogTitle,
+} from './ui/Dialog';
+import { Button } from './ui/Button';
+import { IconLog, IconTerminalWindow } from './ui/Icons';
+import { Separator } from './ui/Separator';
+import { ResultPayload } from '@/lib/types';
+import Img from './ui/Img';
+
+export interface CodeResultDisplayProps {}
+
+const CodeResultDisplay: React.FC<{
+  codeResult: CodeResult;
+}> = ({ codeResult }) => {
+  const { code, test, result } = codeResult;
+  const getDetail = () => {
+    if (!result) return {};
+    try {
+      const detail = JSON.parse(result) as ResultPayload;
+      return {
+        results: detail.results,
+        stderr: detail.logs.stderr,
+        stdout: detail.logs.stdout,
+      };
+    } catch {
+      return {};
+    }
+  };
+
+  const { results, stderr, stdout } = getDetail();
+
+  return (
+    <div className="rounded-lg overflow-hidden relative max-w-5xl">
+      <CodeBlock language="python" value={code} />
+      <div className="rounded-lg relative">
+        <div className="absolute left-1/2 -translate-x-1/2 -top-4 z-10">
+          <Dialog>
+            <DialogTrigger asChild>
+              <Button variant="ghost" size="icon" className="size-8">
+                <IconTerminalWindow className="text-teal-500 size-4" />
+              </Button>
+            </DialogTrigger>
+            <DialogContent className="max-w-5xl">
+              <DialogHeader>
+                <DialogTitle>Test Code</DialogTitle>
+              </DialogHeader>
+              <CodeBlock language="python" value={test} />
+            </DialogContent>
+          </Dialog>
+          {Array.isArray(stderr) && !!stderr.join('').trim() && (
+            <Dialog>
+              <DialogTrigger asChild>
+                <Button variant="ghost" size="icon" className="size-8">
+                  <IconLog className="text-gray-500 size-4" />
+                </Button>
+              </DialogTrigger>
+              <DialogContent className="max-w-5xl">
+                <CodeBlock language="vim" value={stderr.join('').trim()} />
+              </DialogContent>
+            </Dialog>
+          )}
+        </div>
+      </div>
+      {Array.isArray(stdout) && !!stdout.join('').trim() && (
+        <>
+          <Separator />
+          <CodeBlock language="print" value={stdout.join('').trim()} />
+        </>
+      )}
+      {Array.isArray(results) && !!results.length && (
+        <>
+          <Separator />
+          {results.map((result, index) => {
+            if (result.png) {
+              return (
+                <Img
+                  key={'png' + index}
+                  src={result.png}
+                  alt={'answer-image'}
+                  quality={100}
+                  sizes="(min-width: 66em) 15vw,
+                        (min-width: 44em) 20vw,
+                        100vw"
+                />
+              );
+            } else if (result.mp4) {
+              return (
+                <video
+                  key={'mp4' + index}
+                  src={result.mp4}
+                  controls
+                  width={500}
+                  height={500}
+                />
+              );
+            } else if (result.text) {
+              return (
+                <CodeBlock
+                  key={'text' + index}
+                  language="output"
+                  value={result.text}
+                />
+              );
+            } else {
+              return null;
+            }
+          })}
+        </>
+      )}
+      <Separator />
+    </div>
+  );
+};
+
+export default CodeResultDisplay;

components/chat/ChatClient.tsx

@@ -1,85 +0,0 @@
-'use client';
-
-// import { ChatList } from '@/components/chat/ChatList';
-import Composer from '@/components/chat/Composer';
-import useVisionAgent from '@/lib/hooks/useVisionAgent';
-import { useScrollAnchor } from '@/lib/hooks/useScrollAnchor';
-import { Session } from 'next-auth';
-import { useEffect } from 'react';
-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 { dbPostCreateMessage } from '@/lib/db/functions';
-
-export interface ChatClientProps {
-  chat: ChatWithMessages;
-}
-
-export const SCROLL_BOTTOM = 120;
-
-const ChatClient: React.FC<ChatClientProps> = ({ chat }) => {
-  const { id, messages: dbMessages } = chat;
-  const { messages, append, isLoading } = useVisionAgent(chat);
-
-  const { messagesRef, scrollRef, visibilityRef, isVisible, scrollToBottom } =
-    useScrollAnchor(SCROLL_BOTTOM);
-
-  // Scroll to bottom when messages are loading
-  useEffect(() => {
-    if (isLoading && messages.length) {
-      scrollToBottom();
-    }
-  }, [isLoading, scrollToBottom, messages]);
-
-  return (
-    <div
-      className="h-full overflow-auto mx-auto w-[1024px] max-w-full border rounded-lg relative"
-      ref={scrollRef}
-    >
-      <div className="overflow-auto h-full pt-6 px-6 z-10" ref={messagesRef}>
-        {messages
-          // .filter(message => message.role !== 'system')
-          .map((message, index) => (
-            <ChatMessage
-              key={index}
-              message={message}
-              isLoading={isLoading && index === messages.length - 1}
-            />
-          ))}
-        <div
-          className="w-full"
-          style={{ height: SCROLL_BOTTOM }}
-          ref={visibilityRef}
-        />
-      </div>
-      <div className="absolute bottom-4 w-full">
-        <Composer
-          // Use the last message mediaUrl as the initial mediaUrl
-          initMediaUrl={dbMessages[dbMessages.length - 1]?.mediaUrl}
-          isLoading={isLoading}
-          onSubmit={async ({ input, mediaUrl: newMediaUrl }) => {
-            append({
-              prompt: input,
-              mediaUrl: newMediaUrl,
-            });
-          }}
-        />
-      </div>
-      {/* Scroll to bottom Icon */}
-      <Button
-        size="icon"
-        className={cn(
-          'absolute bottom-3 right-3 transition-opacity duration-300 size-6',
-          isVisible ? 'opacity-0' : 'opacity-100',
-        )}
-        onClick={() => scrollToBottom()}
-      >
-        <IconArrowDown className="size-3" />
-      </Button>
-    </div>
-  );
-};
-
-export default ChatClient;

components/chat/ChatList.tsx

@@ -1,63 +1,103 @@
 'use client';
 
-import { Separator } from '@/components/ui/Separator';
-import { ChatMessage } from '@/components/chat/ChatMessage';
-import { MessageUI } from '@/lib/types';
+// import { ChatList } from '@/components/chat/ChatList';
+import Composer from '@/components/chat/Composer';
+import useVisionAgent from '@/lib/hooks/useVisionAgent';
+import { useScrollAnchor } from '@/lib/hooks/useScrollAnchor';
 import { Session } from 'next-auth';
-import { IconExclamationTriangle } from '../ui/Icons';
-import Link from 'next/link';
+import { useEffect } from 'react';
+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 { dbPostCreateMessage } from '@/lib/db/functions';
+import { Card } from '../ui/Card';
+import { useSetAtom } from 'jotai';
+import { selectedMessageId } from '@/state/chat';
 
-export interface ChatList {
-  messages: MessageUI[];
-  session: Session | null;
-  isLoading: boolean;
+export interface ChatListProps {
+  chat: ChatWithMessages;
 }
 
-export function ChatList({ messages, session, isLoading }: ChatList) {
+export const SCROLL_BOTTOM = 120;
+
+const ChatList: React.FC<ChatListProps> = ({ chat }) => {
+  const { id, messages: dbMessages } = chat;
+  const { messages, append, isLoading } = useVisionAgent(chat);
+
+  const lastMessage = messages[messages.length - 1];
+  const lastDbMessage = dbMessages[dbMessages.length - 1];
+  const setMessageId = useSetAtom(selectedMessageId);
+
+  const { messagesRef, scrollRef, visibilityRef, isVisible, scrollToBottom } =
+    useScrollAnchor(SCROLL_BOTTOM);
+
+  // Scroll to bottom on init and highlight last message
+  useEffect(() => {
+    scrollToBottom();
+    if (lastDbMessage.result) {
+      setMessageId(lastDbMessage.id);
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, []);
+
+  // Scroll to bottom when messages are loading
+  useEffect(() => {
+    if (isLoading && messages.length) {
+      scrollToBottom();
+    }
+  }, [isLoading, scrollToBottom, messages]);
+
   return (
-    <div className="relative mx-auto max-w-5xl px-8 pt-6 border rounded-lg">
-      {!session && (
-        <>
-          <div className="group relative mb-6 flex items-center">
-            <div className="bg-background flex size-8 shrink-0 select-none items-center justify-center rounded-md border shadow">
-              <IconExclamationTriangle />
-            </div>
-            <div className="flex-1 px-1 ml-4 space-y-2 overflow-hidden">
-              {process.env.NEXT_PUBLIC_IS_HUGGING_FACE ? (
-                <p className="text-muted-foreground leading-normal">
-                  Please visit and login into{' '}
-                  <Link
-                    href="https://va.landing.ai/"
-                    target="_blank"
-                    className="underline"
-                  >
-                    our landing website
-                  </Link>{' '}
-                  to save and revisit your chat history!
-                </p>
-              ) : (
-                <p className="text-muted-foreground leading-normal">
-                  Please{' '}
-                  <Link href="/sign-in" className="underline">
-                    log in
-                  </Link>{' '}
-                  to save and revisit your chat history!
-                </p>
-              )}
-            </div>
-          </div>
-          <Separator className="my-4" />
-        </>
-      )}
-      {messages
-        // .filter(message => message.role !== 'system')
-        .map((message, index) => (
+    <Card
+      className="size-full max-w-5xl overflow-auto relative"
+      ref={scrollRef}
+    >
+      <div className="overflow-auto h-full p-4 z-10" ref={messagesRef}>
+        {dbMessages.map((message, index) => (
           <ChatMessage
             key={index}
             message={message}
-            isLoading={isLoading && index === messages.length - 1}
+            wipAssistantMessage={
+              isLoading && lastMessage.role === 'assistant'
+                ? lastMessage
+                : undefined
+            }
           />
         ))}
-    </div>
+        <div
+          className="w-full"
+          style={{ height: SCROLL_BOTTOM }}
+          ref={visibilityRef}
+        />
+      </div>
+      <div className="absolute bottom-4 w-full">
+        <Composer
+          // Use the last message mediaUrl as the initial mediaUrl
+          initMediaUrl={dbMessages[dbMessages.length - 1]?.mediaUrl}
+          isLoading={isLoading}
+          onSubmit={async ({ input, mediaUrl: newMediaUrl }) => {
+            append({
+              prompt: input,
+              mediaUrl: newMediaUrl,
+            });
+          }}
+        />
+      </div>
+      {/* Scroll to bottom Icon */}
+      <Button
+        size="icon"
+        className={cn(
+          'absolute bottom-3 right-3 transition-opacity duration-300 size-6',
+          isVisible ? 'opacity-0' : 'opacity-100',
+        )}
+        onClick={() => scrollToBottom()}
+      >
+        <IconArrowDown className="size-3" />
+      </Button>
+    </Card>
   );
-}
+};
+
+export default ChatList;

components/chat/ChatMessage.tsx

@@ -8,8 +8,6 @@ import {
   IconListUnordered,
   IconTerminalWindow,
   IconUser,
-  IconOutput,
-  IconLog,
   IconGlowingDot,
 } from '@/components/ui/Icons';
 import { MessageUI } from '@/lib/types';
@@ -23,59 +21,93 @@ import {
   TableRow,
 } from '../ui/Table';
 import { Button } from '../ui/Button';
-import { Separator } from '../ui/Separator';
-import {
-  Tooltip,
-  TooltipContent,
-  TooltipTrigger,
-} from '@/components/ui/Tooltip';
-
 import { Dialog, DialogContent, DialogTrigger } from '../ui/Dialog';
-import { Markdown } from './MemoizedReactMarkdown';
 import Img from '../ui/Img';
+import CodeResultDisplay from '../CodeResultDisplay';
+import { useAtom, useSetAtom } from 'jotai';
+import { selectedMessageId } from '@/state/chat';
+import { Message } from '@prisma/client';
+import { Separator } from '../ui/Separator';
+import { cn } from '@/lib/utils';
 
 export interface ChatMessageProps {
-  message: MessageUI;
-  isLoading: boolean;
+  message: Message;
+  wipAssistantMessage?: MessageUI;
 }
 
-export function ChatMessage({ message, isLoading }: ChatMessageProps) {
-  const { role, content, mediaUrl } = message;
-
-  return role === 'user' ? (
-    <UserChatMessage content={content} mediaUrl={mediaUrl} />
-  ) : (
-    <AssistantChatMessage content={content} />
+export const ChatMessage: React.FC<ChatMessageProps> = ({
+  message,
+  wipAssistantMessage,
+}) => {
+  const [messageId, setMessageId] = useAtom(selectedMessageId);
+  const { id, mediaUrl, prompt, response, result } = message;
+  const [formattedSections, codeResult] = useMemo(
+    () => formatStreamLogs(response ?? wipAssistantMessage?.content),
+    [response, wipAssistantMessage?.content],
   );
-}
-
-const UserChatMessage: React.FC<{
-  content: string;
-  mediaUrl?: string;
-}> = ({ content, mediaUrl }) => {
   return (
-    <div className="group relative mb-6 flex rounded-md bg-muted p-6 ml-auto mr-0 w-3/5">
-      <div className="flex size-8 shrink-0 select-none items-center justify-center rounded-md border shadow bg-background">
-        <IconUser />
+    <div
+      className={cn(
+        'rounded-md bg-muted border border-muted p-4 mb-4',
+        messageId === id && 'lg:border-primary/50',
+        result && 'lg:cursor-pointer',
+      )}
+      onClick={() => {
+        if (result) {
+          setMessageId(id);
+        }
+      }}
+    >
+      <div className="flex">
+        <div className="flex size-8 shrink-0 select-none items-center justify-center rounded-md border shadow bg-background">
+          <IconUser />
+        </div>
+        <div className="flex-1 px-1 ml-4 space-y-2 overflow-hidden">
+          <p>{prompt}</p>
+          {mediaUrl && (
+            <>
+              {mediaUrl?.endsWith('.mp4') ? (
+                <video src={mediaUrl} controls width={500} height={500} />
+              ) : (
+                <Img src={mediaUrl} alt={mediaUrl} quality={100} width={300} />
+              )}
+            </>
+          )}
+        </div>
       </div>
-      <div className="flex-1 px-1 ml-4 space-y-3 overflow-hidden">
-        <p>{content}</p>
-        {mediaUrl && (
-          <>
-            {mediaUrl?.endsWith('.mp4') ? (
-              <video src={mediaUrl} controls width={500} height={500} />
-            ) : (
-              <Img
-                src={mediaUrl}
-                alt={mediaUrl}
-                quality={100}
-                sizes="(min-width: 66em) 15vw,
-                        (min-width: 44em) 20vw,
-                        100vw"
-              />
-            )}
-          </>
-        )}
+      <Separator className="bg-primary/30 my-4" />
+      <div className="flex">
+        <div className="flex size-8 shrink-0 select-none items-center justify-center rounded-md border shadow bg-primary text-primary-foreground">
+          <IconLandingAI />
+        </div>
+        <div className="flex-1 px-1 space-y-4 ml-4 overflow-hidden">
+          <Table className="w-[400px]">
+            <TableBody>
+              {formattedSections.map(section => (
+                <TableRow
+                  className="border-primary/50 h-[56px]"
+                  key={section.type}
+                >
+                  <TableCell className="text-center text-webkit-center">
+                    {ChunkStatusToIconDict[section.status]}
+                  </TableCell>
+                  <TableCell className="font-medium">
+                    {ChunkTypeToTextDict[section.type]}
+                  </TableCell>
+                  <TableCell className="text-right">
+                    <ChunkPayloadAction payload={section.payload} />
+                  </TableCell>
+                </TableRow>
+              ))}
+            </TableBody>
+          </Table>
+          {codeResult && (
+            <div className="xl:hidden">
+              <CodeResultDisplay codeResult={codeResult} />
+            </div>
+          )}
+          {codeResult && <p>✨ Coding complete</p>}
+        </div>
       </div>
     </div>
   );
@@ -113,7 +145,10 @@ const ChunkPayloadAction: React.FC<{
             <IconListUnordered />
           </Button>
         </DialogTrigger>
-        <DialogContent className="max-w-5xl" onOpenAutoFocus={e => e.preventDefault()}>
+        <DialogContent
+          className="max-w-5xl"
+          onOpenAutoFocus={e => e.preventDefault()}
+        >
           <Table>
             <TableHeader>
               <TableRow className="border-primary/50">
@@ -170,144 +205,4 @@ const ChunkPayloadAction: React.FC<{
   }
 };
 
-const CodeResultDisplay: React.FC<{
-  codeResult: CodeResult;
-}> = ({ codeResult }) => {
-  const { code, test, result } = codeResult;
-  const getDetail = () => {
-    if (!result) return {};
-    try {
-      const detail = JSON.parse(result);
-      return {
-        results: detail.results,
-        stderr: detail.logs.stderr,
-        stdout: detail.logs.stdout,
-      };
-    } catch {
-      return {};
-    }
-  };
-
-  const { results, stderr, stdout } = getDetail();
-
-  return (
-    <div className="rounded-lg overflow-hidden relative max-w-5xl">
-      <CodeBlock language="python" value={code} />
-      <div className="rounded-lg relative">
-        <div className="absolute left-1/2 -translate-x-1/2 -top-4 z-10">
-          <Dialog>
-            <DialogTrigger asChild>
-              <Button variant="ghost" size="icon" className="size-8">
-                <IconTerminalWindow className="text-teal-500 size-4" />
-              </Button>
-            </DialogTrigger>
-            <DialogContent className="max-w-5xl">
-              <CodeBlock language="python" value={test} />
-            </DialogContent>
-          </Dialog>
-          {Array.isArray(stderr) && (
-            <Dialog>
-              <DialogTrigger asChild>
-                <Button variant="ghost" size="icon" className="size-8">
-                  <IconLog className="text-gray-500 size-4" />
-                </Button>
-              </DialogTrigger>
-              <DialogContent className="max-w-5xl">
-                <CodeBlock language="vim" value={stderr.join('').trim()} />
-              </DialogContent>
-            </Dialog>
-          )}
-        </div>
-      </div>
-      {Array.isArray(stdout) && !!stdout.join('').trim() && (
-        <>
-          <Separator />
-          <CodeBlock language="print" value={stdout.join('').trim()} />
-        </>
-      )}
-      {Array.isArray(results) && !!results.length && (
-        <>
-          <Separator />
-          {results.map((result, index) => {
-            if (result.png) {
-              return (
-                <Img
-                  key={'png' + index}
-                  src={result.png}
-                  alt={'answer-image'}
-                  quality={100}
-                  sizes="(min-width: 66em) 15vw,
-                        (min-width: 44em) 20vw,
-                        100vw"
-                />
-              );
-            } else if (result.mp4) {
-              return (
-                <video
-                  key={'mp4' + index}
-                  src={result.mp4}
-                  controls
-                  width={500}
-                  height={500}
-                />
-              );
-            } else if (result.text) {
-              return (
-                <CodeBlock
-                  key={'text' + index}
-                  language="output"
-                  value={result.text}
-                />
-              );
-            } else {
-              return null;
-            }
-          })}
-        </>
-      )}
-      <Separator />
-    </div>
-  );
-};
-
-const AssistantChatMessage: React.FC<{
-  content: string;
-}> = ({ content }) => {
-  const [formattedSections, codeResult] = useMemo(
-    () => formatStreamLogs(content),
-    [content],
-  );
-
-  return (
-    <div className="group relative mb-6 flex rounded-md bg-muted p-6 w-full">
-      <div className="flex size-8 shrink-0 select-none items-center justify-center rounded-md border shadow bg-primary text-primary-foreground">
-        <IconLandingAI />
-      </div>
-      <div className="flex-1 px-1 space-y-4 ml-4 overflow-hidden">
-        <Table className="w-[400px]">
-          <TableBody>
-            {formattedSections.map(section => (
-              <TableRow
-                className="border-primary/50 h-[56px]"
-                key={section.type}
-              >
-                <TableCell className="text-center text-webkit-center">
-                  {ChunkStatusToIconDict[section.status]}
-                </TableCell>
-                <TableCell className="font-medium">
-                  {ChunkTypeToTextDict[section.type]}
-                </TableCell>
-                <TableCell className="text-right">
-                  <ChunkPayloadAction payload={section.payload} />
-                </TableCell>
-              </TableRow>
-            ))}
-          </TableBody>
-        </Table>
-        {codeResult && <CodeResultDisplay codeResult={codeResult} />}
-      </div>
-    </div>
-  );
-};
-
-export default UserChatMessage;
+export default ChatMessage;

components/chat/Composer.tsx

@@ -75,7 +75,7 @@ const Composer = forwardRef<ComposerRef, ComposerProps>(
       <div
         {...getRootProps()}
         className={cn(
-          'w-full mx-auto max-w-2xl px-6 py-4 bg-zinc-600 rounded-xl relative shadow-lg shadow-zinc-600/40 z-50',
+          'mx-auto w-[42rem] max-w-full px-6 py-4 bg-zinc-600 rounded-xl relative shadow-lg shadow-zinc-600/40 z-50',
           isDragActive && 'bg-indigo-700/50',
         )}
       >

components/chat/LoginPrompt.tsx

@@ -0,0 +1,25 @@
+'use client';
+
+import { Card } from '../ui/Card';
+import { IconExclamationTriangle } from '../ui/Icons';
+import Link from 'next/link';
+
+export interface LoginPrompt {}
+
+export function LoginPrompt() {
+  return (
+    <Card className="group py-2 px-4 flex items-center">
+      <div className="bg-background flex size-8 shrink-0 select-none items-center justify-center rounded-md">
+        <IconExclamationTriangle className="font-medium" />
+      </div>
+      <div className="flex-1 px-1 ml-2 overflow-hidden">
+        <p className="leading-normal font-medium">
+          <Link href="/sign-in" className="underline">
+            Sign in
+          </Link>{' '}
+          to save and revisit your chat history!
+        </p>
+      </div>
+    </Card>
+  );
+}

components/project/ClassBar.tsx

@@ -1,22 +0,0 @@
-import { ClassDetails } from '@/lib/fetch';
-import Chip from '../ui/Chip';
-
-export interface ClassBarProps {
-  classList: ClassDetails[];
-}
-
-export default async function ClassBar({ classList }: ClassBarProps) {
-  return (
-    <div className="border-b border-gray-300 px-3 pb-3 max-w-3xl mx-auto">
-      {classList.map(classItem => {
-        return (
-          <Chip
-            key={classItem.id}
-            value={classItem.name}
-            className="px-3 py-1 my-1"
-          />
-        );
-      })}
-    </div>
-  );
-}

components/project/MediaGrid.tsx

@@ -1,18 +0,0 @@
-import { MediaDetails } from '@/lib/fetch';
-import MediaTile from './MediaTile';
-
-export default function MediaGrid({
-  mediaList,
-}: {
-  mediaList: MediaDetails[];
-}) {
-  return (
-    <div className="relative size-full p-3 max-w-3xl mx-auto">
-      <div className="columns-1 sm:columns-1 md:columns-2 lg:columns-2 xl:columns:3 gap-3 [&>img:not(:first-child)]:mt-3">
-        {mediaList.map(media => (
-          <MediaTile key={media.id} media={media} />
-        ))}
-      </div>
-    </div>
-  );
-}

components/project/MediaTile.tsx

@@ -1,53 +0,0 @@
-'use client';
-
-import React from 'react';
-import Image from 'next/image';
-import {
-  Tooltip,
-  TooltipContent,
-  TooltipTrigger,
-} from '@/components/ui/Tooltip';
-import { MediaDetails } from '@/lib/fetch';
-import { useAtom } from 'jotai';
-import { selectedMediaIdAtom } from '@/state/media';
-import { cn } from '@/lib/utils';
-
-export interface MediaTileProps {
-  media: MediaDetails;
-}
-
-export default function MediaTile({ media }: MediaTileProps) {
-  const {
-    url,
-    thumbnails,
-    id,
-    name,
-    properties: { width, height },
-  } = media;
-  const [selectedMediaId, setSelectedMediaId] = useAtom(selectedMediaIdAtom);
-  const selected = selectedMediaId === id;
-  // const imageSrc = thumbnails.length ? thumbnails[thumbnails.length - 1] : url;
-  const imageSrc = url;
-  return (
-    <Tooltip>
-      <TooltipTrigger asChild>
-        <Image
-          src={imageSrc}
-          draggable={false}
-          alt="dataset images"
-          width={width}
-          height={height}
-          onClick={() => setSelectedMediaId(id)}
-          className={cn(
-            'w-full h-auto relative rounded-xl overflow-hidden shadow-md cursor-pointer transition-transform hover:scale-105 box-content',
-            selected && 'border-2 border-primary',
-          )}
-        />
-      </TooltipTrigger>
-      <TooltipContent>
-        <p>{name}</p>
-        <p className="font-light text-xs">{`${width} x ${height}`}</p>
-      </TooltipContent>
-    </Tooltip>
-  );
-}

components/ui/Card.tsx

@@ -0,0 +1,83 @@
+import * as React from 'react';
+
+import { cn } from '@/lib/utils';
+
+const Card = React.forwardRef<
+  HTMLDivElement,
+  React.HTMLAttributes<HTMLDivElement>
+>(({ className, ...props }, ref) => (
+  <div
+    ref={ref}
+    className={cn(
+      'rounded-xl border bg-card text-card-foreground shadow',
+      className,
+    )}
+    {...props}
+  />
+));
+Card.displayName = 'Card';
+
+const CardHeader = React.forwardRef<
+  HTMLDivElement,
+  React.HTMLAttributes<HTMLDivElement>
+>(({ className, ...props }, ref) => (
+  <div
+    ref={ref}
+    className={cn('flex flex-col space-y-1.5 p-6', className)}
+    {...props}
+  />
+));
+CardHeader.displayName = 'CardHeader';
+
+const CardTitle = React.forwardRef<
+  HTMLParagraphElement,
+  React.HTMLAttributes<HTMLHeadingElement>
+>(({ className, ...props }, ref) => (
+  <h3
+    ref={ref}
+    className={cn('font-semibold leading-none tracking-tight', className)}
+    {...props}
+  />
+));
+CardTitle.displayName = 'CardTitle';
+
+const CardDescription = React.forwardRef<
+  HTMLParagraphElement,
+  React.HTMLAttributes<HTMLParagraphElement>
+>(({ className, ...props }, ref) => (
+  <p
+    ref={ref}
+    className={cn('text-sm text-muted-foreground', className)}
+    {...props}
+  />
+));
+CardDescription.displayName = 'CardDescription';
+
+const CardContent = React.forwardRef<
+  HTMLDivElement,
+  React.HTMLAttributes<HTMLDivElement>
+>(({ className, ...props }, ref) => (
+  <div ref={ref} className={cn('p-6 pt-0', className)} {...props} />
+));
+CardContent.displayName = 'CardContent';
+
+const CardFooter = React.forwardRef<
+  HTMLDivElement,
+  React.HTMLAttributes<HTMLDivElement>
+>(({ className, ...props }, ref) => (
+  <div
+    ref={ref}
+    className={cn('flex items-center p-6 pt-0', className)}
+    {...props}
+  />
+));
+CardFooter.displayName = 'CardFooter';
+
+export {
+  Card,
+  CardHeader,
+  CardFooter,
+  CardTitle,
+  CardDescription,
+  CardContent,
+};

components/ui/CodeBlock.tsx

@@ -45,9 +45,14 @@ export const programmingLanguages: languageMap = {
   sql: '.sql',
   html: '.html',
   css: '.css',
+  print: '.txt',
   // add more file extensions here, make sure the key is same as language prop in CodeBlock.tsx component
 };
 
+const customSyntax: languageMap = {
+  print: 'vim',
+};
+
 export const generateRandomString = (length: number, lowercase = false) => {
   const chars = 'ABCDEFGHJKLMNPQRSTUVWXY3456789'; // excluding similar looking characters like Z, 2, I, 1, O, 0
   let result = '';
@@ -118,7 +123,7 @@ const CodeBlock: FC<Props> = memo(({ language, value }) => {
         </div>
       </div>
       <SyntaxHighlighter
-        language={language}
+        language={customSyntax[language] ?? language}
         style={coldarkDark}
         PreTag="div"
         showLineNumbers

components/ui/Icons.tsx

@@ -585,6 +585,8 @@ function IconExclamationTriangle({
       viewBox="0 0 15 15"
       fill="none"
       xmlns="http://www.w3.org/2000/svg"
+      className={cn('size-4', className)}
+      {...props}
     >
       <path
         d="M8.4449 0.608765C8.0183 -0.107015 6.9817 -0.107015 6.55509 0.608766L0.161178 11.3368C-0.275824 12.07 0.252503 13 1.10608 13H13.8939C14.7475 13 15.2758 12.07 14.8388 11.3368L8.4449 0.608765ZM7.4141 1.12073C7.45288 1.05566 7.54712 1.05566 7.5859 1.12073L13.9798 11.8488C14.0196 11.9154 13.9715 12 13.8939 12H1.10608C1.02849 12 0.980454 11.9154 1.02018 11.8488L7.4141 1.12073ZM6.8269 4.48611C6.81221 4.10423 7.11783 3.78663 7.5 3.78663C7.88217 3.78663 8.18778 4.10423 8.1731 4.48612L8.01921 8.48701C8.00848 8.766 7.7792 8.98664 7.5 8.98664C7.2208 8.98664 6.99151 8.766 6.98078 8.48701L6.8269 4.48611ZM8.24989 10.476C8.24989 10.8902 7.9141 11.226 7.49989 11.226C7.08567 11.226 6.74989 10.8902 6.74989 10.476C6.74989 10.0618 7.08567 9.72599 7.49989 9.72599C7.9141 9.72599 8.24989 10.0618 8.24989 10.476Z"

lib/db/functions.ts

@@ -60,6 +60,9 @@ export async function dbGetMyChatList(): Promise<Chat[]> {
 
   return prisma.chat.findMany({
     where: { userId },
+    orderBy: {
+      createdAt: 'desc',
+    },
   });
 }
 
@@ -72,7 +75,11 @@ export async function dbGetChat(id: string): Promise<ChatWithMessages | null> {
   return prisma.chat.findUnique({
     where: { id },
     include: {
-      messages: true,
+      messages: {
+        orderBy: {
+          createdAt: 'asc',
+        },
+      },
     },
   });
 }

lib/db/prisma.ts

@@ -10,18 +10,18 @@ declare global {
       payload: {
         code: string;
         test: string;
-        result: {
-          logs: {
-            stderr: string[];
-            stdout: string[];
-          };
-          results: Array<{
-            png?: string;
-            mp4?: string;
-            text: string;
-            is_main_result: boolean;
-          }>;
-        };
+        result: string; // TODO To be fixed to JSON below
+        // result: {
+        //   logs: {
+        //     stderr: string[];
+        //     stdout: string[];
+        //   };
+        //   results: Array<{
+        //     png?: string;
+        //     text: string;
+        //     is_main_result: boolean;
+        //   }>;
+        // };
       };
     };
   }

lib/hooks/useVisionAgent.ts

@@ -10,10 +10,13 @@ import {
   convertAssistantUIMessageToDBMessageResponse,
   convertDBMessageToUIMessage,
 } from '../utils/message';
+import { useSetAtom } from 'jotai';
+import { selectedMessageId } from '@/state/chat';
 
 const useVisionAgent = (chat: ChatWithMessages) => {
   const { messages: dbMessages, id, mediaUrl } = chat;
   const latestDbMessage = dbMessages[dbMessages.length - 1];
+  const setMessageId = useSetAtom(selectedMessageId);
 
   // Temporary solution for now while single we have to pass mediaUrl separately outside of the messages
   const currMediaUrl = useRef<string>(mediaUrl);
@@ -32,6 +35,7 @@ const useVisionAgent = (chat: ChatWithMessages) => {
         currMessageId.current,
         convertAssistantUIMessageToDBMessageResponse(message),
       );
+      setMessageId(currMessageId.current);
     },
     sendExtraMessageFields: true,
     initialMessages: convertDBMessageToUIMessage(dbMessages),
@@ -63,7 +67,9 @@ const useVisionAgent = (chat: ChatWithMessages) => {
   return {
     messages: messages as MessageUI[],
     append: async (messageInput: MessageUserInput) => {
+      const resp = await dbPostCreateMessage(id, messageInput);
       currMediaUrl.current = messageInput.mediaUrl;
+      currMessageId.current = resp.id;
       append({
         id,
         role: 'user',
@@ -71,8 +77,6 @@ const useVisionAgent = (chat: ChatWithMessages) => {
         // @ts-ignore valid when setting sendExtraMessageFields
         mediaUrl: messageInput.mediaUrl,
       });
-      const resp = await dbPostCreateMessage(id, messageInput);
-      currMessageId.current = resp.id;
     },
     reload,
     isLoading,

lib/types.ts

@@ -18,3 +18,16 @@ export interface SignedPayload {
   signedUrl: string;
   fields: Record<string, string>;
 }
+
+export type ResultPayload = {
+  logs: {
+    stderr: string[];
+    stdout: string[];
+  };
+  results: Array<{
+    png?: string;
+    mp4?: string;
+    text: string;
+    is_main_result: boolean;
+  }>;
+};

lib/utils/content.ts

@@ -77,34 +77,6 @@ const generateCodeExecutionMarkdown = (
   return message;
 };
 
-const generateFinalCodeMarkdown = (
-  code: string,
-  test: string,
-  result: PrismaJson.FinalChatResult['payload']['result'],
-) => {
-  let message = 'Final Code: \n';
-  message += `\`\`\`python\n${code}\n\`\`\`\n`;
-  message += 'Final test: \n';
-  message += `\`\`\`python\n${test}\n\`\`\`\n`;
-  message += `Final result: \n`;
-  const images = result.results.map(result => result.png).filter(png => !!png);
-  if (images.length > 0) {
-    message += `Visualization output:\n`;
-    images.forEach((image, index) => {
-      message += generateAnswersImageMarkdown(index, image!);
-    });
-  }
-  if (result.logs.stderr.length > 0) {
-    message += `Error output:\n`;
-    message += `\`\`\`\n${result.logs.stderr.join('\n')}\n\`\`\`\n`;
-  }
-  if (result.logs.stdout.length > 0) {
-    message += `Output:\n`;
-    message += `\`\`\`\n${result.logs.stdout.join('\n')}\n\`\`\`\n`;
-  }
-  return message;
-};
-
 type PlansBody =
   | {
       type: 'plans';
@@ -255,8 +227,9 @@ export type ChunkBody =
  * @returns An array of grouped sections and an optional final code result.
  */
 export const formatStreamLogs = (
-  content: string,
+  content: string | null | undefined,
 ): [ChunkBody[], CodeResult?] => {
+  if (!content) return [[], undefined];
   const streamLogs = content.split('\n').filter(log => !!log);
 
   const buffer = streamLogs.pop();

state/chat.ts

@@ -1,3 +1,3 @@
 import { atom } from 'jotai';
 
-export const chatViewMode = atom<'chat' | 'chat-all'>('chat');
+export const selectedMessageId = atom<string | undefined>(undefined);

state/media.ts

@@ -1,3 +0,0 @@
-import { atom } from 'jotai';
-
-export const selectedMediaIdAtom = atom<number | null>(null);