From fd0a922b4af146d219f087532db97af68dc64f58 Mon Sep 17 00:00:00 2001
From: n4ze3m <n4ze3m@gmail.com>
Date: Mon, 17 Jun 2024 11:50:21 +0530
Subject: [PATCH] feat: Add usePuppeteerFetch option to QSource and
 DialoqbaseSettings

---
 app/ui/src/routes/settings/application.tsx    |   7 ++
 .../20240617050823_q_12_4/migration.sql       |   2 +
 server/prisma/schema.prisma                   |   1 +
 server/prisma/seed.ts                         | 116 +++++++++---------
 server/src/handlers/api/v1/admin/type.ts      |   1 +
 server/src/loader/web.ts                      |  19 ++-
 .../src/queue/controllers/crawl.controller.ts |   7 +-
 .../queue/controllers/website.controller.ts   |   2 +
 server/src/queue/index.ts                     |   3 +-
 server/src/queue/type.ts                      |   2 +
 server/src/schema/api/v1/admin/index.ts       |   2 +
 server/src/utils/crawl.ts                     |  35 ++++--
 server/src/utils/puppeteer-fetch.ts           |  16 ++-
 server/src/utils/rag-settings.ts              |   2 +
 14 files changed, 136 insertions(+), 79 deletions(-)
 create mode 100644 server/prisma/migrations/20240617050823_q_12_4/migration.sql

diff --git a/app/ui/src/routes/settings/application.tsx b/app/ui/src/routes/settings/application.tsx
index ec1579f1..8106e7ad 100644
--- a/app/ui/src/routes/settings/application.tsx
+++ b/app/ui/src/routes/settings/application.tsx
@@ -154,6 +154,13 @@ export default function SettingsApplicationRoot() {
                   >
                     <Switch />
                   </Form.Item>
+                  <Form.Item
+                    label="Enhanced Website loader"
+                    name="usePuppeteerFetch"
+                    valuePropName="checked"
+                  >
+                    <Switch />
+                  </Form.Item>
                 </div>
                 <div className="bg-gray-50 border-x border-b rounded-b-md rounded-x-md px-4 py-3 text-right sm:px-6 dark:bg-[#141414] dark:border-gray-600">
                   <button
diff --git a/server/prisma/migrations/20240617050823_q_12_4/migration.sql b/server/prisma/migrations/20240617050823_q_12_4/migration.sql
new file mode 100644
index 00000000..de1e26fc
--- /dev/null
+++ b/server/prisma/migrations/20240617050823_q_12_4/migration.sql
@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "DialoqbaseSettings" ADD COLUMN     "usePuppeteerFetch" BOOLEAN DEFAULT false;
diff --git a/server/prisma/schema.prisma b/server/prisma/schema.prisma
index 6a99a300..0a71ff23 100644
--- a/server/prisma/schema.prisma
+++ b/server/prisma/schema.prisma
@@ -99,6 +99,7 @@ model DialoqbaseSettings {
   defaultChatModel             String   @default("gpt-3.5-turbo-dbase")
   defaultEmbeddingModel        String   @default("dialoqbase_eb_text-embedding-ada-002")
   ollamaURL                    String?  @default("http://host.docker.internal:11434")
+  usePuppeteerFetch            Boolean? @default(false)
 }
 
 model BotIntegration {
diff --git a/server/prisma/seed.ts b/server/prisma/seed.ts
index 5c2f34b1..d3753a72 100644
--- a/server/prisma/seed.ts
+++ b/server/prisma/seed.ts
@@ -488,70 +488,70 @@ const removeTensorflowSupport = async () => {
   });
 };
 
