done

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,
@@ -92,7 +93,7 @@ export const POST = withLogging(
     const fetchResponse = await fetch(
       // `https://api.dev.landing.ai/v1/agent/chat?agent_class=vision_agent&self_reflection=false`,
       `https://api.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: {
@@ -150,6 +151,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);
@@ -160,6 +162,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 (
@@ -168,9 +187,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 ?? '';

components/chat/ChatMessage.tsx

@@ -29,7 +29,6 @@ import { selectedMessageId } from '@/state/chat';
 import { Message } from '@prisma/client';
 import { Separator } from '../ui/Separator';
 import { cn } from '@/lib/utils';
-import { usePrevious } from '@/lib/hooks/usePrevious';
 
 export interface ChatMessageProps {
   message: Message;