feat: Update useMessage hook with website embedding options and improvements

src/hooks/useMessage.tsx

@@ -44,8 +44,16 @@ export const useMessage = () => {
   const { t } = useTranslation("option")
   const [selectedModel, setSelectedModel] = useStorage("selectedModel")
   const currentChatModelSettings = useStoreChatModelSettings()
-  const { setIsSearchingInternet, webSearch, setWebSearch, isSearchingInternet } =
-    useStoreMessageOption()
+  const {
+    setIsSearchingInternet,
+    webSearch,
+    setWebSearch,
+    isSearchingInternet
+  } = useStoreMessageOption()
+
+
+  const [chatWithWebsiteEmbedding] = useStorage("chatWithWebsiteEmbedding", true)
+  const [maxWebsiteContext] = useStorage("maxWebsiteContext", 4028)
 
   const {
     history,
@@ -150,35 +158,34 @@ export const useMessage = () => {
     setMessages(newMessage)
     let fullText = ""
     let contentToSave = ""
-    let isAlreadyExistEmbedding: MemoryVectorStore
     let embedURL: string, embedHTML: string, embedType: string
     let embedPDF: { content: string; page: number }[] = []
 
-    if (messages.length === 0) {
-      const { content: html, url, type, pdf } = await getDataFromCurrentTab()
+    let isAlreadyExistEmbedding: MemoryVectorStore
+    const {
+      content: html,
+      url: websiteUrl,
+      type,
+      pdf
+    } = await getDataFromCurrentTab()
+
     embedHTML = html
-      embedURL = url
+    embedURL = websiteUrl
     embedType = type
     embedPDF = pdf
-      setCurrentURL(url)
+
+    console.log(embedHTML)
+    if (messages.length === 0) {
+      setCurrentURL(websiteUrl)
       isAlreadyExistEmbedding = keepTrackOfEmbedding[currentURL]
     } else {
-      const { content: html, url, type, pdf } = await getDataFromCurrentTab()
-      if (currentURL !== url) {
-        embedHTML = html
-        embedURL = url
-        embedType = type
-        embedPDF = pdf
-        setCurrentURL(url)
+      if (currentURL !== websiteUrl) {
+        setCurrentURL(websiteUrl)
       } else {
-        embedHTML = html
         embedURL = currentURL
-        embedType = type
-        embedPDF = pdf
       }
-      isAlreadyExistEmbedding = keepTrackOfEmbedding[url]
+      isAlreadyExistEmbedding = keepTrackOfEmbedding[websiteUrl]
     }
-
     setMessages(newMessage)
     const ollamaUrl = await getOllamaURL()
     const embeddingModle = await defaultEmbeddingModelForRag()
@@ -198,6 +205,7 @@ export const useMessage = () => {
         vectorstore = isAlreadyExistEmbedding
         console.log("Embedding already exist")
       } else {
+        if (chatWithWebsiteEmbedding) {
           vectorstore = await memoryEmbedding({
             html: embedHTML,
             keepTrackOfEmbedding: keepTrackOfEmbedding,
@@ -208,7 +216,7 @@ export const useMessage = () => {
             type: embedType,
             url: embedURL
           })
-
+        }
         console.log("Embedding created")
       }
       let query = message
@@ -247,9 +255,20 @@ export const useMessage = () => {
         query = response.content.toString()
       }
 
+      let context: string = ""
+      let source: {
+        name: any
+        type: any
+        mode: string
+        url: string
+        pageContent: string
+        metadata: Record<string, any>
+      }[] = []
+
+      if (chatWithWebsiteEmbedding) {
         const docs = await vectorstore.similaritySearch(query, 4)
-      const context = formatDocs(docs)
-      const source = docs.map((doc) => {
+        context = formatDocs(docs)
+        source = docs.map((doc) => {
           return {
             ...doc,
             name: doc?.metadata?.source || "untitled",
@@ -258,14 +277,37 @@ export const useMessage = () => {
             url: ""
           }
         })
-      // message = message.trim().replaceAll("\n", " ")
+      } else {
+        if (type === "html") {
+          context = embedHTML.slice(0, maxWebsiteContext)
+        } else {
+          context = embedPDF
+            .map((pdf) => pdf.content)
+            .join(" ")
+            .slice(0, maxWebsiteContext)
+        }
+
+        source = [
+          {
+            name: embedURL,
+            type: type,
+            mode: "chat",
+            url: embedURL,
+            pageContent: context,
+            metadata: {
+              source: embedURL,
+              url: embedURL
+            }
+          }
+        ]
+      }
 
       let humanMessage = new HumanMessage({
         content: [
           {
             text: systemPrompt
               .replace("{context}", context)
-              .replace("{question}", message),
+              .replace("{question}", query),
             type: "text"
           }
         ]
@@ -299,7 +341,7 @@ export const useMessage = () => {
         })
         count++
       }
-      // update the message with the full text
+
       setMessages((prev) => {
         return prev.map((message) => {
           if (message.id === generateMessageId) {
@@ -976,6 +1018,6 @@ export const useMessage = () => {
     regenerateLastMessage,
     webSearch,
     setWebSearch,
-    isSearchingInternet,
+    isSearchingInternet
   }
 }