feat: add timestamp to msg body (#86)

and refactor the agent endpoint code a bit

<img width="1443" alt="image"
src="https://github.com/landing-ai/vision-agent-ui/assets/132986242/bd79c37a-3fd4-47cb-a0ca-16dfe7c2c2ab">

app/api/vision-agent/route.ts

@@ -1,4 +1,4 @@
-import { StreamingTextResponse, experimental_StreamData } from 'ai';
+import { JSONValue, StreamingTextResponse, experimental_StreamData } from 'ai';
 
 // import { auth } from '@/auth';
 import { MessageUI, ResultPayload, SignedPayload } from '@/lib/types';
@@ -12,6 +12,15 @@ import { getPresignedUrl } from '@/lib/aws';
 export const dynamic = 'force-dynamic';
 export const maxDuration = 300; // This function can run for a maximum of 5 minutes
 const TIMEOUT_MILI_SECONDS = 2 * 60 * 1000;
+const FINAL_TIMEOUT_ERROR = {
+  type: 'final_error',
+  status: 'failed',
+  payload: {
+    name: 'AgentTimeout',
+    value: `Haven't received any response in last ${TIMEOUT_MILI_SECONDS / 60000} minutes, agent timed out.`,
+    traceback_raw: [],
+  },
+};
 
 const uploadBase64 = async (
   base64: string,
@@ -45,6 +54,54 @@ const uploadBase64 = async (
   }
 };
 
+const modifyCodePayload = async (
+  msg: Record<string, any>,
+  messageId: string,
+  chatId: string,
+  user: string,
+) => {
+  if (
+    msg.type !== 'final_code' &&
+    (msg.type !== 'code' ||
+      msg.status === 'started' ||
+      msg.status === 'running')
+  ) {
+    return msg;
+  }
+  const result = JSON.parse(msg.payload.result) as ResultPayload;
+  if (msg.type === 'code') {
+    if (result && result.results) {
+      msg.payload.result = {
+        ...result,
+        results: result.results.map((_result: any) => {
+          return {
+            ..._result,
+            png: undefined,
+            mp4: undefined,
+          };
+        }),
+      };
+    }
+    return msg;
+  }
+  for (let index = 0; index < result.results.length; index++) {
+    const png = result.results[index].png ?? '';
+    const mp4 = result.results[index].mp4 ?? '';
+    if (!png && !mp4) continue;
+    const resp = await uploadBase64(
+      png ? 'data:image/png;base64,' + png : 'data:video/mp4;base64,' + mp4,
+      messageId,
+      chatId,
+      index,
+      user,
+    );
+    if (png) result.results[index].png = resp;
+    if (mp4) result.results[index].mp4 = resp;
+  }
+  msg.payload.result = result;
+  return msg;
+};
+
 export const POST = withLogging(
   async (
     session,
@@ -90,220 +147,168 @@ export const POST = withLogging(
     );
     formData.append('image', mediaUrl);
 
-    const agentHost = process.env.LND_TIER ? 'http://publicrestapi-app-lndsvc.publicrestapi.svc.cluster.local:5000' : 'https://api.dev.landing.ai';
+    const agentHost = process.env.LND_TIER
+      ? 'http://publicrestapi-app-lndsvc.publicrestapi.svc.cluster.local:5000'
+      : 'https://api.dev.landing.ai';
 
     const fetchResponse = await fetch(
       `${agentHost}/v1/agent/chat?agent_class=vision_agent&self_reflection=false`,
-      // `https://api.landing.ai/v1/agent/chat?agent_class=vision_agent&self_reflection=false`,
+      // `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=false`,
       {
         method: 'POST',
         headers: {
           // default to dev apikey
-          apikey: process.env.LND_TIER === 'production' ? 'land_sk_nMnUf8xiJJUjyw1l5QaIJJ4ZyrvPthzVmPAIG7TtJY7F9CW6lu' : 'land_sk_DKeoYtaZZrYqJ9TMMiXe4BIQgJcZ0s3XAoB0JT3jv73FFqnr6k'
+          apikey:
+            process.env.LND_TIER === 'production'
+              ? 'land_sk_nMnUf8xiJJUjyw1l5QaIJJ4ZyrvPthzVmPAIG7TtJY7F9CW6lu'
+              : 'land_sk_DKeoYtaZZrYqJ9TMMiXe4BIQgJcZ0s3XAoB0JT3jv73FFqnr6k',
         },
         body: formData,
       },
     );
 
-    if (!fetchResponse.ok) {
-      if (fetchResponse.body) {
-        const reader = fetchResponse.body.getReader();
-        return new StreamingTextResponse(
-          new ReadableStream({
-            async start(controller) {
-              try {
-                const { done, value } = await reader.read();
-                if (!done) {
-                  const errorText = new TextDecoder().decode(value);
-                  logger.error(session, { message: errorText }, request);
-                  controller.error(new Error(`Response error: ${errorText}`));
-                }
-              } catch (e) {
-                logger.error(session, (e as Error).message, request);
+    if (!fetchResponse.ok && fetchResponse.body) {
+      const reader = fetchResponse.body.getReader();
+      return new StreamingTextResponse(
+        new ReadableStream({
+          async start(controller) {
+            try {
+              const { done, value } = await reader?.read();
+              if (!done) {
+                const errorText = new TextDecoder().decode(value);
+                logger.error(session, { message: errorText }, request);
+                controller.error(new Error(`Response error: ${errorText}`));
               }
-            },
-          }),
-          {
-            status: 400,
-          },
-        );
-      } else {
-        return new StreamingTextResponse(
-          new ReadableStream({
-            start(controller) {
-              logger.error(
-                session,
-                { message: 'Response error: No response body' },
-                request,
-              );
-              controller.error(new Error('Response error: No response body'));
-            },
-          }),
-          {
-            status: 400,
+            } catch (e) {
+              logger.error(session, (e as Error).message, request);
+            }
           },
-        );
-      }
+        }),
+        {
+          status: 400,
+        },
+      );
     }
     // const streamData = new experimental_StreamData();
 
-    if (fetchResponse.body) {
-      const encoder = new TextEncoder();
-      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);
-          for await (const chunk of fetchResponse.body as any) {
-            const data = decoder.decode(chunk);
-            buffer += data;
-            maxChunkSize = Math.max(data.length, maxChunkSize);
-            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.enqueue(
-                  encoder.encode(
-                    JSON.stringify({
-                      type: 'final_error',
-                      status: 'failed',
-                      payload: {
-                        name: 'AgentTimeout',
-                        value: `Haven't received any response in last ${TIMEOUT_MILI_SECONDS / 60000} minutes, agent timed out.`,
-                        traceback_raw: [],
-                      },
-                    }) + '\n',
-                  ),
-                );
-                controller.close();
-                return;
-              }
-            } else {
-              time = Date.now();
-            }
-            buffer = lines.pop() ?? ''; // Save the last incomplete line back to the buffer
-            let done = false;
-            const parseLine = async (
-              line: string,
-              errorCallback?: (e: Error) => void,
-            ) => {
-              try {
-                const msg = JSON.parse(line);
-                if (msg.type === 'final_code' || msg.type === 'final_error') {
-                  done = true;
-                }
-                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 ?? '';
-                  const mp4 = result.results[index].mp4 ?? '';
-                  if (!png && !mp4) continue;
-                  const resp = await uploadBase64(
-                    png
-                      ? 'data:image/png;base64,' + png
-                      : 'data:video/mp4;base64,' + mp4,
-                    messages[messages.length - 1].id,
-                    json.id,
-                    index,
-                    user,
-                  );
-                  if (png) result.results[index].png = resp;
-                  if (mp4) result.results[index].mp4 = resp;
-                }
-                msg.payload.result = JSON.stringify(result);
-                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,
-                  {
-                    line,
-                    message: e.message,
-                  },
-                  request,
-                );
-                controller.error(e);
-              });
+    if (!fetchResponse.body) {
+      return fetchResponse;
+    }
+    const encoder = new TextEncoder();
+    const decoder = new TextDecoder('utf-8');
+    let maxChunkSize = 0;
+    let buffer = '';
+    let time = Date.now();
+    let done = false;
+    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 = 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.enqueue(
-                encoder.encode(
-                  parsedLine?.trim() ? parsedLine?.trim() + '\n' : '',
-                ),
+                encoder.encode(JSON.stringify(FINAL_TIMEOUT_ERROR) + '\n'),
               );
+              controller.close();
+              return;
             }
-            if (buffer) {
-              const parsedBuffer = await parseLine(buffer);
-              if (parsedBuffer?.trim()) {
-                buffer = '';
-                controller.enqueue(encoder.encode(parsedBuffer.trim() + '\n'));
-              }
-            }
-            if (done) {
-              logger.info(
+          } else {
+            time = Date.now();
+          }
+          buffer = lines.pop() ?? ''; // Save the last incomplete line back to the buffer
+          const parseLine = async (
+            line: string,
+            ignoreParsingError = false,
+          ) => {
+            const handleError = (e: Error) => {
+              console.error(e);
+              logger.error(
                 session,
                 {
-                  message: 'Streaming finished',
-                  maxChunkSize,
+                  line,
+                  message: e.message,
                 },
                 request,
-                '__AGENT_DONE',
               );
-              controller.close();
+              controller.error(e);
+            };
+            let msg = null;
+            try {
+              msg = JSON.parse(line);
+              if (msg.type === 'final_code' || msg.type === 'final_error') {
+                done = true;
+              }
+            } catch (e) {
+              if (ignoreParsingError) return;
+              handleError(e as Error);
             }
+            if (!msg) return;
+            try {
+              const modifiedMsg = await modifyCodePayload(
+                {
+                  ...msg,
+                  timestamp: new Date(),
+                },
+                messages[messages.length - 1].id,
+                json.id,
+                user,
+              );
+              return modifiedMsg;
+            } catch (e) {
+              handleError(e as Error);
+            }
+          };
+          for (let line of lines) {
+            const parsedMsg = await parseLine(line);
+            controller.enqueue(
+              parsedMsg ? encoder.encode(JSON.stringify(parsedMsg) + '\n') : '',
+            );
           }
-        },
-      });
-      return new Response(stream, {
-        headers: {
-          'Content-Type': 'application/x-ndjson',
-        },
-      });
-    } else {
-      return fetchResponse;
-    }
+          if (buffer) {
+            const parsedBuffer = await parseLine(buffer, true);
+            if (parsedBuffer) {
+              buffer = '';
+              controller.enqueue(
+                encoder.encode(JSON.stringify(parsedBuffer) + '\n'),
+              );
+            } else {
+              controller.enqueue('');
+            }
+          }
+          if (done) {
+            logger.info(
+              session,
+              {
+                message: 'Streaming finished',
+                maxChunkSize,
+              },
+              request,
+              '__AGENT_DONE',
+            );
+            controller.close();
+          }
+        }
+      },
+    });
+    return new Response(stream, {
+      headers: {
+        'Content-Type': 'application/x-ndjson',
+      },
+    });
   },
 );

lib/utils/content.ts

@@ -83,6 +83,7 @@ export type ChunkBody =
   | {
       type: 'plans' | 'tools' | 'code' | 'final_code' | 'final_error';
       status: 'started' | 'completed' | 'failed' | 'running';
+      timestamp: string;
       payload:
         | Array<Record<string, string>> // PlansBody | ToolsBody
         | CodeResult; // CodeBody