From 707699bc231d7c9abdba8c6659ba1ee9d3f9d7c1 Mon Sep 17 00:00:00 2001
From: STetsing <41009393+STetsing@users.noreply.github.com>
Date: Fri, 18 Oct 2024 12:45:14 +0200
Subject: [PATCH] fixed streaming error leaving out some text in chat UI

---
 .../src/helpers/streamHandler.ts              |  2 +-
 .../src/types/remix-project.code-workspace    |  3 ++
 libs/remix-ai-core/src/types/types.ts         | 45 ++++++++++---------
 .../remix-ai/src/lib/components/Default.tsx   |  1 +
 4 files changed, 30 insertions(+), 21 deletions(-)

diff --git a/libs/remix-ai-core/src/helpers/streamHandler.ts b/libs/remix-ai-core/src/helpers/streamHandler.ts
index af094fe564..a42db7b645 100644
--- a/libs/remix-ai-core/src/helpers/streamHandler.ts
+++ b/libs/remix-ai-core/src/helpers/streamHandler.ts
@@ -19,7 +19,7 @@ export const HandleSimpleResponse = async (response,
 }
 
 export const HandleStreamResponse = async (streamResponse,
-  cb?: (streamText: string) => void,
+  cb: (streamText: string) => void,
   done_cb?: (result: string) => void) => {
   try {
     let resultText = ''
diff --git a/libs/remix-ai-core/src/types/remix-project.code-workspace b/libs/remix-ai-core/src/types/remix-project.code-workspace
index 596ef7a482..01fe49386a 100644
--- a/libs/remix-ai-core/src/types/remix-project.code-workspace
+++ b/libs/remix-ai-core/src/types/remix-project.code-workspace
@@ -2,6 +2,9 @@
     "folders": [
         {
             "path": "../../../.."
+        },
+        {
+            "path": "../../../../../remix-wildcard"
         }
     ]
 }
\ No newline at end of file
diff --git a/libs/remix-ai-core/src/types/types.ts b/libs/remix-ai-core/src/types/types.ts
index 0454798ce5..b170fc6259 100644
--- a/libs/remix-ai-core/src/types/types.ts
+++ b/libs/remix-ai-core/src/types/types.ts
@@ -88,6 +88,10 @@ export enum RemoteBackendOPModel{
   MISTRAL
 }
 
+interface GeneratedTextObject {
+  generatedText: string;
+  isGenerating: boolean;
+}
 export class JsonStreamParser {
   buffer: string
   constructor() {
@@ -97,33 +101,34 @@ export class JsonStreamParser {
   safeJsonParse<T>(chunk: string): T[] | null {
     this.buffer += chunk;
     const results = [];
+    let startIndex = 0;
+    let endIndex: number;
+    while ((endIndex = this.buffer.indexOf('}', startIndex)) !== -1) {
+      // check if next character is a opening curly bracket
+      let modifiedEndIndex = endIndex;
+      if ((modifiedEndIndex = this.buffer.indexOf('{', endIndex)) !== -1 ) {
+        endIndex = modifiedEndIndex - 1;
+      }
 
-    // eslint-disable-next-line no-constant-condition
-    while (true) {
+      if (((modifiedEndIndex = this.buffer.indexOf('{', endIndex)) === -1) &&
+          (this.buffer.indexOf('}', endIndex) < this.buffer.length)) {
+        endIndex = this.buffer.indexOf('}', endIndex+1) <0 ? this.buffer.length - 1 : this.buffer.indexOf('}', endIndex+1);
+      }
+
+      const jsonStr = this.buffer.slice(startIndex, endIndex + 1);
       try {
-        const result = JSON.parse(this.buffer);
-        results.push(result);
-        this.buffer = '';
-        break;
+        const obj: GeneratedTextObject = JSON.parse(jsonStr);
+        console.log('parsed:', obj);
+        results.push(obj);
       } catch (error) {
-        // eslint-disable-next-line no-useless-escape
-        const match = /^([^\{]*\{[^\}]*\})(.*)/.exec(this.buffer);
-        if (match) {
-          try {
-            const result = JSON.parse(match[1]);
-            results.push(result);
-            this.buffer = match[2];
-          } catch (e) {
-            break;
-          }
-        } else {
-          break;
-        }
+        console.error('Error parsing JSON:', error);
       }
+      startIndex = endIndex + 1;
     }
-
+    this.buffer = this.buffer.slice(startIndex);
     return results;
   }
+
   safeJsonParseSingle<T>(chunk: string): T[] | null {
     return JSON.parse(this.buffer);
   }
diff --git a/libs/remix-ui/remix-ai/src/lib/components/Default.tsx b/libs/remix-ui/remix-ai/src/lib/components/Default.tsx
index 784af476ab..b2555f1802 100644
--- a/libs/remix-ui/remix-ai/src/lib/components/Default.tsx
+++ b/libs/remix-ui/remix-ai/src/lib/components/Default.tsx
@@ -29,6 +29,7 @@ export const Default = (props) => {
     if (GenerationParams.return_stream_response) HandleStreamResponse(response,
       (text) => {observer.next(text)},
       (result) => {
+        observer.next(' ') // Add a space to flush
         ChatHistory.pushHistory(prompt, result)
         observer.complete()
       }