feat: support image rendering on answer (#19)

https://app.asana.com/0/1204554785675703/1207155143920392/f

<img width="769" alt="image"
src="https://github.com/landing-ai/vision-agent-ui/assets/132986242/e20eac86-73ec-4dc2-8764-afb0f0c153da">

app/api/sign/route.ts

@@ -16,14 +16,13 @@ export async function POST(req: Request): Promise<Response> {
   // }
 
   try {
-    const { fileName, fileType } = (await req.json()) as {
+    const { fileName, fileType, id } = (await req.json()) as {
+      id?: string;
       fileName: string;
       fileType: string;
     };
 
-    const id = nanoid();
-
-    const signedFileName = `${user}/${id}/${fileName}`;
+    const signedFileName = `${user}/${id ?? nanoid()}/${fileName}`;
     const res = await getPresignedUrl(signedFileName, fileType);
     return Response.json({
       id,

app/api/vision-agent/route.ts

@@ -27,7 +27,7 @@ export async function POST(req: Request) {
   formData.append('image', url);
 
   const fetchResponse = await fetch(
-    'https://api.dev.landing.ai/v1/agent/chat?agent_class=vision_agent',
+    'https://api.dev.landing.ai/v1/agent/chat?agent_class=vision_agent&visualize_output=true',
     // 'http://localhost:5050/v1/agent/chat?agent_class=vision_agent',
     {
       method: 'POST',

components/chat/index.tsx

@@ -20,7 +20,6 @@ export function Chat({ chat }: ChatProps) {
 
   const { messagesRef, scrollRef, visibilityRef, isAtBottom, scrollToBottom } =
     useScrollAnchor();
-  console.log('[Ming] ~ Chat ~ isAtBottom:', isAtBottom);
 
   return (
     <>

lib/hooks/useCleanedUpMessages.ts

@@ -1,5 +1,6 @@
-import { useMemo } from 'react';
-import { MessageBase } from '../types';
+import { useMemo, useEffect } from 'react';
+import { MessageBase, SignedPayload } from '../types';
+import { fetcher } from '../utils';
 
 const PAIRS: Record<string, string> = {
   '┍': '┑',
@@ -10,54 +11,68 @@ const PAIRS: Record<string, string> = {
 
 const MIDDLE_STARTER = '┝';
 const MIDDLE_SEPARATOR = '┿';
+export const CLEANED_SEPARATOR = '|CLEANED|';
 
-export const useCleanedUpMessages = ({ content, role }: MessageBase) => {
-  return useMemo(() => {
-    if (role === 'user') {
-      return {
-        content,
-      };
-    }
-    const [logs = '', answer = ''] = content.split('<ANSWER>');
-    const cleanedLogs = [];
-    let left = 0;
-    let right = 0;
-    while (right < logs.length) {
-      if (Object.keys(PAIRS).includes(content[right])) {
-        cleanedLogs.push(content.substring(left, right));
-        left = right++;
-        while (
-          right < content.length &&
-          PAIRS[content[left]] !== content[right]
-        ) {
-          right++;
-        }
-        if (content[left] === MIDDLE_STARTER) {
-          // add the text alignment so it can be shown as a table
-          const separators = logs
-            .substring(left, right)
-            .split(MIDDLE_SEPARATOR).length;
-          if (separators > 0) {
-            cleanedLogs.push(
-              Array(separators + 1)
-                .fill('|')
-                .join(' :- '),
-            );
-          }
-        }
-        left = ++right;
-      } else {
+export const getCleanedUpMessages = ({
+  content,
+  role,
+}: Pick<MessageBase, 'role' | 'content'>) => {
+  if (role === 'user') {
+    return {
+      content,
+    };
+  }
+  if (content.split(CLEANED_SEPARATOR).length === 2) {
+    return {
+      logs: content.split(CLEANED_SEPARATOR)[0],
+      content: content.split(CLEANED_SEPARATOR)[1],
+    };
+  }
+  const [logs = '', answer = ''] = content.split('<ANSWER>');
+  const cleanedLogs = [];
+  let left = 0;
+  let right = 0;
+  while (right < logs.length) {
+    if (Object.keys(PAIRS).includes(content[right])) {
+      cleanedLogs.push(content.substring(left, right));
+      left = right++;
+      while (
+        right < content.length &&
+        PAIRS[content[left]] !== content[right]
+      ) {
         right++;
       }
+      if (content[left] === MIDDLE_STARTER) {
+        // add the text alignment so it can be shown as a table
+        const separators = logs
+          .substring(left, right)
+          .split(MIDDLE_SEPARATOR).length;
+        if (separators > 0) {
+          cleanedLogs.push(
+            Array(separators + 1)
+              .fill('|')
+              .join(' :- '),
+          );
+        }
+      }
+      left = ++right;
+    } else {
+      right++;
     }
-    cleanedLogs.push(content.substring(left, right));
-    return {
-      logs: cleanedLogs
-        .join('')
-        .replace(/│/g, '|')
-        .split('|\n\n|')
-        .join('|\n|'),
-      content: answer.replace('</</ANSWER>', '').replace('</ANSWER>', ''),
-    };
+  }
+  cleanedLogs.push(content.substring(left, right));
+  const [answerText, imagesStr = ''] = answer.split('<VIZ>');
+  const images = imagesStr.split('</IMG>').map(str => str.replace('<IMG>', ''));
+  return {
+    logs: cleanedLogs.join('').replace(/│/g, '|').split('|\n\n|').join('|\n|'),
+    content: answerText.replace('</</ANSWER>', '').replace('</ANSWER>', ''),
+    images: images.slice(0, -1),
+  };
+};
+
+export const useCleanedUpMessages = ({ content, role }: MessageBase) => {
+  const cleanedMessage = useMemo(() => {
+    return getCleanedUpMessages({ content, role });
   }, [content, role]);
+  return cleanedMessage;
 };

lib/hooks/useVisionAgent.tsx

@@ -1,8 +1,49 @@
 import { useChat, type Message, UseChatHelpers } from 'ai/react';
 import { toast } from 'react-hot-toast';
 import { useEffect, useState } from 'react';
-import { ChatEntity, MessageBase } from '../types';
+import { ChatEntity, MessageBase, SignedPayload } from '../types';
 import { saveKVChatMessage } from '../kv/chat';
+import { fetcher } from '../utils';
+import {
+  getCleanedUpMessages,
+  CLEANED_SEPARATOR,
+} from './useCleanedUpMessages';
+
+const uploadBase64 = async (
+  base64: string,
+  messageId: string,
+  chatId: string,
+  index: number,
+) => {
+  const res = await fetch('data:image/png;base64,' + base64);
+  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');
+  }
+};
 
 const useVisionAgent = (chat: ChatEntity) => {
   const { messages: initialMessages, id, url } = chat;
@@ -23,8 +64,32 @@ const useVisionAgent = (chat: ChatEntity) => {
         toast.error(response.statusText);
       }
     },
-    onFinish(message) {
-      saveKVChatMessage(id, message);
+    onFinish: async message => {
+      const { logs = '', content, images } = getCleanedUpMessages(message);
+      if (images?.length) {
+        const publicUrls = await Promise.all(
+          images.map((image, index) =>
+            uploadBase64(image, message.id, id, index),
+          ),
+        );
+        const newMessage = {
+          ...message,
+          content:
+            logs +
+            CLEANED_SEPARATOR +
+            content +
+            '\n' +
+            publicUrls
+              .map((url, index) => `![image-${index}](${url})`)
+              .join('\n'),
+        };
+        saveKVChatMessage(id, newMessage);
+      } else {
+        saveKVChatMessage(id, {
+          ...message,
+          content: logs + CLEANED_SEPARATOR + content,
+        });
+      }
     },
     initialMessages: initialMessages,
     body: {