feat: support png and mp4 rendering (#73)

<img width="1164" alt="image"
src="https://github.com/landing-ai/vision-agent-ui/assets/132986242/2311e390-fa0c-4acf-bcb7-5da1f7d34c36">

app/api/sign/route.ts

@@ -25,14 +25,8 @@ export const POST = withLogging(
     try {
       const { fileName, fileType, id = nanoid() } = json;
 
-      const signedFileName = `${user}/${id}/${fileName}`;
-      const res = await getPresignedUrl(signedFileName, fileType);
-      return Response.json({
-        id,
-        signedUrl: res.url,
-        publicUrl: `https://${process.env.AWS_BUCKET_NAME}.s3.${process.env.AWS_REGION}.amazonaws.com/${signedFileName}`,
-        fields: res.fields,
-      });
+      const res = await getPresignedUrl(fileName, fileType, id, user);
+      return Response.json(res);
     } catch (error) {
       return new Response((error as Error).message, {
         status: 400,

app/api/vision-agent/route.ts

@@ -6,7 +6,7 @@ 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 { fetcher } from '@/lib/utils';
+import { getPresignedUrl } from '@/lib/aws';
 
 // export const runtime = 'edge';
 export const dynamic = 'force-dynamic';
@@ -17,21 +17,15 @@ const uploadBase64 = async (
   messageId: string,
   chatId: string,
   index: number,
+  user: string,
 ) => {
-  const res = await fetch(
-    'data:image/png;base64,' + base64.replace('base:64', ''),
-  );
+  const res = await fetch(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 { signedUrl, publicUrl, fields } = await getPresignedUrl(
+    `answer-${index}.${blob.type.split('/')[1]}`,
+    blob.type,
+    `${chatId}/${messageId}`,
+    user,
   );
   const formData = new FormData();
   Object.entries(fields).forEach(([key, value]) => {
@@ -61,6 +55,7 @@ export const POST = withLogging(
     request,
   ) => {
     const { messages, mediaUrl } = json;
+    const user = session?.user?.email ?? 'anonymous';
 
     // const session = await auth();
     // if (!session?.user?.email) {
@@ -152,60 +147,85 @@ export const POST = withLogging(
       const encoder = new TextEncoder();
       const decoder = new TextDecoder('utf-8');
       let maxChunkSize = 0;
+      let buffer = '';
       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 = data.split('\n');
-            const results = [];
+            const lines = buffer
+              .split('\n')
+              .filter(line => line.trim().length > 0);
+            buffer = lines.pop() ?? ''; // Save the last incomplete line back to the buffer
             let done = false;
-            for (let line of lines) {
-              if (!line.trim()) {
-                continue;
-              }
+            const parseLine = async (
+              line: string,
+              errorCallback?: (e: Error) => void,
+            ) => {
               try {
-                const msg = JSON.parse(line.trim());
+                const msg = JSON.parse(line);
                 if (msg.type !== 'final_code') {
-                  results.push(line);
-                  continue;
+                  return line;
                 }
                 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 png = result.results[index].png ?? '';
+                  const mp4 = result.results[index].mp4 ?? '';
+                  if (!png && !mp4) continue;
                   const resp = await uploadBase64(
-                    png,
+                    png
+                      ? 'data:image/png;base64,' + png
+                      : 'data:video/mp4;base64,' + mp4,
                     messages[messages.length - 1].id,
                     json.id,
                     index,
+                    user,
                   );
-                  result.results[index].png = resp;
+                  if (png) result.results[index].png = resp;
+                  if (mp4) result.results[index].mp4 = resp;
                 }
                 msg.payload.result = JSON.stringify(result);
-                results.push(JSON.stringify(msg));
                 done = true;
+                return JSON.stringify(msg);
               } catch (e) {
+                errorCallback?.(e as Error);
+              }
+            };
+            for (let line of lines) {
+              if (!line.trim()) {
+                continue;
+              }
+              const parsedLine = await parseLine(line, (e: Error) => {
                 console.error(e);
                 logger.error(
                   session,
                   {
-                    message: (e as Error).message,
+                    line,
+                    message: e.message,
                   },
                   request,
                 );
                 controller.error(e);
+              });
+              controller.enqueue(
+                encoder.encode(
+                  parsedLine?.trim() ? parsedLine?.trim() + '\n' : '',
+                ),
+              );
+            }
+            if (buffer) {
+              const parsedBuffer = await parseLine(buffer);
+              if (parsedBuffer?.trim()) {
+                buffer = '';
+                controller.enqueue(encoder.encode(parsedBuffer.trim() + '\n'));
               }
             }
-            controller.enqueue(
-              encoder.encode(
-                results.length === 0 ? '' : results.join('\n') + '\n',
-              ),
-            );
             if (done) {
+              console.log(done);
               logger.info(
                 session,
                 {

components/chat/ChatMessage.tsx

@@ -225,10 +225,44 @@ const CodeResultDisplay: React.FC<{
           <CodeBlock language="print" value={stdout.join('').trim()} />
         </>
       )}
-      {!!results.length && (
+      {Array.isArray(results) && !!results.length && (
         <>
           <Separator />
-          <CodeBlock language="output" value={results} />
+          {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 />

lib/aws.ts

@@ -9,10 +9,16 @@ const s3Client = new S3Client({
   credentials: fromEnv(),
 });
 
-export const getPresignedUrl = async (fileName: string, fileType: string) => {
-  return createPresignedPost(s3Client, {
+export const getPresignedUrl = async (
+  fileName: string,
+  fileType: string,
+  id: string,
+  user: string,
+) => {
+  const signedFileName = `${user}/${id}/${fileName}`;
+  const res = await createPresignedPost(s3Client, {
     Bucket: process.env.AWS_BUCKET_NAME ?? 'vision-agent-dev',
-    Key: fileName,
+    Key: signedFileName,
     Conditions: [
       ['content-length-range', 0, FILE_SIZE_LIMIT],
       ['starts-with', '$Content-Type', fileType],
@@ -23,24 +29,10 @@ export const getPresignedUrl = async (fileName: string, fileType: string) => {
     },
     Expires: 600,
   });
-};
-
-export const upload = async (
-  base64: string,
-  fileName: string,
-  fileType: string,
-) => {
-  const { url, fields } = await getPresignedUrl(fileName, fileType);
-  const formData = new FormData();
-  Object.entries(fields).forEach(([key, value]) => {
-    formData.append(key, value as string);
-  });
-  const res = await fetch(base64);
-  const blob = await res.blob();
-  formData.append('file', blob);
-
-  return fetch(url, {
-    method: 'POST',
-    body: formData,
-  });
+  return {
+    id,
+    signedUrl: res.url,
+    publicUrl: `https://${process.env.AWS_BUCKET_NAME}.s3.${process.env.AWS_REGION}.amazonaws.com/${signedFileName}`,
+    fields: res.fields,
+  };
 };

lib/db/prisma.ts

@@ -17,6 +17,7 @@ declare global {
           };
           results: Array<{
             png?: string;
+            mp4?: string;
             text: string;
             is_main_result: boolean;
           }>;

next.config.js

@@ -12,10 +12,8 @@ module.exports = {
   },
   experimental: {
     serverActions: {
-      bodySizeLimit: '10mb',
+      bodySizeLimit: '30mb',
     },
-  },
-  experimental: {
     serverComponentsExternalPackages: ['pino', 'pino-loki'],
   },
   ...(process.env.USE_STANDALONE_BUILD ? { output: 'standalone' } : {}),