Refactor useMessage hook to improve message handling and performance

n4ze3m · Jun 17, 2024 · 83e634f · 83e634f
1 parent dcc378f
commit 83e634f
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 24 deletions.
diff --git a/app/widget/src/hooks/useMessage.tsx b/app/widget/src/hooks/useMessage.tsx
@@ -15,8 +15,6 @@ export type BotResponse = {
 };
 
 const parsesStreamingResponse = (text: string) => {
-  // event: chunk or result\ndata:  been or object\n\n
-  //   console.log(`text: ${text}`);
   const REGEX = /event: (.+)\ndata: (.+)/g;
   const matches = text.matchAll(REGEX);
   const result = [];
@@ -149,18 +147,20 @@ export const useMessage = () => {
           if (type === "chunk") {
             const jsonMessage = JSON.parse(message);
             if (count === 0) {
-              newMessage[appendingIndex].message = jsonMessage.message;
+              newMessage[appendingIndex].message = jsonMessage.message + "▋";
               setMessages(newMessage);
               localStorage.setItem("DS_MESSAGE", JSON.stringify(newMessage));
             } else {
-              newMessage[appendingIndex].message += jsonMessage.message;
+              newMessage[appendingIndex].message =
+                newMessage[appendingIndex].message.slice(0, -1) +
+                jsonMessage.message +
+                "▋";
               setMessages(newMessage);
               localStorage.setItem("DS_MESSAGE", JSON.stringify(newMessage));
             }
             count++;
           } else if (type === "result") {
             const responseData = JSON.parse(message) as BotResponse;
-            console.log(responseData);
             newMessage[appendingIndex].message = responseData.bot.text;
             newMessage[appendingIndex].sources =
               responseData.bot.sourceDocuments;

diff --git a/server/src/utils/puppeteer-fetch.ts b/server/src/utils/puppeteer-fetch.ts
@@ -2,7 +2,7 @@ import type { Browser } from 'puppeteer';
 import puppeteer from 'puppeteer-extra';
 import puppeteerBlockResources from 'puppeteer-extra-plugin-block-resources';
 import puppeteerPageProxy from 'puppeteer-extra-plugin-page-proxy';
-import puppeteerStealth from 'puppeteer-extra-plugin-stealth'
+import puppeteerStealth from 'puppeteer-extra-plugin-stealth';
 import * as fs from 'fs';
 
 const readabilityJsStr = fs.readFileSync(
@@ -15,9 +15,8 @@ function executor() {
     return new Readability({}, document).parse();
 }
 
+// Use puppeteer plugins
 puppeteer.use(puppeteerStealth());
-
-
 puppeteer.use(puppeteerBlockResources({
     blockedTypes: new Set(['media']),
     interceptResolutionPriority: 1,
@@ -28,7 +27,6 @@ puppeteer.use(puppeteerPageProxy({
 
 let browser: Browser;
 
-
 const init = async () => {
     if (!browser || !browser.connected) {
         browser = await puppeteer.launch({
@@ -39,45 +37,48 @@ const init = async () => {
     }
 }
 
-
 const puppeteerFetch = async (url: string, useReadability = false) => {
+    await init();
+    let page = null;
     try {
-        await init();
-        const page = await browser.newPage();
+        page = await browser.newPage();
         await page.goto(url, { waitUntil: 'networkidle2' });
+
         if (useReadability) {
             const resultArticle = await page.evaluate(`
                 (function(){
-                  ${readabilityJsStr}
-                  ${executor}
-                  return executor();
+                    ${readabilityJsStr}
+                    ${executor}
+                    return executor();
                 }())
-              `) as { content?: string, title?: string };
+            `) as { content?: string, title?: string };
+
+            await page.close();
+
             if (resultArticle?.content) {
-                await page.close();
-                return `<!DOCTYPE html><html><head><title>${resultArticle.title}</title></head><body>${resultArticle.content}</body></html>`
+                return `<!DOCTYPE html><html><head><title>${resultArticle.title}</title></head><body>${resultArticle.content}</body></html>`;
             }
             console.error(`[puppeteerFetch] Error fetching ${url}: Readability failed`);
         }
+
         const html = await page.content();
         await page.close();
         return html;
     } catch (error) {
+        if (page) await page.close();
         console.error(`[puppeteerFetch] Error fetching ${url}: ${error.message}`);
         return '';
     }
 }
 
 export const closePuppeteer = async () => {
-    try {
-        if (browser.connected) {
+    if (browser && browser.connected) {
+        try {
             await browser.close();
+        } catch (error) {
+            console.error(`[closePuppeteer] Error closing browser: ${error.message}`);
         }
-    } catch (error) {
-        console.error(`[closePuppeteer] Error closing browser: ${error.message}`);
     }
 }
 
-
-
 export default puppeteerFetch;