feat: vision agent streaming (#32)

<img width="1493" alt="image"
src="https://github.com/landing-ai/vision-agent-ui/assets/132986242/bcaf6c02-6076-41a0-a1b0-83197bc90af7">

Not able to test error logging at the moment, will add later

app/api/vision-agent/route.ts

@@ -1,8 +1,12 @@
-import { StreamingTextResponse } from 'ai';
+import {
+  StreamingTextResponse,
+  experimental_StreamData,
+  createStreamDataTransformer,
+} from 'ai';
 
 // import { auth } from '@/auth';
 import { MessageBase } from '../../../lib/types';
-import { withLogging } from '@/lib/logger';
+import { logger, withLogging } from '@/lib/logger';
 import { CLEANED_SEPARATOR } from '@/lib/constants';
 import { cleanAnswerMessage, cleanInputMessage } from '@/lib/messageUtils';
 
@@ -12,12 +16,13 @@ export const maxDuration = 300; // This function can run for a maximum of 5 minu
 
 export const POST = withLogging(
   async (
-    _session,
+    session,
     json: {
       messages: MessageBase[];
       id: string;
       url: string;
     },
+    request,
   ) => {
     const { messages, url } = json;
 
@@ -65,8 +70,57 @@ export const POST = withLogging(
       },
     );
 
+    if (!fetchResponse.ok) {
+      if (fetchResponse.body) {
+        const reader = fetchResponse.body.getReader();
+        return new StreamingTextResponse(
+          new ReadableStream({
+            async start(controller) {
+              const { done, value } = await reader.read();
+              if (!done) {
+                const errorText = new TextDecoder().decode(value);
+                logger.error(session, errorText, request);
+                controller.error(new Error(`Response error: ${errorText}`));
+              }
+            },
+          }),
+          {
+            status: 400,
+          },
+        );
+      } else {
+        return new StreamingTextResponse(
+          new ReadableStream({
+            start(controller) {
+              logger.error(
+                session,
+                'Response error: No response body',
+                request,
+              );
+              controller.error(new Error('Response error: No response body'));
+            },
+          }),
+          {
+            status: 400,
+          },
+        );
+      }
+    }
+
     if (fetchResponse.body) {
-      return new StreamingTextResponse(fetchResponse.body);
+      const encoder = new TextEncoder();
+      const decoder = new TextDecoder();
+
+      const stream = fetchResponse.body.pipeThrough(
+        new TransformStream({
+          transform: async (chunk, controller) => {
+            const message = decoder.decode(chunk);
+            logger.info(session, { message }, request, '__AGENT_RESPONSE');
+            controller.enqueue(encoder.encode(message));
+          },
+        }),
+      );
+      return new StreamingTextResponse(stream);
     } else {
       return fetchResponse;
     }

lib/logger.ts

@@ -38,16 +38,20 @@ class Logger {
   }
 
   private createLogMessage(
-    request: Request,
     message: Object | string,
     session: Session | null,
+    request?: Request,
     logDescription?: string,
   ) {
     const body = {
       type: logDescription || '__DEFAULT',
       context: {
-        method: request.method,
-        url: request.url,
+        ...(request
+          ? {
+              method: request?.method,
+              url: request.url,
+            }
+          : {}),
         sessionUserId: session?.user.id,
         sessionUserName: session?.user.name,
         sessionUserEmail: session?.user.email,
@@ -58,62 +62,71 @@ class Logger {
   }
 
   info(
-    request: Request,
     session: Session | null,
     message: string | Object,
+    request?: Request,
     logDescription?: string,
   ) {
     this.lokiLogger.info(
-      this.createLogMessage(request, message, session, logDescription),
+      this.createLogMessage(message, session, request, logDescription),
     );
   }
   warn(
-    request: Request,
     session: Session | null,
     message: string | Object,
+    request?: Request,
     logDescription?: string,
   ) {
     this.lokiLogger.warn(
-      this.createLogMessage(request, message, session, logDescription),
+      this.createLogMessage(message, session, request, logDescription),
     );
   }
   debug(
-    request: Request,
     session: Session | null,
     message: string | Object,
+    request?: Request,
     logDescription?: string,
   ) {
     this.lokiLogger.debug(
-      this.createLogMessage(request, message, session, logDescription),
+      this.createLogMessage(message, session, request, logDescription),
     );
   }
   error(
-    request: Request,
     session: Session | null,
     message: string | Object,
+    request?: Request,
     logDescription?: string,
   ) {
     this.lokiLogger.error(
-      this.createLogMessage(request, message, session, logDescription),
+      this.createLogMessage(
+        message,
+        session,
+        request,
+        logDescription ?? '__ERROR',
+      ),
     );
   }
 }
 
 export const withLogging = (
-  handler: (session: Session | null, json: any) => Promise<Response>,
+  handler: (
+    session: Session | null,
+    json: any,
+    request: Request,
+  ) => Promise<Response>,
 ) => {
   return async (req: Request) => {
     const session = await auth();
     const json = await req.json();
     logger.info(
-      req,
       session,
       {
         params: json,
       },
+      req,
       '_API_REQUEST',
     );
-    return handler(session, json);
+    return handler(session, json, req);
   };
 };