feat: change db message format and parser (#66)

<img width="767" alt="image"
src="https://github.com/landing-ai/vision-agent-ui/assets/132986242/070987c0-994b-4223-a030-a05d352c8b4f">

app/api/vision-agent/route.ts

@@ -1,15 +1,54 @@
 import { StreamingTextResponse } from 'ai';
 
 // import { auth } from '@/auth';
-import { MessageBase } from '../../../lib/types';
+import { MessageBase, SignedPayload } from '../../../lib/types';
 import { logger, withLogging } from '@/lib/logger';
 import { CLEANED_SEPARATOR } from '@/lib/constants';
 import { cleanAnswerMessage, cleanInputMessage } from '@/lib/messageUtils';
+import { fetcher } from '@/lib/utils';
 
 // export const runtime = 'edge';
 export const dynamic = 'force-dynamic';
 export const maxDuration = 300; // This function can run for a maximum of 5 minutes
 
+const uploadBase64 = async (
+  base64: string,
+  messageId: string,
+  chatId: string,
+  index: number,
+) => {
+  const res = await fetch(
+    'data:image/png;base64,' + base64.replace('base:64', ''),
+  );
+  const blob = await res.blob();
+  const { signedUrl, publicUrl, fields } = await fetcher<SignedPayload>(
+    '/api/sign',
+    {
+      method: 'POST',
+      body: JSON.stringify({
+        id: `${chatId}/${messageId}`,
+        fileType: blob.type,
+        fileName: `answer-${index}.${blob.type.split('/')[1]}`,
+      }),
+    },
+  );
+  const formData = new FormData();
+  Object.entries(fields).forEach(([key, value]) => {
+    formData.append(key, value as string);
+  });
+  formData.append('file', blob);
+
+  const uploadResponse = await fetch(signedUrl, {
+    method: 'POST',
+    body: formData,
+  });
+  if (uploadResponse.ok) {
+    return publicUrl;
+  } else {
+    throw new Error('Upload failed');
+  }
+};
+
 export const POST = withLogging(
   async (
     session,
@@ -110,7 +149,6 @@ export const POST = withLogging(
     if (fetchResponse.body) {
       const encoder = new TextEncoder();
       const decoder = new TextDecoder('utf-8');
-      let buffer = '';
       let maxChunkSize = 0;
       const stream = fetchResponse.body.pipeThrough(
         new TransformStream({
@@ -128,7 +166,50 @@ export const POST = withLogging(
           transform: async (chunk, controller) => {
             const data = decoder.decode(chunk, { stream: true });
             maxChunkSize = Math.max(data.length, maxChunkSize);
-            controller.enqueue(encoder.encode(data));
+            const lines = data.split('\n');
+            const results = [];
+            for (let line of lines) {
+              if (!line.trim()) {
+                continue;
+              }
+              try {
+                const msg = JSON.parse(line.trim());
+                if (msg.type !== 'final_code') {
+                  results.push(line);
+                  continue;
+                }
+                const result = JSON.parse(
+                  msg.payload.result,
+                ) as PrismaJson.FinalChatResult['payload']['result'];
+                for (let index = 0; index < result.results.length; index++) {
+                  const png = result.results[index].png;
+                  if (!png) continue;
+                  const resp = await uploadBase64(
+                    png,
+                    messages[messages.length - 1].id,
+                    json.id,
+                    index,
+                  );
+                  result.results[index].png = resp;
+                }
+                msg.payload.result = JSON.stringify(result);
+                results.push(JSON.stringify(msg));
+              } catch (e) {
+                console.error(e);
+                logger.error(
+                  session,
+                  {
+                    message: (e as Error).message,
+                  },
+                  request,
+                );
+              }
+            }
+            controller.enqueue(
+              encoder.encode(
+                results.length === 0 ? '' : results.join('\n') + '\n',
+              ),
+            );
           },
         }),
       );

app/page.tsx

@@ -56,6 +56,7 @@ export default function Page() {
                       (mediaUrl
                         ? '\n\n' + generateInputImageMarkdown(mediaUrl)
                         : ''),
+                    result: null,
                   },
                 ],
               });

components/chat/ChatMessage.tsx

@@ -10,9 +10,9 @@ import { cn } from '@/lib/utils';
 import { CodeBlock } from '@/components/ui/CodeBlock';
 import { MemoizedReactMarkdown } from '@/components/chat/MemoizedReactMarkdown';
 import { IconLandingAI, IconUser } from '@/components/ui/Icons';
-import { MessageBase } from '../../lib/types';
+import type { ChatMessage } from '@/lib/db/types';
 import Img from '../ui/Img';
-import { getCleanedUpMessages } from '@/lib/messageUtils';
+import { getFormattedMessage } from '@/lib/messageUtils';
 import Loading from '../ui/Loading';
 import {
   Table,
@@ -23,9 +23,11 @@ import {
 } from '../ui/Table';
 import { Button } from '../ui/Button';
 import { Dialog, DialogContent } from '../ui/Dialog';
+import { Message } from 'ai';
+import { MessageRole } from '@prisma/client';
 
 export interface ChatMessageProps {
-  message: MessageBase;
+  message: Message;
   isLoading: boolean;
 }
 
@@ -140,9 +142,9 @@ const Markdown: React.FC<{
 
 export function ChatMessage({ message, isLoading }: ChatMessageProps) {
   const { content } = useMemo(() => {
-    return getCleanedUpMessages({
+    return getFormattedMessage({
       content: message.content,
-      role: message.role,
+      role: message.role as MessageRole,
     });
   }, [message.content, message.role]);
   const [details, setDetails] = useState<string>('');

lib/db/functions.ts

@@ -119,6 +119,7 @@ export async function dbPostCreateChat({
           create: initMessages.map(message => ({
             ...message,
             ...userConnect,
+            result: undefined,
           })),
         },
       },

lib/db/prisma.ts

@@ -2,6 +2,28 @@ import { PrismaClient } from '@prisma/client';
 
 declare global {
   var prisma: PrismaClient | undefined;
+  namespace PrismaJson {
+    // you can use classes, interfaces, types, etc.
+    type FinalChatResult = {
+      type: 'final_code';
+      status: 'completed' | 'failed';
+      payload: {
+        code: string;
+        test: string;
+        result: {
+          logs: {
+            stderr: string[];
+            stdout: string[];
+          };
+          results: Array<{
+            png?: string;
+            text: string;
+            is_main_result: boolean;
+          }>;
+        };
+      };
+    };
+  }
 }
 
 const db = globalThis.prisma || new PrismaClient();

lib/db/types.ts

@@ -1,8 +1,10 @@
 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,55 +1,13 @@
-import { useChat, type Message, UseChatHelpers } from 'ai/react';
+import { useChat, UseChatHelpers } from 'ai/react';
 import { toast } from 'react-hot-toast';
-import { useEffect, useState } from 'react';
-import { MessageBase, SignedPayload } from '../types';
-import { fetcher, nanoid } from '../utils';
-import {
-  getCleanedUpMessages,
-  generateAnswersImageMarkdown,
-  generateInputImageMarkdown,
-} from '../messageUtils';
-import { CLEANED_SEPARATOR } from '../constants';
+import { useEffect } from 'react';
 import { useSearchParams } from 'next/navigation';
-import { ChatWithMessages, MessageRaw } from '../db/types';
+import { ChatWithMessages } from '../db/types';
 import { dbPostCreateMessage } from '../db/functions';
-
-const uploadBase64 = async (
-  base64: string,
-  messageId: string,
-  chatId: string,
-  index: number,
-) => {
-  const res = await fetch(
-    'data:image/png;base64,' + base64.replace('base:64', ''),
-  );
-  const blob = await res.blob();
-  const { signedUrl, publicUrl, fields } = await fetcher<SignedPayload>(
-    '/api/sign',
-    {
-      method: 'POST',
-      body: JSON.stringify({
-        id: `${chatId}/${messageId}`,
-        fileType: blob.type,
-        fileName: `answer-${index}.${blob.type.split('/')[1]}`,
-      }),
-    },
-  );
-  const formData = new FormData();
-  Object.entries(fields).forEach(([key, value]) => {
-    formData.append(key, value as string);
-  });
-  formData.append('file', blob);
-
-  const uploadResponse = await fetch(signedUrl, {
-    method: 'POST',
-    body: formData,
-  });
-  if (uploadResponse.ok) {
-    return publicUrl;
-  } else {
-    throw new Error('Upload failed');
-  }
-};
+import {
+  convertDbMessageToMessage,
+  convertMessageToDbMessage,
+} from '../messageUtils';
 
 const useVisionAgent = (chat: ChatWithMessages) => {
   const { messages: initialMessages, id, mediaUrl } = chat;
@@ -69,16 +27,16 @@ const useVisionAgent = (chat: ChatWithMessages) => {
       }
     },
     onFinish: async message => {
-      await dbPostCreateMessage(id, {
-        role: message.role as 'user' | 'assistant',
-        content: message.content,
-      });
+      await dbPostCreateMessage(id, convertMessageToDbMessage(message));
     },
-    initialMessages: initialMessages,
+    initialMessages: initialMessages.map(convertDbMessageToMessage),
     body: {
       mediaUrl,
       id,
     },
+    onError: err => {
+      err && toast.error(err.message);
+    },
   });
 
   /**
@@ -94,12 +52,13 @@ const useVisionAgent = (chat: ChatWithMessages) => {
     dbPostCreateMessage(id, {
       role: message.role as 'user' | 'assistant',
       content: message.content,
+      result: null,
     });
     return appendRaw(message);
   };
 
   return {
-    messages: messages as MessageBase[],
+    messages,
     append,
     reload,
     isLoading,

lib/messageUtils.ts

@@ -1,5 +1,5 @@
-import { MessageBase } from './types';
-import { CLEANED_SEPARATOR } from './constants';
+import type { ChatMessage } from '@/lib/db/types';
+import { Message } from 'ai';
 
 const PAIRS: Record<string, string> = {
   '┍': '┑',
@@ -89,31 +89,30 @@ const generateCodeExecutionMarkdown = (
 };
 
 const generateFinalCodeMarkdown = (
-  message: string,
-  payload: FinalCodeBody['payload'],
+  code: string,
+  test: string,
+  result: PrismaJson.FinalChatResult['payload']['result'],
 ) => {
-  message += 'Final Code: \n';
-  message += `\`\`\`python\n${payload.code}\n\`\`\`\n`;
+  let message = 'Final Code: \n';
+  message += `\`\`\`python\n${code}\n\`\`\`\n`;
   message += 'Final test: \n';
-  message += `\`\`\`python\n${payload.test}\n\`\`\`\n`;
+  message += `\`\`\`python\n${test}\n\`\`\`\n`;
   message += `Final result: \n`;
-  message += `\`\`\`\n${payload.result}\n\`\`\`\n`;
-  return message;
-};
-
-const generateFinalResultMarkdown = (
-  message: string,
-  payload: {
-    success: boolean;
-    reach_max_retries: boolean;
-  },
-) => {
-  if (payload.reach_max_retries) {
-    message += 'Reach max debug retries!\n';
-  } else if (payload.success) {
-    message += 'Success!';
+  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`;
   }
-  message += '\n';
   return message;
 };
 
@@ -162,26 +161,6 @@ type CodeBody =
       };
     };
 
-type FinalCodeBody = {
-  type: 'final_code';
-  status: 'completed' | 'failed';
-  payload: {
-    code: string;
-    test: string;
-    result: {
-      logs: {
-        stderr: string[];
-        stdout: string[];
-      };
-      results: Array<{
-        png: string;
-        text: string;
-        is_main_result: boolean;
-      }>;
-    };
-  };
-};
-
 // this will return if self_reflection flag is true
 type ReflectionBody =
   | {
@@ -199,7 +178,7 @@ type MessageBody =
   | ToolsBody
   | CodeBody
   | ReflectionBody
-  | FinalCodeBody;
+  | PrismaJson.FinalChatResult;
 
 const getMessageTitle = (json: MessageBody) => {
   switch (json.type) {
@@ -258,16 +237,16 @@ const parseLine = (json: MessageBody) => {
     case 'self_reflection':
       return generateJSONArrayMarkdown(title, [json.payload]);
     case 'final_code':
-      return generateFinalCodeMarkdown(title, json.payload);
+      return '';
     default:
       throw 'Not supported type';
   }
 };
 
-export const getCleanedUpMessages = ({
+export const getFormattedMessage = ({
   content,
   role,
-}: Pick<MessageBase, 'role' | 'content'>) => {
+}: Pick<ChatMessage, 'role' | 'content'>) => {
   if (role === 'user') {
     return {
       content,
@@ -275,14 +254,25 @@ export const getCleanedUpMessages = ({
   }
   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 (
+      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'
@@ -298,6 +288,40 @@ export const getCleanedUpMessages = ({
   }
   jsons.forEach(json => (formattedLogs += parseLine(json)));
   return {
-    content: formattedLogs,
+    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,
   };
 };

package.json

@@ -46,6 +46,7 @@
     "pino": "^9.0.0",
     "pino-loki": "^2.2.1",
     "prisma": "^5.14.0",
+    "prisma-json-types-generator": "^3.0.4",
     "react": "^18.2.0",
     "react-dom": "^18.2.0",
     "react-dropzone": "^14.2.3",

pnpm-lock.yaml

@@ -104,6 +104,9 @@ importers:
       prisma:
         specifier: ^5.14.0
         version: 5.14.0
+      prisma-json-types-generator:
+        specifier: ^3.0.4
+        version: 3.0.4([email protected])([email protected])
       react:
         specifier: ^18.2.0
         version: 18.2.0
@@ -736,6 +739,9 @@ packages:
   '@prisma/[email protected]':
     resolution: {integrity: sha512-iq56qBZuFfX3fCxoxT8gBX33lQzomBU0qIUaEj1RebsKVz1ob/BVH1XSBwwwvRVtZEV1b7Fxx2eVu34Ge/mg3w==}
 
+  '@prisma/[email protected]':
+    resolution: {integrity: sha512-yAHFSFCg8KVoL0oRUno3m60GAjsUKYUDkQ+9BA2X2JfVR3kRVSJFc/GpQ2fSORi4pSHZR9orfM4UC9OVXIFFTA==}
+
   '@prisma/engines-version@5.14.0-25.e9771e62de70f79a5e1c604a2d7c8e2a0a874b48':
     resolution: {integrity: sha512-ip6pNkRo1UxWv+6toxNcYvItNYaqQjXdFNGJ+Nuk2eYtRoEdoF13wxo7/jsClJFFenMPVNVqXQDV0oveXnR1cA==}
 
@@ -745,6 +751,9 @@ packages:
   '@prisma/[email protected]':
     resolution: {integrity: sha512-VrheA9y9DMURK5vu8OJoOgQpxOhas3qF0IBHJ8G/0X44k82kc8E0w98HCn2nhnbOOMwbWsJWXfLC2/F8n5u0gQ==}
 
+  '@prisma/[email protected]':
+    resolution: {integrity: sha512-WMdEUPpPYxUGruRQM6e6IVTWXFjt1hHdF/m2TO7pWxhPo7/ZeoTOF9fH8JsvVSV78DYLOQkx9osjFLXZu447Kw==}
+
   '@prisma/[email protected]':
     resolution: {integrity: sha512-/yAyBvcEjRv41ynZrhdrPtHgk47xLRRq/o5eWGcUpBJ1YrUZTYB8EoPiopnP7iQrMATK8stXQdPOoVlrzuTQZw==}
 
@@ -3153,6 +3162,14 @@ packages:
   [email protected]:
     resolution: {integrity: sha512-WuxUnVtlWL1OfZFQFuqvnvs6MiAGk9UNsBostyBOB0Is9wb5uRESevA6rnl/rkksXaGX3GzZhPup5d6Vp1nFew==}
 
+  [email protected]:
+    resolution: {integrity: sha512-W53OpjBdGZxCsYv7MlUX69d7TPA9lEsQbDf9ddF0J93FX5EvaIRDMexdFPe0KTxiuquGvZTDJgeNXb3gIqEhJw==}
+    engines: {node: '>=14.0'}
+    hasBin: true
+    peerDependencies:
+      prisma: ^5.1
+      typescript: ^5.1
+
   [email protected]:
     resolution: {integrity: sha512-gCNZco7y5XtjrnQYeDJTiVZmT/ncqCr5RY1/Cf8X2wgLRmyh9ayPAGBNziI4qEE4S6SxCH5omQLVo9lmURaJ/Q==}
     engines: {node: '>=16.13'}
@@ -4764,6 +4781,8 @@ snapshots:
 
   '@prisma/[email protected]': {}
 
+  '@prisma/[email protected]': {}
+
   '@prisma/engines-version@5.14.0-25.e9771e62de70f79a5e1c604a2d7c8e2a0a874b48': {}
 
   '@prisma/[email protected]':
@@ -4779,6 +4798,10 @@ snapshots:
       '@prisma/engines-version': 5.14.0-25.e9771e62de70f79a5e1c604a2d7c8e2a0a874b48
       '@prisma/get-platform': 5.14.0
 
+  '@prisma/[email protected]':
+    dependencies:
+      '@prisma/debug': 5.9.1
+
   '@prisma/[email protected]':
     dependencies:
       '@prisma/debug': 5.14.0
@@ -7724,6 +7747,13 @@ snapshots:
 
   [email protected]: {}
 
+  [email protected]([email protected])([email protected]):
+    dependencies:
+      '@prisma/generator-helper': 5.9.1
+      prisma: 5.14.0
+      tslib: 2.6.2
+      typescript: 5.4.5
+
   [email protected]:
     dependencies:
       '@prisma/engines': 5.14.0

prisma/schema.prisma

@@ -2,6 +2,10 @@ generator client {
   provider = "prisma-client-js"
 }
 
+generator json {
+  provider = "prisma-json-types-generator"
+}
+
 datasource db {
   provider  = "postgresql"
   url       = env("POSTGRES_PRISMA_URL")
@@ -40,6 +44,8 @@ model Message {
   userId    String?
   chatId    String
   content   String
+  /// [FinalChatResult]
+  result    Json?
   role      MessageRole
   chat      Chat        @relation(fields: [chatId], references: [id], onDelete: Cascade)
   user      User?       @relation(fields: [userId], references: [id])