-const replaceOldEmbeddings = async () => {
-  await prisma.bot.updateMany({
-    where: {
-      embedding: "openai",
-    },
-    data: {
-      embedding: "dialoqbase_eb_text-embedding-ada-002",
-    },
-  });
+// const replaceOldEmbeddings = async () => {
+//   await prisma.bot.updateMany({
+//     where: {
+//       embedding: "openai",
+//     },
+//     data: {
+//       embedding: "dialoqbase_eb_text-embedding-ada-002",
+//     },
+//   });
 
-  await prisma.bot.updateMany({
-    where: {
-      embedding: "cohere",
-    },
-    data: {
-      embedding: "dialoqbase_eb_small",
-    },
-  });
+//   await prisma.bot.updateMany({
+//     where: {
+//       embedding: "cohere",
+//     },
+//     data: {
+//       embedding: "dialoqbase_eb_small",
+//     },
+//   });
 
-  await prisma.bot.updateMany({
-    where: {
-      embedding: "transformer",
-    },
-    data: {
-      embedding: "dialoqbase_eb_Xenova/all-MiniLM-L6-v2",
-    },
-  });
+//   await prisma.bot.updateMany({
+//     where: {
+//       embedding: "transformer",
+//     },
+//     data: {
+//       embedding: "dialoqbase_eb_Xenova/all-MiniLM-L6-v2",
+//     },
+//   });
 
-  await prisma.bot.updateMany({
-    where: {
-      embedding: "google-gecko",
-    },
-    data: {
-      embedding: "dialoqbase_eb_models/embedding-gecko-001",
-    },
-  });
+//   await prisma.bot.updateMany({
+//     where: {
+//       embedding: "google-gecko",
+//     },
+//     data: {
+//       embedding: "dialoqbase_eb_models/embedding-gecko-001",
+//     },
+//   });
 
-  await prisma.bot.updateMany({
-    where: {
-      embedding: "jina-api",
-    },
-    data: {
-      embedding: "dialoqbase_eb_jina-embeddings-v2-base-en",
-    },
-  });
+//   await prisma.bot.updateMany({
+//     where: {
+//       embedding: "jina-api",
+//     },
+//     data: {
+//       embedding: "dialoqbase_eb_jina-embeddings-v2-base-en",
+//     },
+//   });
 
-  await prisma.bot.updateMany({
-    where: {
-      embedding: "jina",
-    },
-    data: {
-      embedding: "dialoqbase_eb_Xenova/jina-embeddings-v2-small-en",
-    },
-  });
+//   await prisma.bot.updateMany({
+//     where: {
+//       embedding: "jina",
+//     },
+//     data: {
+//       embedding: "dialoqbase_eb_Xenova/jina-embeddings-v2-small-en",
+//     },
+//   });
 
-  await prisma.bot.updateMany({
-    where: {
-      embedding: "google",
-    },
-    data: {
-      embedding: "dialoqbase_eb_embedding-001",
-    },
-  });
-};
+//   await prisma.bot.updateMany({
+//     where: {
+//       embedding: "google",
+//     },
+//     data: {
+//       embedding: "dialoqbase_eb_embedding-001",
+//     },
+//   });
+// };
 
 const updateGeminiStreamingToTrue = async () => {
   await prisma.dialoqbaseModels.update({
diff --git a/server/src/handlers/api/v1/admin/type.ts b/server/src/handlers/api/v1/admin/type.ts
index 904b50b3..faf10b26 100644
--- a/server/src/handlers/api/v1/admin/type.ts
+++ b/server/src/handlers/api/v1/admin/type.ts
@@ -3,6 +3,7 @@ export type UpdateDialoqbaseSettingsRequest = {
     noOfBotsPerUser: number;
     allowUserToCreateBots: boolean;
     allowUserToRegister: boolean;
+    usePuppeteerFetch: boolean;
   };
 };
 
diff --git a/server/src/loader/web.ts b/server/src/loader/web.ts
index 616a9687..4d10850e 100644
--- a/server/src/loader/web.ts
+++ b/server/src/loader/web.ts
@@ -1,23 +1,37 @@
 import { BaseDocumentLoader } from "langchain/document_loaders/base";
 import { Document } from "langchain/document";
 import { websiteParser } from "../utils/website-parser";
-// import puppeteerFetch from "../utils/puppeteer-fetch";
+import puppeteerFetch, { closePuppeteer } from "../utils/puppeteer-fetch";
 
 export interface WebLoaderParams {
   url: string;
+  usePuppeteerFetch?: boolean;
+  doNotClosePuppeteer?: boolean;
 }
 
 export class DialoqbaseWebLoader
   extends BaseDocumentLoader
   implements WebLoaderParams {
   url: string;
+  usePuppeteerFetch?: boolean;
+  doNotClosePuppeteer?: boolean;
 
-  constructor({ url }: WebLoaderParams) {
+  constructor({ url, usePuppeteerFetch, doNotClosePuppeteer }: WebLoaderParams) {
     super();
     this.url = url;
+    this.usePuppeteerFetch = usePuppeteerFetch;
+    this.doNotClosePuppeteer = doNotClosePuppeteer;
   }
 
   async _fetchHTML(): Promise<string> {
+    if (this.usePuppeteerFetch) {
+      console.log(`[DialoqbaseWebLoader] Using puppeteer to fetch ${this.url}`)
+      const response = await puppeteerFetch(this.url, true);
+      if (!this.doNotClosePuppeteer) {
+        await closePuppeteer();
+      }
+      return response;
+    }
     const response = await fetch(this.url);
     return await response.text();
   }
@@ -25,7 +39,6 @@ export class DialoqbaseWebLoader
   async load(): Promise<Document<Record<string, any>>[]> {
     const html = await this._fetchHTML();
     const text = websiteParser(html);
-    console.log(text)
     const metadata = { source: this.url };
     return [new Document({ pageContent: text, metadata })];
   }
diff --git a/server/src/queue/controllers/crawl.controller.ts b/server/src/queue/controllers/crawl.controller.ts
index 494d0058..18bff46c 100644
--- a/server/src/queue/controllers/crawl.controller.ts
+++ b/server/src/queue/controllers/crawl.controller.ts
@@ -2,12 +2,13 @@ import { PrismaClient } from "@prisma/client";
 import { QSource } from "../type";
 import { crawl } from "../../utils/crawl";
 import { websiteQueueController } from "./website.controller";
+import { closePuppeteer } from "../../utils/puppeteer-fetch";
 const prisma = new PrismaClient();
 
 export const crawlQueueController = async (source: QSource) => {
   let maxDepth = source.maxDepth || 1;
   let maxLinks = source.maxLinks || 1;
-  const data = await crawl(source.content!, maxDepth, maxLinks);
+  const data = await crawl(source.content!, maxDepth, maxLinks, source.usePuppeteerFetch);
   const links = Array.from(data?.links || []);
 
   for (const link of links) {
@@ -27,6 +28,8 @@ export const crawlQueueController = async (source: QSource) => {
         embedding: source.embedding,
         chunkOverlap: source.chunkOverlap,
         chunkSize: source.chunkSize,
+        usePuppeteerFetch: source.usePuppeteerFetch,
+        doNotClosePuppeteer: true,
       },
       prisma
     );
@@ -41,4 +44,6 @@ export const crawlQueueController = async (source: QSource) => {
       },
     });
   }
+
+  await closePuppeteer()
 };
diff --git a/server/src/queue/controllers/website.controller.ts b/server/src/queue/controllers/website.controller.ts
index 23d0dc22..38997656 100644
--- a/server/src/queue/controllers/website.controller.ts
+++ b/server/src/queue/controllers/website.controller.ts
@@ -73,6 +73,8 @@ export const websiteQueueController = async (
     } else {
       const loader = new DialoqbaseWebLoader({
         url: source.content!,
+        usePuppeteerFetch: source.usePuppeteerFetch,
+        doNotClosePuppeteer: source.doNotClosePuppeteer,
       });
       docs = await loader.load();
     }
diff --git a/server/src/queue/index.ts b/server/src/queue/index.ts
index a0d12338..6d2c63df 100644
--- a/server/src/queue/index.ts
+++ b/server/src/queue/index.ts
@@ -34,9 +34,10 @@ export default async function queueHandler(job: SandboxedJob) {
             status: "PROCESSING",
           },
         });
