fixes #1

NilFoundation · Feb 4, 2025 · dcf8363 · dcf8363
1 parent 5f15d3a
commit dcf8363
Show file tree

Hide file tree

Showing 2 changed files with 77 additions and 55 deletions.
diff --git a/docs/src/components/CustomChatRuntimeProvider/index.tsx b/docs/src/components/CustomChatRuntimeProvider/index.tsx
@@ -20,20 +20,6 @@ function processChunk(chunk: string): string {
   return cleanedText;
 }
 
-function asAsyncIterable<T>(source: ReadableStream<T>): AsyncIterable<T> {
-  return {
-    [Symbol.asyncIterator]: () => {
-      const reader = source.getReader();
-      return {
-        async next(): Promise<IteratorResult<T, undefined>> {
-          const { done, value } = await reader.read();
-          return done ? { done: true, value: undefined } : { done: false, value };
-        },
-      };
-    },
-  };
-}
-
 const CustomModelAdapter: (string, Function) => ChatModelAdapter = (
   token,
   handleReCaptchaVerify,
@@ -46,6 +32,7 @@ const CustomModelAdapter: (string, Function) => ChatModelAdapter = (
         .map((c) => c.text)
         .join(" "),
     }));
+
     const response = await fetch("https://docs.nil.foundation/bot/api/chat", {
       method: "POST",
       headers: {
@@ -57,19 +44,51 @@ const CustomModelAdapter: (string, Function) => ChatModelAdapter = (
       }),
       signal: abortSignal,
     });
+
     handleReCaptchaVerify();
 
-    let text = "";
-    for await (const chunk of asAsyncIterable(
-      response.body!.pipeThrough(new TextDecoderStream()),
-    )) {
-      const matches = [...chunk.matchAll(pattern)];
-      const tokens = matches.map((match) => match[1]);
-
-      let cleanedText = tokens.join("");
-      cleanedText = cleanedText.replace(/\\n/g, "\n");
-      text += cleanedText;
-      yield { content: [{ type: "text", text }] };
+    if (!response.ok) {
+      throw new Error(`HTTP error! status: ${response.status}`);
+    }
+
+    if (!response.body) {
+      throw new Error('No response body received');
+    }
+
+    const reader = response.body.getReader();
+    const decoder = new TextDecoder();
+    let buffer = '';
+    let accumulatedText = '';
+
+    try {
+      while (true) {
+        const { value, done } = await reader.read();
+
+        if (done) {
+          if (buffer.trim()) {
+            const cleanedText = processChunk(buffer);
+            if (cleanedText) {
+              accumulatedText += cleanedText;
+              yield { content: [{ type: "text", text: accumulatedText }] };
+            }
+          }
+          break;
+        }
+
+        buffer += decoder.decode(value, { stream: true });
+
+        const cleanedText = processChunk(buffer);
+        if (cleanedText) {
+          accumulatedText += cleanedText;
+          yield { content: [{ type: "text", text: accumulatedText }] };
+          buffer = '';
+        }
+      }
+    } catch (error) {
+      console.error('Stream reading error:', error);
+      throw error;
+    } finally {
+      reader.releaseLock();
     }
   },
 });

diff --git a/docs_ai_backend/app/api/chat/route.ts b/docs_ai_backend/app/api/chat/route.ts
@@ -49,40 +49,43 @@ export async function POST(req: Request) {
 
       const stream = new TransformStream();
       const writer = stream.writable.getWriter();
-      const encoder = new TextEncoder();
 
       (async () => {
-        try {
-          const chain = intent === "Question" 
-            ? handlerToUse.genericLllmChain 
-            : handlerToUse.generatorLllmChain;
-
-          const result = await chain.stream({
-            query: query,
-            context: retrievedDocs,
-            sources: sources
-          });
-
-          for await (const chunk of result) {
-            await writer.write(encoder.encode(JSON.stringify(chunk) + '\n'));
+      try {
+        const chain = intent === "Question" 
+          ? handlerToUse.genericLllmChain 
+          : handlerToUse.generatorLllmChain;
+
+        const result = await chain.stream({
+          query: query,
+          context: retrievedDocs,
+          sources: sources
+        });
+
+        for await (const chunk of result) {
+          await writer.write(chunk);
+
+          if (writer.ready) {
+            await writer.ready;
           }
-        } catch (error) {
-          console.error('Streaming error:', error);
-          const errorChunk = JSON.stringify({ error: 'Streaming error occurred' }) + '\n';
-          await writer.write(encoder.encode(errorChunk));
-        } finally {
-          await writer.close();
         }
-      })();
-      return new Response(stream.readable, {
-        headers: {
-          'Content-Type': 'text/event-stream',
-          'Cache-Control': 'no-cache',
-          'Connection': 'keep-alive',
-          'Access-Control-Allow-Origin': '*',
-          'Access-Control-Allow-Methods': 'GET,OPTIONS,PATCH,DELETE,POST,PUT'
-        },
-      });
+      } catch (error) {
+        console.error('Streaming error:', error);
+      } finally {
+        await writer.close();
+      }
+    })();
+
+    return new Response(stream.readable, {
+      headers: {
+        'Content-Type': 'text/event-stream',
+        'Cache-Control': 'no-cache, no-transform',
+        'Connection': 'keep-alive',
+        'X-Accel-Buffering': 'no', 
+        'Access-Control-Allow-Origin': '*',
+        'Access-Control-Allow-Methods': 'GET,OPTIONS,PATCH,DELETE,POST,PUT'
+      },
+    });
 
     } else {
       return new Response(JSON.stringify({ message: "Failed to verify" }), {