feat: add responseBody field and serving data to hooks (#93)

New chat loaded properly and refresh still loaded:
<img width="1477" alt="image"
src="https://github.com/landing-ai/vision-agent-ui/assets/132986242/5fa60f22-e609-49e4-9a80-5ffde16bce8c">

Previous chat can still be loaded properly:
<img width="1500" alt="image"
src="https://github.com/landing-ai/vision-agent-ui/assets/132986242/e526e549-711c-433d-a70f-459c8af7bb9b">

app/api/vision-agent/route.ts

@@ -1,16 +1,17 @@
-import { JSONValue, StreamingTextResponse, experimental_StreamData } from 'ai';
+import { StreamingTextResponse } from 'ai';
 
 // import { auth } from '@/auth';
-import { MessageUI, SignedPayload } from '@/lib/types';
+import { MessageUI } from '@/lib/types';
 
 import { logger, withLogging } from '@/lib/logger';
 import { getPresignedUrl } from '@/lib/aws';
+import { dbPostUpdateMessageResponse } from '@/lib/db/functions';
 
 // export const runtime = 'edge';
 export const dynamic = 'force-dynamic';
 export const maxDuration = 300; // This function can run for a maximum of 5 minutes
 const TIMEOUT_MILI_SECONDS = 2 * 60 * 1000;
-const FINAL_TIMEOUT_ERROR = {
+const FINAL_TIMEOUT_ERROR: PrismaJson.FinalErrorBody = {
   type: 'final_error',
   status: 'failed',
   payload: {
@@ -53,20 +54,25 @@ const uploadBase64 = async (
 };
 
 const modifyCodePayload = async (
-  msg: Record<string, any>,
+  msg: PrismaJson.MessageBody,
   messageId: string,
   chatId: string,
   user: string,
-) => {
+): Promise<PrismaJson.MessageBody> => {
   if (
-    msg.type !== 'final_code' &&
-    (msg.type !== 'code' ||
-      msg.status === 'started' ||
-      msg.status === 'running')
+    (msg.type !== 'final_code' &&
+      (msg.type !== 'code' ||
+        msg.status === 'started' ||
+        msg.status === 'running')) ||
+    !msg.payload?.result
   ) {
     return msg;
   }
-  const result = JSON.parse(msg.payload.result) as PrismaJson.StructuredResult;
+  const result = (
+    typeof msg.payload.result === 'string'
+      ? JSON.parse(msg.payload.result)
+      : msg.payload.result
+  ) as PrismaJson.StructuredResult;
   if (msg.type === 'code') {
     if (result && result.results) {
       msg.payload.result = {
@@ -105,17 +111,18 @@ export const POST = withLogging(
     session,
     json: {
       apiMessages: string;
-      id: string;
+      chatId: string;
+      messageId: string;
       mediaUrl: string;
     },
     request,
   ) => {
-    const { apiMessages, mediaUrl } = json;
+    const { apiMessages, mediaUrl, chatId, messageId } = json;
     const messages: MessageUI[] = JSON.parse(apiMessages);
     const user = session?.user?.email ?? 'anonymous';
 
     const formData = new FormData();
-    formData.append('input', JSON.stringify(messages));
+    formData.append('input', apiMessages);
     formData.append('image', mediaUrl);
 
     const agentHost = process.env.LND_TIER
@@ -171,102 +178,147 @@ export const POST = withLogging(
     let maxChunkSize = 0;
     let buffer = '';
     let time = Date.now();
-    let done = false;
+    const results: PrismaJson.MessageBody[] = [];
     const stream = new ReadableStream({
       async start(controller) {
-        // const parser = createParser(streamParser);
-        for await (const chunk of fetchResponse.body as any) {
-          const data = decoder.decode(chunk);
-          buffer += data;
-          maxChunkSize = Math.max(data.length, maxChunkSize);
-          const lines = buffer
-            .split('\n')
-            .filter(line => line.trim().length > 0);
+        const parseLine = async (
+          line: string,
+          ignoreParsingError = false,
+        ): Promise<{ data?: PrismaJson.MessageBody; error?: Error }> => {
+          let msg = null;
+          try {
+            msg = JSON.parse(line);
+          } catch (e) {
+            if (ignoreParsingError) return {};
+            else {
+              return { error: e as Error };
+            }
+          }
+          if (!msg) return {};
+          try {
+            const modifiedMsg = await modifyCodePayload(
+              {
+                ...msg,
+                timestamp: new Date(),
+              },
+              messageId,
+              chatId,
+              user,
+            );
+            return { data: modifiedMsg };
+          } catch (e) {
+            return { error: e as Error };
+          }
+        };
+
+        const processChunk = async (lines: string[]) => {
           if (lines.length === 0) {
             if (Date.now() - time > TIMEOUT_MILI_SECONDS) {
-              // logger.info(
-              //   session,
-              //   {
-              //     message: 'Agent timed out',
-              //   },
-              //   request,
-              //   '__Agent_timeout__',
-              // );
+              results.push(FINAL_TIMEOUT_ERROR);
+              // https://github.com/vercel/ai/blob/f7002ad2c5aa58ce6ed83e8d31fe22f71ebdb7d7/packages/ui-utils/src/stream-parts.ts#L62
               controller.enqueue(
-                encoder.encode(JSON.stringify(FINAL_TIMEOUT_ERROR) + '\n'),
+                '2:' +
+                  encoder.encode(JSON.stringify(FINAL_TIMEOUT_ERROR) + '\n'),
               );
-              controller.close();
-              return;
+              return { done: true, reason: 'timeout' };
             }
           } else {
             time = Date.now();
           }
           buffer = lines.pop() ?? ''; // Save the last incomplete line back to the buffer
-          const parseLine = async (
-            line: string,
-            ignoreParsingError = false,
-          ) => {
-            const handleError = (e: Error) => {
-              console.error(e);
-              logger.error(
-                session,
-                {
-                  line,
-                  message: e.message,
+          for (let line of lines) {
+            const { data: parsedMsg, error } = await parseLine(line);
+            if (error) {
+              results.push({
+                type: 'final_error',
+                status: 'failed',
+                payload: {
+                  name: 'ParsingError',
+                  value: line,
+                  traceback_raw: [],
                 },
-                request,
+              });
+              return { done: true, reason: 'api_error', error };
+            } else if (parsedMsg) {
+              results.push(parsedMsg);
+              controller.enqueue(
+                encoder.encode('2:' + JSON.stringify([parsedMsg]) + '\n'),
               );
-              controller.error(e);
-            };
-            let msg = null;
-            try {
-              msg = JSON.parse(line);
-              if (msg.type === 'final_code' || msg.type === 'final_error') {
-                done = true;
+              if (parsedMsg.type === 'final_code') {
+                return { done: true, reason: 'agent_concluded' };
+              } else if (parsedMsg.type === 'final_error') {
+                return {
+                  done: true,
+                  reason: 'agent_error',
+                  error: parsedMsg.payload,
+                };
               }
-            } catch (e) {
-              if (ignoreParsingError) return;
-              handleError(e as Error);
-            }
-            if (!msg) return;
-            try {
-              const modifiedMsg = await modifyCodePayload(
-                {
-                  ...msg,
-                  timestamp: new Date(),
-                },
-                messages[messages.length - 1].id,
-                json.id,
-                user,
-              );
-              return modifiedMsg;
-            } catch (e) {
-              handleError(e as Error);
+            } else {
+              controller.enqueue(encoder.encode(''));
             }
-          };
-          for (let line of lines) {
-            const parsedMsg = await parseLine(line);
-            controller.enqueue(
-              parsedMsg ? encoder.encode(JSON.stringify(parsedMsg) + '\n') : '',
-            );
           }
           if (buffer) {
-            const parsedBuffer = await parseLine(buffer, true);
-            if (parsedBuffer) {
+            const { data: parsedBuffer, error } = await parseLine(buffer, true);
+            if (error) {
+              results.push({
+                type: 'final_error',
+                status: 'failed',
+                payload: {
+                  name: 'ParsingError',
+                  value: buffer,
+                  traceback_raw: [],
+                },
+              });
+              return { done: true, reason: 'api_error', error };
+            } else if (parsedBuffer) {
               buffer = '';
+              results.push(parsedBuffer);
               controller.enqueue(
-                encoder.encode(JSON.stringify(parsedBuffer) + '\n'),
+                encoder.encode('2:' + JSON.stringify([parsedBuffer]) + '\n'),
               );
+              if (parsedBuffer.type === 'final_code') {
+                return { done: true, reason: 'agent_concluded' };
+              } else if (parsedBuffer.type === 'final_error') {
+                return {
+                  done: true,
+                  reason: 'agent_error',
+                  error: parsedBuffer.payload,
+                };
+              }
             } else {
-              controller.enqueue('');
+              controller.enqueue(encoder.encode(''));
             }
           }
+          return { done: false };
+        };
+
+        // const parser = createParser(streamParser);
+        for await (const chunk of fetchResponse.body as any) {
+          const data = decoder.decode(chunk);
+          buffer += data;
+          maxChunkSize = Math.max(data.length, maxChunkSize);
+          const lines = buffer
+            .split('\n')
+            .filter(line => line.trim().length > 0);
+          const { done, reason, error } = await processChunk(lines);
           if (done) {
+            const processMsgs = results.filter(
+              res => res.type !== 'final_code',
+            ) as PrismaJson.AgentResponseBodies;
+            await dbPostUpdateMessageResponse(messageId, {
+              response: processMsgs.map(res => JSON.stringify(res)).join('\n'),
+              result: results.find(
+                res => res.type === 'final_code',
+              ) as PrismaJson.FinalCodeBody,
+              responseBody: processMsgs,
+            });
             logger.info(
               session,
               {
-                message: 'Streaming finished',
+                message: 'Streaming ended',
                 maxChunkSize,
+                reason,
+                error,
               },
               request,
               '__AGENT_DONE',

components/CodeResultDisplay.tsx

@@ -23,7 +23,7 @@ import {
 export interface CodeResultDisplayProps {}
 
 const CodeResultDisplay: React.FC<{
-  codeResult: PrismaJson.FinalChatResult['payload'];
+  codeResult: PrismaJson.FinalCodeBody['payload'];
 }> = ({ codeResult }) => {
   const { code, test, result } = codeResult;
   const getDetail = () => {

components/chat/ChatList.tsx

@@ -4,9 +4,8 @@
 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 { ChatWithMessages } from '@/lib/types';
 import { ChatMessage } from './ChatMessage';
 import { Button } from '../ui/Button';
 import { cn } from '@/lib/utils';
@@ -25,16 +24,11 @@ export const SCROLL_BOTTOM = 120;
 
 const ChatList: React.FC<ChatListProps> = ({ chat, userId }) => {
   const { id, messages: dbMessages, userId: chatUserId } = chat;
-  const { messages, append, isLoading } = useVisionAgent(chat);
+  const { append, isLoading, data } = useVisionAgent(chat);
 
   // Only login and chat owner can compose
   const canCompose = !chatUserId || userId === chatUserId;
 
-  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);
 
@@ -63,11 +57,7 @@ const ChatList: React.FC<ChatListProps> = ({ chat, userId }) => {
               key={message.id}
               message={message}
               loading={isLastMessage && isLoading}
-              wipAssistantMessage={
-                lastAssistantMessage && isLastMessage
-                  ? lastAssistantMessage
-                  : undefined
-              }
+              wipAssistantMessage={data}
             />
           );
         })}

components/chat/ChatMessage.tsx

@@ -33,7 +33,7 @@ import { cn } from '@/lib/utils';
 export interface ChatMessageProps {
   message: Message;
   loading?: boolean;
-  wipAssistantMessage?: string;
+  wipAssistantMessage?: PrismaJson.MessageBody[];
 }
 
 export const ChatMessage: React.FC<ChatMessageProps> = ({
@@ -42,10 +42,15 @@ export const ChatMessage: React.FC<ChatMessageProps> = ({
   loading,
 }) => {
   const [messageId, setMessageId] = useAtom(selectedMessageId);
-  const { id, mediaUrl, prompt, response, result } = message;
+  const { id, mediaUrl, prompt, response, result, responseBody } = message;
   const [formattedSections, finalResult, finalError] = useMemo(
-    () => formatStreamLogs(response ?? wipAssistantMessage),
-    [response, wipAssistantMessage],
+    () =>
+      formatStreamLogs(
+        responseBody ??
+          wipAssistantMessage ??
+          (response ? JSON.parse(response) : ''),
+      ),
+    [response, wipAssistantMessage, result, responseBody],
   );
   // prioritize the result from the message over the WIP message
   const codeResult = result?.payload ?? finalResult;
@@ -266,7 +271,7 @@ const ChunkPayloadAction: React.FC<{
         </DialogContent>
       </Dialog>
     );
-  } else if ((payload as PrismaJson.FinalChatResult['payload']).code) {
+  } else if ((payload as PrismaJson.FinalCodeBody['payload']).code) {
     return (
       <Dialog>
         <DialogTrigger asChild>
@@ -276,7 +281,7 @@ const ChunkPayloadAction: React.FC<{
         </DialogTrigger>
         <DialogContent className="max-w-5xl">
           <CodeResultDisplay
-            codeResult={payload as PrismaJson.FinalChatResult['payload']}
+            codeResult={payload as PrismaJson.FinalCodeBody['payload']}
           />
         </DialogContent>
       </Dialog>

lib/db/functions.ts

@@ -188,6 +188,7 @@ export async function dbPostUpdateMessageResponse(
     data: {
       response: messageResponse.response,
       result: messageResponse.result ?? undefined,
+      responseBody: messageResponse.responseBody,
     },
     where: {
       id: messageId,

lib/db/prisma.ts

@@ -5,6 +5,12 @@ declare global {
   namespace PrismaJson {
     // you can use classes, interfaces, types, etc.
 
+    interface StructuredError {
+      name: string;
+      value: string;
+      traceback_raw: string[];
+    }
+
     type StructuredResult = {
       logs: {
         stderr: string[];
@@ -16,25 +22,50 @@ declare global {
         text: string;
         is_main_result: boolean;
       }>;
-      error: {
-        name: string;
-        value: string;
-        traceback_raw: string[];
-      };
+      error?: StructuredError;
     };
 
-    type FinalChatResult = {
+    interface PlanAndToolsBody {
+      type: 'plans' | 'tools';
+      status: 'started' | 'completed';
+      payload?: Array<Record<string, string>>;
+    }
+
+    interface CodeBody {
+      type: 'code';
+      status: 'completed' | 'failed' | 'started' | 'running';
+      payload?: {
+        code: string;
+        test: string;
+        result?: StructuredResult | string;
+      };
+    }
+
+    interface FinalCodeBody {
       type: 'final_code';
       status: 'completed' | 'failed';
       payload: {
         code: string;
         test: string;
-        // 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;
+        result: StructuredResult | string;
       };
-    };
+    }
+
+    interface FinalErrorBody {
+      type: 'final_error';
+      status: 'failed';
+      payload: StructuredError;
+    }
+
+    type MessageBody =
+      | PlanAndToolsBody
+      | CodeBody
+      | FinalErrorBody
+      | FinalCodeBody;
+
+    type AgentResponseBodies = Array<
+      PlanAndToolsBody | CodeBody | FinalErrorBody
+    >;
   }
 }
 

lib/hooks/useVisionAgent.ts

@@ -1,12 +1,8 @@
 import { useChat } from 'ai/react';
 import { toast } from 'react-hot-toast';
-import { useEffect, useRef, useState } from 'react';
-import { ChatWithMessages, MessageUI, MessageUserInput } from '../types';
-import { dbPostUpdateMessageResponse } from '../db/functions';
-import {
-  convertAssistantUIMessageToDBMessageResponse,
-  convertDBMessageToAPIMessage,
-} from '../utils/message';
+import { useEffect, useRef } from 'react';
+import { ChatWithMessages, MessageUI } from '../types';
+import { convertDBMessageToAPIMessage } from '../utils/message';
 import { useSetAtom } from 'jotai';
 import { selectedMessageId } from '@/state/chat';
 import { Message } from '@prisma/client';
@@ -20,24 +16,20 @@ const useVisionAgent = (chat: ChatWithMessages) => {
   const currMediaUrl = useRef<string>(mediaUrl);
   const currMessageId = useRef<string>(latestDbMessage?.id);
 
-  const { messages, append, isLoading, reload } = useChat({
+  const { append, isLoading, data, reload } = useChat({
     api: '/api/vision-agent',
-    streamMode: 'text',
     onResponse(response) {
       if (response.status !== 200) {
         toast.error(response.statusText);
       }
     },
-    onFinish: async message => {
-      await dbPostUpdateMessageResponse(
-        currMessageId.current,
-        convertAssistantUIMessageToDBMessageResponse(message),
-      );
+    onFinish: message => {
       setMessageId(currMessageId.current);
     },
     body: {
       mediaUrl: currMediaUrl.current,
-      id,
+      chatId: id,
+      messageId: currMessageId.current,
       // for some reason, the messages has to be stringified to be sent to the API
       apiMessages: JSON.stringify(convertDBMessageToAPIMessage(dbMessages)),
     },
@@ -56,10 +48,9 @@ const useVisionAgent = (chat: ChatWithMessages) => {
       once.current = false;
       reload();
     }
-  }, [isLoading, latestDbMessage.result, messages, reload]);
+  }, [isLoading, latestDbMessage.result, reload]);
 
   return {
-    messages: messages as MessageUI[],
     append: (message: Message) => {
       currMediaUrl.current = message.mediaUrl;
       currMessageId.current = message.id;
@@ -69,6 +60,7 @@ const useVisionAgent = (chat: ChatWithMessages) => {
         content: message.prompt,
       });
     },
+    data: data as unknown as PrismaJson.MessageBody[],
     reload,
     isLoading,
   };

lib/types.ts

@@ -4,9 +4,11 @@ import { type Message as MessageAI } from 'ai';
 export type ChatWithMessages = Chat & { messages: Message[] };
 
 export type MessageUserInput = Pick<Message, 'prompt' | 'mediaUrl'>;
-export type MessageAssistantResponse = Partial<
-  Pick<Message, 'response' | 'result'>
->;
+export type MessageAssistantResponse = {
+  result?: PrismaJson.FinalCodeBody;
+  response: string;
+  responseBody: PrismaJson.AgentResponseBodies;
+};
 
 export type MessageUI = Pick<MessageAI, 'role' | 'content' | 'id'>;
 

lib/utils/content.ts

@@ -15,11 +15,12 @@ export type ChunkBody = {
   timestamp?: string;
   payload:
     | Array<Record<string, string>> // PlansBody | ToolsBody
-    | PrismaJson.FinalChatResult['payload'] // CodeBody & FinalCodeBody
-    | PrismaJson.StructuredResult['error']; // ErrorBody
+    | PrismaJson.FinalCodeBody['payload'] // CodeBody & FinalCodeBody
+    | PrismaJson.StructuredError; // ErrorBody
 };
 
-export type WIPChunkBodyGroup = ChunkBody & {
+export type WIPChunkBodyGroup = PrismaJson.MessageBody & {
+  timestamp?: string;
   duration?: number;
 };
 
@@ -30,37 +31,16 @@ export type WIPChunkBodyGroup = ChunkBody & {
  * @returns An array of grouped sections and an optional final code result.
  */
 export const formatStreamLogs = (
-  content: string | null | undefined,
+  content: WIPChunkBodyGroup[] | null,
 ): [
   WIPChunkBodyGroup[],
-  PrismaJson.FinalChatResult['payload']?,
-  PrismaJson.StructuredResult['error']?,
+  PrismaJson.FinalCodeBody['payload']?,
+  PrismaJson.StructuredError?,
 ] => {
-  if (!content) return [[], undefined];
-  const streamLogs = content.split('\n').filter(log => !!log);
-
-  const buffer = streamLogs.pop();
-  const parsedStreamLogs: ChunkBody[] = [];
-  try {
-    streamLogs.forEach(streamLog =>
-      parsedStreamLogs.push(JSON.parse(streamLog)),
-    );
-  } catch {
-    toast.error('Error parsing stream logs');
-    return [[], undefined];
-  }
-
-  if (buffer) {
-    try {
-      const lastLog = JSON.parse(buffer);
-      parsedStreamLogs.push(lastLog);
-    } catch {
-      console.log(buffer);
-    }
-  }
+  if (!content) return [[], undefined, undefined];
 
   // Merge consecutive logs of the same type to the latest status
-  const groupedSections = parsedStreamLogs.reduce((acc, curr) => {
+  const groupedSections = content.reduce((acc: WIPChunkBodyGroup[], curr) => {
     const lastGroup = acc[acc.length - 1];
     if (
       acc.length > 0 &&
@@ -81,13 +61,13 @@ export const formatStreamLogs = (
       acc.push(curr);
     }
     return acc;
-  }, [] as WIPChunkBodyGroup[]);
+  }, []);
 
   return [
     groupedSections.filter(section => WIPLogTypes.includes(section.type)),
     groupedSections.find(section => section.type === 'final_code')
-      ?.payload as PrismaJson.FinalChatResult['payload'],
+      ?.payload as PrismaJson.FinalCodeBody['payload'],
     groupedSections.find(section => section.type === 'final_error')
-      ?.payload as PrismaJson.StructuredResult['error'],
+      ?.payload as PrismaJson.StructuredError,
   ];
 };

lib/utils/message.ts

@@ -1,6 +1,5 @@
 import { Message } from '@prisma/client';
 import { MessageAssistantResponse, MessageUI } from '../types';
-import { ChunkBody } from './content';
 
 /**
  * The Message we saved to database consists of a prompt and a response
@@ -28,25 +27,3 @@ export const convertDBMessageToAPIMessage = (
     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/migrations/20240612024011_add_response_body_in_message/migration.sql

@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "message" ADD COLUMN     "responseBody" JSONB;

prisma/schema.prisma

@@ -47,8 +47,10 @@ model Message {
   mediaUrl  String
   prompt    String
   response  String?
-  /// [FinalChatResult]
+  /// [FinalCodeBody]
   result    Json?
+  /// [AgentResponseBodies]
+  responseBody Json?
   chat      Chat     @relation(fields: [chatId], references: [id], onDelete: Cascade)
   user      User?    @relation(fields: [userId], references: [id])