Update chatStream.js

server/chatStream.js

@@ -3,6 +3,7 @@ import { Worker } from "worker_threads";
 import { fileURLToPath } from "url";
 import path from "path";
 import { LIGHTNING_BASE } from "./config.js";
+import WebSocket from "ws";
 
 const __dirname = path.dirname(fileURLToPath(import.meta.url));
 const WORKER_PATH = path.join(__dirname, "searchWorker.js");
@@ -54,8 +55,11 @@ function makeClient(accessToken, clientId) {
 
 async function websocketChatStream(body, headers, onToken) {
 
-  const ws = new WebSocket(`${LIGHTNING_BASE.replace("http","ws")}/ws/chat`);
-
+  const wsURL =
+    (LIGHTNING_BASE.startsWith("https")
+      ? LIGHTNING_BASE.replace("https", "wss")
+      : LIGHTNING_BASE.replace("http", "ws")
+    ) + "/ws/chat";
   await new Promise(resolve => {
     ws.onopen = resolve;
   });
@@ -80,9 +84,9 @@ async function websocketChatStream(body, headers, onToken) {
 
   return new Promise((resolve, reject) => {
 
-    ws.onmessage = (event) => {
+    ws.on("message", (data) => {
 
-      const line = event.data;
+      const line = data.toString();
 
       if (!line.startsWith("data:")) return;
 
@@ -99,7 +103,7 @@ async function websocketChatStream(body, headers, onToken) {
         ws.close();
         resolve({ assistantText });
       }
-    };
+    });
 
     ws.onerror = reject;
   });