-        const { chunkOverlap, chunkSize } = await getRagSettings(prisma);
+        const { chunkOverlap, chunkSize , usePuppeteerFetch} = await getRagSettings(prisma);
         source.chunkOverlap = chunkOverlap;
         source.chunkSize = chunkSize;
+        source.usePuppeteerFetch = usePuppeteerFetch;
         switch (source.type.toLowerCase()) {
           case "website":
             await websiteQueueController(source, prisma);
diff --git a/server/src/queue/type.ts b/server/src/queue/type.ts
index f74f4286..72eb0a93 100644
--- a/server/src/queue/type.ts
+++ b/server/src/queue/type.ts
@@ -6,4 +6,6 @@ export interface QSource extends BotSource {
   maxLinks?: number;
   chunkSize: number;
   chunkOverlap: number;
+  usePuppeteerFetch?: boolean;
+  doNotClosePuppeteer?: boolean;
 }
diff --git a/server/src/schema/api/v1/admin/index.ts b/server/src/schema/api/v1/admin/index.ts
index 424b6e6d..992be4d9 100644
--- a/server/src/schema/api/v1/admin/index.ts
+++ b/server/src/schema/api/v1/admin/index.ts
@@ -22,6 +22,7 @@ export const dialoqbaseSettingsSchema: FastifySchema = {
       dynamicallyFetchOllamaModels: { type: "boolean" },
       hideDefaultModels: { type: "boolean" },
       ollamaURL: { type: "string" },
+      usePuppeteerFetch: { type: "boolean" },
     },
   },
 };
