feat: conversation restructure, assistant message should be just code (#89)





Also fix backwards compatibly from
https://github.com/landing-ai/vision-agent-ui/pull/86

app/api/vision-agent/route.ts

@@ -1,11 +1,9 @@
 import { JSONValue, StreamingTextResponse, experimental_StreamData } from 'ai';
 
 // import { auth } from '@/auth';
-import { MessageUI, ResultPayload, 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/utils/content';
 import { getPresignedUrl } from '@/lib/aws';
 
 // export const runtime = 'edge';
@@ -68,7 +66,7 @@ const modifyCodePayload = async (
   ) {
     return msg;
   }
-  const result = JSON.parse(msg.payload.result) as ResultPayload;
+  const result = JSON.parse(msg.payload.result) as PrismaJson.StructuredResult;
   if (msg.type === 'code') {
     if (result && result.results) {
       msg.payload.result = {
@@ -106,45 +104,18 @@ export const POST = withLogging(
   async (
     session,
     json: {
-      messages: MessageUI[];
+      apiMessages: string;
       id: string;
       mediaUrl: string;
     },
     request,
   ) => {
-    const { messages, mediaUrl } = json;
+    const { apiMessages, mediaUrl } = json;
+    const messages: MessageUI[] = JSON.parse(apiMessages);
     const user = session?.user?.email ?? 'anonymous';
 
-    // const session = await auth();
-    // if (!session?.user?.email) {
-    //   return new Response('Unauthorized', {
-    //     status: 401,
-    //   });
-    // }
-
     const formData = new FormData();
-    formData.append(
-      'input',
-      JSON.stringify(
-        messages.map(message => {
-          if (message.role !== 'assistant') {
-            return {
-              ...message,
-              content: cleanInputMessage(message.content),
-            };
-          } else {
-            const splitedContent = message.content.split(CLEANED_SEPARATOR);
-            return {
-              ...message,
-              content:
-                splitedContent.length > 1
-                  ? cleanAnswerMessage(splitedContent[1])
-                  : message.content,
-            };
-          }
-        }),
-      ),
-    );
+    formData.append('input', JSON.stringify(messages));
     formData.append('image', mediaUrl);
 
     const agentHost = process.env.LND_TIER
@@ -213,14 +184,14 @@ export const POST = withLogging(
             .filter(line => line.trim().length > 0);
           if (lines.length === 0) {
             if (Date.now() - time > TIMEOUT_MILI_SECONDS) {
-              logger.info(
-                session,
-                {
-                  message: 'Agent timed out',
-                },
-                request,
-                '__Agent_timeout__',
-              );
+              // logger.info(
+              //   session,
+              //   {
+              //     message: 'Agent timed out',
+              //   },
+              //   request,
+              //   '__Agent_timeout__',
+              // );
               controller.enqueue(
                 encoder.encode(JSON.stringify(FINAL_TIMEOUT_ERROR) + '\n'),
               );

app/page.tsx

@@ -15,7 +15,7 @@ const EXAMPLES = [
     mediaUrl:
       'https://vision-agent-dev.s3.us-east-2.amazonaws.com/examples/flower.png',
     prompt:
-      'Draw box and output the image, return the number of flowers as output.',
+      'Detect flowers in this image, draw boxes and output the image, also return total number of flowers',
   },
   {
     title: 'Detecting sharks in video',

components/CodeResultDisplay.tsx

@@ -1,6 +1,5 @@
 import React from 'react';
 
-import { ChunkBody, CodeResult, formatStreamLogs } from '@/lib/utils/content';
 import { CodeBlock } from './ui/CodeBlock';
 import {
   Dialog,
@@ -12,7 +11,6 @@ import {
 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';
 import {
   Carousel,
@@ -25,13 +23,18 @@ import {
 export interface CodeResultDisplayProps {}
 
 const CodeResultDisplay: React.FC<{
-  codeResult: CodeResult;
+  codeResult: PrismaJson.FinalChatResult['payload'];
 }> = ({ codeResult }) => {
   const { code, test, result } = codeResult;
   const getDetail = () => {
     if (!result) return {};
     try {
-      const detail = JSON.parse(result) as ResultPayload;
+      // IMPORTANT: This is for backwards compatibility with old chat that save result as JSON string
+      // updated in https://github.com/landing-ai/vision-agent-ui/pull/86
+      const detail =
+        typeof result === 'object'
+          ? result
+          : (JSON.parse(result) as PrismaJson.StructuredResult);
       return {
         results: detail.results,
         stderr: detail.logs.stderr,
@@ -108,7 +111,9 @@ const CodeResultDisplay: React.FC<{
             <p>image output</p>
             <Dialog>
               <DialogTrigger asChild>
-                <Button variant="ghost">View all</Button>
+                <Button variant="outline" size="sm">
+                  View all
+                </Button>
               </DialogTrigger>
               <DialogContent className="max-w-5xl flex justify-center items-center">
                 <Carousel className="w-3/4">
@@ -138,7 +143,7 @@ const CodeResultDisplay: React.FC<{
         </div>
       )}
       {!!videoResults.length && (
-        <div className="p-4 text-xs lowercase bg-zinc-900 space-y-4">
+        <div className="p-4 text-xs lowercase bg-zinc-900 space-y-4 border-t border-muted">
           <p>video output</p>
           <div className="flex flex-row space-x-4 overflow-auto">
             {videoResults.map((mp4, index) => (

components/chat/ChatList.tsx

@@ -30,7 +30,11 @@ const ChatList: React.FC<ChatListProps> = ({ chat, userId }) => {
   // Only login and chat owner can compose
   const canCompose = !chatUserId || userId === chatUserId;
 
-  const lastMessage = messages[messages.length - 1];
+  const lastAssistantMessage = messages.length
+    ? messages[messages.length - 1]?.role === 'assistant'
+      ? messages[messages.length - 1]?.content
+      : undefined
+    : undefined;
   const lastDbMessage = dbMessages[dbMessages.length - 1];
   const setMessageId = useSetAtom(selectedMessageId);
 
@@ -67,8 +71,8 @@ const ChatList: React.FC<ChatListProps> = ({ chat, userId }) => {
               message={message}
               loading={isLastMessage && isLoading}
               wipAssistantMessage={
-                lastMessage.role === 'assistant' && isLastMessage
-                  ? lastMessage
+                lastAssistantMessage && isLastMessage
+                  ? lastAssistantMessage
                   : undefined
               }
             />

components/chat/ChatMessage.tsx

@@ -11,11 +11,7 @@ import {
   IconGlowingDot,
 } from '@/components/ui/Icons';
 import { MessageUI } from '@/lib/types';
-import {
-  WIPChunkBodyGroup,
-  CodeResult,
-  formatStreamLogs,
-} from '@/lib/utils/content';
+import { WIPChunkBodyGroup, formatStreamLogs } from '@/lib/utils/content';
 import {
   Table,
   TableBody,
@@ -37,7 +33,7 @@ import { cn } from '@/lib/utils';
 export interface ChatMessageProps {
   message: Message;
   loading?: boolean;
-  wipAssistantMessage?: MessageUI;
+  wipAssistantMessage?: string;
 }
 
 export const ChatMessage: React.FC<ChatMessageProps> = ({
@@ -48,8 +44,8 @@ export const ChatMessage: React.FC<ChatMessageProps> = ({
   const [messageId, setMessageId] = useAtom(selectedMessageId);
   const { id, mediaUrl, prompt, response, result } = message;
   const [formattedSections, finalResult, finalError] = useMemo(
-    () => formatStreamLogs(response ?? wipAssistantMessage?.content),
-    [response, wipAssistantMessage?.content],
+    () => formatStreamLogs(response ?? wipAssistantMessage),
+    [response, wipAssistantMessage],
   );
   // prioritize the result from the message over the WIP message
   const codeResult = result?.payload ?? finalResult;
@@ -100,10 +96,10 @@ export const ChatMessage: React.FC<ChatMessageProps> = ({
             <div className="flex-1 px-1 space-y-4 ml-4 overflow-hidden">
               <Table className="w-[400px]">
                 <TableBody>
-                  {formattedSections.map(section => (
+                  {formattedSections.map((section, index) => (
                     <TableRow
                       className="border-primary/50 h-[56px]"
-                      key={section.type}
+                      key={index}
                     >
                       <TableCell className="text-center text-webkit-center">
                         {ChunkStatusToIconDict[section.status]}
@@ -131,13 +127,11 @@ export const ChatMessage: React.FC<ChatMessageProps> = ({
               )}
               {!codeResult && finalError && (
                 <>
-                  <p>❌ Coding failed </p>
+                  <p>❌ {finalError.name}</p>
                   <div>
                     <CodeBlock
                       language="error"
                       value={
-                        finalError.name +
-                        '\n' +
                         finalError.value +
                         '\n' +
                         finalError.traceback_raw.join('\n')
@@ -181,7 +175,7 @@ const ChunkTypeToText: React.FC<{
   const isExecuting = !['completed', 'failed'].includes(status);
 
   useEffect(() => {
-    if (isExecuting && timestamp && !useTimer) {
+    if (isExecuting && timestamp && useTimer) {
       const timerId = setInterval(() => {
         setMSeconds(Date.now() - Date.parse(timestamp));
       }, 200);
@@ -272,7 +266,7 @@ const ChunkPayloadAction: React.FC<{
         </DialogContent>
       </Dialog>
     );
-  } else {
+  } else if ((payload as PrismaJson.FinalChatResult['payload']).code) {
     return (
       <Dialog>
         <DialogTrigger asChild>
@@ -281,11 +275,14 @@ const ChunkPayloadAction: React.FC<{
           </Button>
         </DialogTrigger>
         <DialogContent className="max-w-5xl">
-          <CodeResultDisplay codeResult={payload as CodeResult} />
+          <CodeResultDisplay
+            codeResult={payload as PrismaJson.FinalChatResult['payload']}
+          />
         </DialogContent>
       </Dialog>
     );
   }
+  return null;
 };
 
 export default ChatMessage;

components/chat/TopPrompt.tsx

@@ -13,7 +13,6 @@ export interface TopPrompt {
 
 export default async function TopPrompt({ chat, userId }: TopPrompt) {
   const authorId = chat.userId;
-  console.log('[Ming] ~ TopPrompt ~ authorId:', authorId);
   // 1. Viewer logged in, Viewer = Author
   if (userId && authorId === userId) {
     return null;
@@ -26,15 +25,11 @@ export default async function TopPrompt({ chat, userId }: TopPrompt) {
   if (authorId && authorId !== userId) {
     const chatAuthor = authorId ? await dbGetUser(authorId) : null;
     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">
+      <Card className="group py-2 px-4 flex flex-row items-center">
+        <p className="leading-normal text-sm">Authored by</p>
+        <div className="flex-1 px-1 ml-2 flex flex-row items-center space-x-2">
           <Avatar name={chatAuthor?.name} avatar={chatAuthor?.avatar} />
-        </div>
-        <div className="flex-1 px-1 ml-2 overflow-hidden">
-          <p className="leading-normal">
-            Code author:{' '}
-            <span className="font-medium">{chatAuthor?.name ?? 'Unknown'}</span>
-          </p>
+          <p className="font-medium">{chatAuthor?.name ?? 'Unknown'}</p>
         </div>
       </Card>
     );

components/ui/CodeBlock.tsx

@@ -47,7 +47,7 @@ export const programmingLanguages: languageMap = {
   css: '.css',
   // custom titles
   print: '.txt',
-  error: 'vim',
+  error: '.txt',
   // add more file extensions here, make sure the key is same as language prop in CodeBlock.tsx component
 };
 

components/ui/Icons.tsx

@@ -222,7 +222,7 @@ function IconCrossCircle({ className, ...props }: React.ComponentProps<'svg'>) {
   return (
     <svg
       height="16"
-      stroke-linejoin="round"
+      strokeLinejoin="round"
       viewBox="0 0 16 16"
       width="16"
       className={cn('size-4', className)}

lib/db/prisma.ts

@@ -4,24 +4,35 @@ declare global {
   var prisma: PrismaClient | undefined;
   namespace PrismaJson {
     // you can use classes, interfaces, types, etc.
+
+    type StructuredResult = {
+      logs: {
+        stderr: string[];
+        stdout: string[];
+      };
+      results: Array<{
+        png?: string;
+        mp4?: string;
+        text: string;
+        is_main_result: boolean;
+      }>;
+      error: {
+        name: string;
+        value: string;
+        traceback_raw: string[];
+      };
+    };
+
     type FinalChatResult = {
       type: 'final_code';
       status: 'completed' | 'failed';
       payload: {
         code: string;
         test: string;
-        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;
-        //   }>;
-        // };
+        // Change introduces https://github.com/landing-ai/vision-agent-ui/pull/86
+        // 1. Backward compatibility, it could be stringified StructuredResult
+        // 2. result not modified in stream server, could still be stringified StructuredResult
+        result: string | StructuredResult;
       };
     };
   }

lib/hooks/useVisionAgent.ts

@@ -2,13 +2,10 @@ import { useChat } from 'ai/react';
 import { toast } from 'react-hot-toast';
 import { useEffect, useRef, useState } from 'react';
 import { ChatWithMessages, MessageUI, MessageUserInput } from '../types';
-import {
-  dbPostCreateMessage,
-  dbPostUpdateMessageResponse,
-} from '../db/functions';
+import { dbPostUpdateMessageResponse } from '../db/functions';
 import {
   convertAssistantUIMessageToDBMessageResponse,
-  convertDBMessageToUIMessage,
+  convertDBMessageToAPIMessage,
 } from '../utils/message';
 import { useSetAtom } from 'jotai';
 import { selectedMessageId } from '@/state/chat';
@@ -38,10 +35,11 @@ const useVisionAgent = (chat: ChatWithMessages) => {
       );
       setMessageId(currMessageId.current);
     },
-    initialMessages: convertDBMessageToUIMessage(dbMessages),
     body: {
       mediaUrl: currMediaUrl.current,
       id,
+      // for some reason, the messages has to be stringified to be sent to the API
+      apiMessages: JSON.stringify(convertDBMessageToAPIMessage(dbMessages)),
     },
     onError: err => {
       err && toast.error(err.message);

lib/types.ts

@@ -16,21 +16,3 @@ 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;
-  }>;
-  error: {
-    name: string;
-    value: string;
-    traceback_raw: string[];
-  };
-};

lib/utils/content.ts

@@ -1,26 +1,4 @@
 import toast from 'react-hot-toast';
-import { ResultPayload } from '../types';
-const ANSWERS_PREFIX = 'answers';
-
-export const generateAnswersImageMarkdown = (index: number, url: string) => {
-  return `![${ANSWERS_PREFIX}-${index}](${url})`;
-};
-
-export const cleanInputMessage = (content: string) => {
-  return content
-    .replace(/!\[input-.*?\)/g, '')
-    .replace(/<video[^>]*>.*?<\/video>/g, '');
-};
-
-export const cleanAnswerMessage = (content: string) => {
-  return content.replace(/!\[answers.*?\.png\)/g, '');
-};
-
-export type CodeResult = {
-  code: string;
-  test: string;
-  result: string;
-};
 
 const WIPLogTypes = ['plans', 'tools', 'code'];
 const AllLogTypes = [
@@ -37,8 +15,8 @@ export type ChunkBody = {
   timestamp?: string;
   payload:
     | Array<Record<string, string>> // PlansBody | ToolsBody
-    | CodeResult // CodeBody
-    | ResultPayload['error']; // ErrorBody
+    | PrismaJson.FinalChatResult['payload'] // CodeBody & FinalCodeBody
+    | PrismaJson.StructuredResult['error']; // ErrorBody
 };
 
 export type WIPChunkBodyGroup = ChunkBody & {
@@ -53,7 +31,11 @@ export type WIPChunkBodyGroup = ChunkBody & {
  */
 export const formatStreamLogs = (
   content: string | null | undefined,
-): [WIPChunkBodyGroup[], CodeResult?, ResultPayload['error']?] => {
+): [
+  WIPChunkBodyGroup[],
+  PrismaJson.FinalChatResult['payload']?,
+  PrismaJson.StructuredResult['error']?,
+] => {
   if (!content) return [[], undefined];
   const streamLogs = content.split('\n').filter(log => !!log);
 
@@ -104,8 +86,8 @@ export const formatStreamLogs = (
   return [
     groupedSections.filter(section => WIPLogTypes.includes(section.type)),
     groupedSections.find(section => section.type === 'final_code')
-      ?.payload as CodeResult,
+      ?.payload as PrismaJson.FinalChatResult['payload'],
     groupedSections.find(section => section.type === 'final_error')
-      ?.payload as ResultPayload['error'],
+      ?.payload as PrismaJson.StructuredResult['error'],
   ];
 };

lib/utils/message.ts

@@ -4,13 +4,13 @@ import { ChunkBody } from './content';
 
 /**
  * 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)
+ * for the API to use, we need to break them to 2 messages, User and Assistant(if responded)
  */
-export const convertDBMessageToUIMessage = (
+export const convertDBMessageToAPIMessage = (
   messages: Message[],
 ): MessageUI[] => {
   return messages.reduce((acc, message) => {
-    const { id, mediaUrl, prompt, response, result } = message;
+    const { id, mediaUrl, prompt, result } = message;
     if (mediaUrl && prompt) {
       acc.push({
         id: id + '-user',
@@ -18,11 +18,11 @@ export const convertDBMessageToUIMessage = (
         content: prompt,
       });
     }
-    if (response) {
+    if (result) {
       acc.push({
         id: id + '-assistant',
         role: 'assistant',
-        content: response,
+        content: result?.payload.code,
       });
     }
     return acc;