feat: add timeout and remove png from code result (#82)

tested locally and see the final result:
<img width="1441" alt="image"
src="https://github.com/landing-ai/vision-agent-ui/assets/132986242/2c015adc-9793-4325-ab13-31eeb65c652b">

app/api/vision-agent/route.ts

@@ -11,6 +11,7 @@ import { getPresignedUrl } from '@/lib/aws';
 // 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 = 5 * 60 * 1000;
 
 const uploadBase64 = async (
   base64: string,
@@ -91,7 +92,7 @@ export const POST = withLogging(
 
     const fetchResponse = await fetch(
       `https://api.dev.landing.ai/v1/agent/chat?agent_class=vision_agent&self_reflection=false`,
-      // `http://localhost:5001/v1/agent/chat?agent_class=vision_agent&self_reflection=${enableSelfReflection}`,
+      // `http://localhost:5001/v1/agent/chat?agent_class=vision_agent&self_reflection=false`,
       {
         method: 'POST',
         headers: {
@@ -148,6 +149,7 @@ export const POST = withLogging(
       const decoder = new TextDecoder('utf-8');
       let maxChunkSize = 0;
       let buffer = '';
+      let time = Date.now();
       const stream = new ReadableStream({
         async start(controller) {
           // const parser = createParser(streamParser);
@@ -158,6 +160,23 @@ export const POST = withLogging(
             const lines = buffer
               .split('\n')
               .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__',
+                );
+                controller.error(
+                  `Haven't received any response in last ${TIMEOUT_MILI_SECONDS / 60000} minutes, agent timed out.`,
+                );
+              }
+            } else {
+              time = Date.now();
+            }
             buffer = lines.pop() ?? ''; // Save the last incomplete line back to the buffer
             let done = false;
             const parseLine = async (
@@ -166,9 +185,32 @@ export const POST = withLogging(
             ) => {
               try {
                 const msg = JSON.parse(line);
-                if (msg.type !== 'final_code') {
+                if (
+                  msg.type !== 'final_code' &&
+                  (msg.type !== 'code' ||
+                    msg.status === 'started' ||
+                    msg.status === 'running')
+                ) {
                   return line;
                 }
+                if (msg.type === 'code') {
+                  const result = JSON.parse(
+                    msg.payload.result,
+                  ) as ResultPayload;
+                  if (result && result.results) {
+                    msg.payload.result = JSON.stringify({
+                      ...result,
+                      results: result.results.map((_result: any) => {
+                        return {
+                          ..._result,
+                          png: undefined,
+                          mp4: undefined,
+                        };
+                      }),
+                    });
+                  }
+                  return JSON.stringify(msg);
+                }
                 const result = JSON.parse(msg.payload.result) as ResultPayload;
                 for (let index = 0; index < result.results.length; index++) {
                   const png = result.results[index].png ?? '';