@@ -47,6 +48,7 @@ export const updateDialoqbaseSettingsSchema: FastifySchema = {
       defaultEmbeddingModel: { type: "string" },
       hideDefaultModels: { type: "boolean" },
       ollamaURL: { type: "string" },
+      usePuppeteerFetch: { type: "boolean" },
     },
   },
   response: {
diff --git a/server/src/utils/crawl.ts b/server/src/utils/crawl.ts
index 4f2508b5..e1764b89 100644
--- a/server/src/utils/crawl.ts
+++ b/server/src/utils/crawl.ts
@@ -1,5 +1,6 @@
 import axios from "axios";
-import { load } from "cheerio";
+import { CheerioAPI, load } from "cheerio";
+import puppeteerFetch from "./puppeteer-fetch";
 
 type CrawlResult = {
   links: Set<string>;
@@ -13,7 +14,8 @@ const queuedLinks: Set<string> = new Set();
 export const crawl = async (
   startUrl: string,
   maxDepth = 2,
-  maxLinks = 20
+  maxLinks = 20,
+  usePuppeteerFetch = false
 ): Promise<CrawlResult> => {
   const queue: { url: string; depth: number }[] = [{ url: startUrl, depth: 0 }];
   const fetchedLinks: Set<string> = new Set();
@@ -28,20 +30,27 @@ export const crawl = async (
         }
 
         try {
-          const response = await axios.get(url, {
-            headers: {
-              Accept: "text/html",
-              "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3",
-            },
-          });
 
-          const contentType = response.headers['content-type'];
-          if (!contentType || !contentType.includes("text/html")) {
-            return;
-          }
+          let $: CheerioAPI;
+
+          if (usePuppeteerFetch) {
+            const response = await puppeteerFetch(url);
+            $ = load(response);
+          } else {
+            const response = await axios.get(url, {
+              headers: {
+                Accept: "text/html",
+                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3",
+              },
+            });
 
-          const $ = load(response.data);
+            const contentType = response.headers['content-type'];
+            if (!contentType || !contentType.includes("text/html")) {
+              return;
+            }
 
+            $ = load(response.data);
+          }
           visitedLinks.add(url);
           fetchedLinks.add(url);
 
diff --git a/server/src/utils/puppeteer-fetch.ts b/server/src/utils/puppeteer-fetch.ts
index 4fcba71f..00f0a1d9 100644
--- a/server/src/utils/puppeteer-fetch.ts
+++ b/server/src/utils/puppeteer-fetch.ts
@@ -30,7 +30,7 @@ let browser: Browser;
 
 
 const init = async () => {
-    if (!browser) {
+    if (!browser || !browser.connected) {
         browser = await puppeteer.launch({
             headless: true,
             args: ['--no-sandbox', '--disable-setuid-sandbox'],
@@ -52,10 +52,10 @@ const puppeteerFetch = async (url: string, useReadability = false) => {
                   ${executor}
                   return executor();
                 }())
-              `) as { content?: string }
+              `) as { content?: string, title?: string };
             if (resultArticle?.content) {
                 await page.close();
-                return resultArticle.content;
+                return `<!DOCTYPE html><html><head><title>${resultArticle.title}</title></head><body>${resultArticle.content}</body></html>`
             }
             console.error(`[puppeteerFetch] Error fetching ${url}: Readability failed`);
         }
@@ -68,6 +68,16 @@ const puppeteerFetch = async (url: string, useReadability = false) => {
     }
 }
 
+export const closePuppeteer = async () => {
+    try {
+        if (browser.connected) {
+            await browser.close();
+        }
+    } catch (error) {
+        console.error(`[closePuppeteer] Error closing browser: ${error.message}`);
+    }
+}
+
 
 
 export default puppeteerFetch;
\ No newline at end of file
diff --git a/server/src/utils/rag-settings.ts b/server/src/utils/rag-settings.ts
index d27c31e9..d8fe09da 100644
--- a/server/src/utils/rag-settings.ts
+++ b/server/src/utils/rag-settings.ts
@@ -5,11 +5,13 @@ export const getRagSettings = async (prisma: PrismaClient) => {
     select: {
       defaultChunkSize: true,
       defaultChunkOverlap: true,
+      usePuppeteerFetch: true,
     },
   });
 
   return {
     chunkSize: data?.defaultChunkSize || 1000,
     chunkOverlap: data?.defaultChunkOverlap || 200,
+    usePuppeteerFetch: data?.usePuppeteerFetch
   };
 };