Fix rendering online with note references post streaming chat response

Previously only the notes references would get rendered post response streaming when when both online and notes references were used to respond to the user's message
khoj-ai · Mar 13, 2024 · a1ce122 · a1ce122
1 parent 1aeea3d
commit a1ce122
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 19 deletions.
diff --git a/src/interface/desktop/chat.html b/src/interface/desktop/chat.html
@@ -357,15 +357,16 @@
 
             let numReferences = 0;
 
-            if (Array.isArray(references)) {
-                numReferences = references.length;
+            if (references.hasOwnProperty("notes")) {
+                numReferences += references["notes"].length;
 
-                references.forEach((reference, index) => {
+                references["notes"].forEach((reference, index) => {
                     let polishedReference = generateReference(reference, index);
                     referenceSection.appendChild(polishedReference);
                 });
-            } else {
-                numReferences += processOnlineReferences(referenceSection, references);
+            }
+            if (references.hasOwnProperty("online")){
+                numReferences += processOnlineReferences(referenceSection, references["online"]);
             }
 
             let referenceExpandButton = document.createElement('button');
@@ -511,16 +512,16 @@
                 // Handle streamed response of type text/event-stream or text/plain
                 const reader = response.body.getReader();
                 const decoder = new TextDecoder();
-                let references = null;
+                let references = {};
 
                 readStream();
 
                 function readStream() {
                     reader.read().then(({ done, value }) => {
                         if (done) {
                             // Append any references after all the data has been streamed
-                            if (references != null) {
-                                newResponseText.appendChild(references);
+                            if (references != {}) {
+                                newResponseText.appendChild(createReferenceSection(references));
                             }
                             document.getElementById("chat-body").scrollTop = document.getElementById("chat-body").scrollHeight;
                             document.getElementById("chat-input").removeAttribute("disabled");
@@ -538,7 +539,11 @@
 
                             const rawReference = chunk.split("### compiled references:")[1];
                             const rawReferenceAsJson = JSON.parse(rawReference);
-                            references = createReferenceSection(rawReferenceAsJson);
+                            if (rawReferenceAsJson instanceof Array) {
+                                references["notes"] = rawReferenceAsJson;
+                            } else if (typeof rawReferenceAsJson === "object" && rawReferenceAsJson !== null) {
+                                references["online"] = rawReferenceAsJson;
+                            }
                             readStream();
                         } else {
                             // Display response from Khoj

diff --git a/src/khoj/interface/web/chat.html b/src/khoj/interface/web/chat.html
@@ -368,15 +368,16 @@
 
             let numReferences = 0;
 
-            if (Array.isArray(references)) {
-                numReferences = references.length;
+            if (references.hasOwnProperty("notes")) {
+                numReferences += references["notes"].length;
 
-                references.forEach((reference, index) => {
+                references["notes"].forEach((reference, index) => {
                     let polishedReference = generateReference(reference, index);
                     referenceSection.appendChild(polishedReference);
                 });
-            } else {
-                numReferences += processOnlineReferences(referenceSection, references);
+            }
+            if (references.hasOwnProperty("online")) {
+                numReferences += processOnlineReferences(referenceSection, references["online"]);
             }
 
             let referenceExpandButton = document.createElement('button');
@@ -518,16 +519,16 @@
                 // Handle streamed response of type text/event-stream or text/plain
                 const reader = response.body.getReader();
                 const decoder = new TextDecoder();
-                let references = null;
+                let references = {};
 
                 readStream();
 
                 function readStream() {
                     reader.read().then(({ done, value }) => {
                         if (done) {
                             // Append any references after all the data has been streamed
-                            if (references != null) {
-                                newResponseText.appendChild(references);
+                            if (references != {}) {
+                                newResponseText.appendChild(createReferenceSection(references));
                             }
                             document.getElementById("chat-body").scrollTop = document.getElementById("chat-body").scrollHeight;
                             document.getElementById("chat-input").removeAttribute("disabled");
@@ -545,7 +546,11 @@
 
                             const rawReference = chunk.split("### compiled references:")[1];
                             const rawReferenceAsJson = JSON.parse(rawReference);
-                            references = createReferenceSection(rawReferenceAsJson);
+                            if (rawReferenceAsJson instanceof Array) {
+                                references["notes"] = rawReferenceAsJson;
+                            } else if (typeof rawReferenceAsJson === "object" && rawReferenceAsJson !== null) {
+                                references["online"] = rawReferenceAsJson;
+                            }
                             readStream();
                         } else {
                             // Display response from Khoj

diff --git a/src/khoj/processor/conversation/utils.py b/src/khoj/processor/conversation/utils.py
@@ -64,7 +64,7 @@ def send(self, data):
     def close(self):
         if self.compiled_references and len(self.compiled_references) > 0:
             self.queue.put(f"### compiled references:{json.dumps(self.compiled_references)}")
-        elif self.online_results and len(self.online_results) > 0:
+        if self.online_results and len(self.online_results) > 0:
             self.queue.put(f"### compiled references:{json.dumps(self.online_results)}")
         self.queue.put(StopIteration)