feat: Add Number of Retrieved Documents on RAG

src/assets/locale/en/settings.json

@@ -311,6 +311,11 @@
         "label": "Knowledge Base Default File Upload Limit",
         "placeholder": "Enter default file upload limit (e.g., 10)",
         "required": "Please enter the default file upload limit"
+      },
+      "noOfRetrievedDocs": {
+        "label": "Number of Retrieved Documents",
+        "placeholder": "Enter Number of Retrieved Documents",
+        "required": "Please enter the number of retrieved documents"
       }
     },
     "prompt": {

src/assets/locale/es/settings.json

@@ -30,7 +30,7 @@
       "sendNotificationAfterIndexing": {
         "label": "Enviar notificación después de terminar el procesamiento de la base de conocimientos"
       },
-      "generateTitle" :{
+      "generateTitle": {
         "label": "Generar título usando IA"
       }
     },
@@ -310,7 +310,13 @@
       "totalFilePerKB": {
         "label": "Límite predeterminado de carga de archivos para la Base de Conocimientos",
         "placeholder": "Ingrese el límite predeterminado de carga de archivos (ej., 10)",
-        "required": "Por favor, ingrese el límite predeterminado de carga de archivos"      }
+        "required": "Por favor, ingrese el límite predeterminado de carga de archivos"
+      },
+      "noOfRetrievedDocs": {
+        "label": "Número de Documentos Recuperados",
+        "placeholder": "Ingrese el Número de Documentos Recuperados",
+        "required": "Por favor, ingrese el número de documentos recuperados"
+      }
     },
     "prompt": {
       "label": "Configurar el Prompt del RAG",

src/assets/locale/fa/settings.json

@@ -30,7 +30,7 @@
       "sendNotificationAfterIndexing": {
         "label": "ارسال نوتیفیکیشن پس از اتمام پردازش پایگاه دانش"
       },
-      "generateTitle" :{
+      "generateTitle": {
         "label": "تولید عنوان با استفاده از هوش مصنوعی"
       }
     },
@@ -309,7 +309,13 @@
       "totalFilePerKB": {
         "label": "محدودیت پیش‌فرض آپلود فایل پایگاه دانش",
         "placeholder": "محدودیت پیش‌فرض آپلود فایل را وارد کنید (مثلاً 10)",
-        "required": "لطفاً محدودیت پیش‌فرض آپلود فایل را وارد کنید"      }
+        "required": "لطفاً محدودیت پیش‌فرض آپلود فایل را وارد کنید"
+      },
+      "noOfRetrievedDocs": {
+        "label": "تعداد اسناد بازیابی شده",
+        "placeholder": "تعداد اسناد بازیابی شده را وارد کنید",
+        "required": "لطفاً تعداد اسناد بازیابی شده را وارد کنید"
+      }
     },
     "prompt": {
       "label": "پیکربندی پرامپت RAG",

src/assets/locale/fr/settings.json

@@ -30,7 +30,7 @@
       "sendNotificationAfterIndexing": {
         "label": "Envoyer une notification après avoir terminé le traitement de la base de connaissances"
       },
-      "generateTitle" :{
+      "generateTitle": {
         "label": "Générer le titre en utilisant l'IA"
       }
     },
@@ -310,7 +310,13 @@
       "totalFilePerKB": {
         "label": "Limite par défaut de téléchargement de fichiers pour la base de connaissances",
         "placeholder": "Entrez la limite par défaut de téléchargement de fichiers (par exemple, 10)",
-        "required": "Veuillez saisir la limite par défaut de téléchargement de fichiers"      }
+        "required": "Veuillez saisir la limite par défaut de téléchargement de fichiers"
+      },
+      "noOfRetrievedDocs": {
+        "label": "Nombre de documents récupérés",
+        "placeholder": "Entrez le nombre de documents récupérés",
+        "required": "Veuillez saisir le nombre de documents récupérés"
+      }
     },
     "prompt": {
       "label": "Configure RAG Prompt",

src/assets/locale/it/settings.json

@@ -30,7 +30,7 @@
       "sendNotificationAfterIndexing": {
         "label": "Inviare notifica dopo aver terminato l'elaborazione della base di conoscenza"
       },
-      "generateTitle" :{
+      "generateTitle": {
         "label": "Genera titolo utilizzando l'IA"
       }
     },
@@ -310,7 +310,13 @@
       "totalFilePerKB": {
         "label": "Limite Predefinito di Caricamento File per la Base di Conoscenza",
         "placeholder": "Inserisci il limite predefinito di caricamento file (es. 10)",
-        "required": "Inserisci il limite predefinito di caricamento file"      }
+        "required": "Inserisci il limite predefinito di caricamento file"
+      },
+      "noOfRetrievedDocs": {
+        "label": "Numero di Documenti Recuperati",
+        "placeholder": "Inserisci il Numero di Documenti Recuperati",
+        "required": "Inserisci il numero di documenti recuperati"
+      }
     },
     "prompt": {
       "label": "Configura il Prompt RAG",

src/assets/locale/ja-JP/settings.json

@@ -314,6 +314,11 @@
         "label": "知識ベースのデフォルトファイルアップロード制限",
         "placeholder": "デフォルトのファイルアップロード制限を入力してください（例：10）",
         "required": "デフォルトのファイルアップロード制限を入力してください"
+      },
+      "noOfRetrievedDocs": {
+        "label": "取得ドキュメント数",
+        "placeholder": "取得ドキュメント数を入力",
+        "required": "取得ドキュメント数を入力してください"
       }
     },
     "prompt": {

src/assets/locale/ml/settings.json

@@ -314,6 +314,11 @@
         "label": "വിജ്ഞാനാധാരത്തിന്റെ സ്ഥിര ഫയൽ അപ്‌ലോഡ് പരിധി",
         "placeholder": "സ്ഥിര ഫയൽ അപ്‌ലോഡ് പരിധി നൽകുക (ഉദാ: 10)",
         "required": "ദയവായി സ്ഥിര ഫയൽ അപ്‌ലോഡ് പരിധി നൽകുക"
+      },
+      "noOfRetrievedDocs": {
+        "label": "വീണ്ടെടുത്ത രേഖകളുടെ എണ്ണം",
+        "placeholder": "വീണ്ടെടുത്ത രേഖകളുടെ എണ്ണം നൽകുക",
+        "required": "ദയവായി വീണ്ടെടുത്ത രേഖകളുടെ എണ്ണം നൽകുക"
       }
     },
     "prompt": {

src/assets/locale/pt-BR/settings.json

@@ -311,6 +311,11 @@
         "label": "Limite Padrão de Upload de Arquivos da Base de Conhecimento",
         "placeholder": "Digite o limite padrão de upload de arquivos (ex: 10)",
         "required": "Por favor, insira o limite padrão de upload de arquivos"
+      },
+      "noOfRetrievedDocs": {
+        "label": "Número de Documentos Recuperados",
+        "placeholder": "Digite o Número de Documentos Recuperados",
+        "required": "Por favor, insira o número de documentos recuperados"
       }
     },
     "prompt": {

src/assets/locale/ru/settings.json

@@ -312,7 +312,13 @@
       "totalFilePerKB": {
         "label": "Лимит загрузки файлов по умолчанию для базы знаний",
         "placeholder": "Введите лимит загрузки файлов по умолчанию (например, 10)",
-        "required": "Пожалуйста, введите лимит загрузки файлов по умолчанию"      }
+        "required": "Пожалуйста, введите лимит загрузки файлов по умолчанию"
+      },
+      "noOfRetrievedDocs": {
+        "label": "Количество извлеченных документов",
+        "placeholder": "Введите количество извлеченных документов",
+        "required": "Пожалуйста, введите количество извлеченных документов"
+      }
     },
     "prompt": {
       "label": "Настройка системной подсказки RAG",

src/assets/locale/zh/settings.json

@@ -316,6 +316,11 @@
         "label": "知识库默认文件上传限制",
         "placeholder": "输入默认文件上传限制（例如：10）",
         "required": "请输入默认文件上传限制"
+      },
+      "noOfRetrievedDocs": {
+        "label": "检索文档数量",
+        "placeholder": "输入检索文档数量",
+        "required": "请输入检索文档数量"
       }
     },
     "prompt": {

src/components/Option/Settings/rag.tsx

@@ -10,7 +10,7 @@ import {
 } from "~/services/ollama"
 import { SettingPrompt } from "./prompt"
 import { useTranslation } from "react-i18next"
-import { getTotalFilePerKB } from "@/services/app"
+import { getNoOfRetrievedDocs, getTotalFilePerKB } from "@/services/app"
 import { SidepanelRag } from "./sidepanel-rag"
 
 export const RagSettings = () => {
@@ -21,20 +21,28 @@ export const RagSettings = () => {
   const { data: ollamaInfo, status } = useQuery({
     queryKey: ["fetchRAGSettings"],
     queryFn: async () => {
-      const [allModels, chunkOverlap, chunkSize, defaultEM, totalFilePerKB] =
+      const [
-        await Promise.all([
+        allModels,
+        chunkOverlap,
+        chunkSize,
+        defaultEM,
+        totalFilePerKB,
+        noOfRetrievedDocs
+      ] = await Promise.all([
         getAllModels({ returnEmpty: true }),
         defaultEmbeddingChunkOverlap(),
         defaultEmbeddingChunkSize(),
         defaultEmbeddingModelForRag(),
-          getTotalFilePerKB()
+        getTotalFilePerKB(),
+        getNoOfRetrievedDocs()
       ])
       return {
         models: allModels,
         chunkOverlap,
         chunkSize,
         defaultEM,
-        totalFilePerKB
+        totalFilePerKB,
+        noOfRetrievedDocs
       }
     }
   })
@@ -45,8 +53,15 @@ export const RagSettings = () => {
       chunkSize: number
       overlap: number
       totalFilePerKB: number
+      noOfRetrievedDocs: number
     }) => {
-      await saveForRag(data.model, data.chunkSize, data.overlap,  data.totalFilePerKB)
+      await saveForRag(
+        data.model,
+        data.chunkSize,
+        data.overlap,
+        data.totalFilePerKB,
+        data.noOfRetrievedDocs
+      )
       return true
     },
     onSuccess: () => {
@@ -75,14 +90,16 @@ export const RagSettings = () => {
                   model: data.defaultEM,
                   chunkSize: data.chunkSize,
                   overlap: data.chunkOverlap,
-                  totalFilePerKB: data.totalFilePerKB
+                  totalFilePerKB: data.totalFilePerKB,
+                  noOfRetrievedDocs: data.noOfRetrievedDocs
                 })
               }}
               initialValues={{
                 chunkSize: ollamaInfo?.chunkSize,
                 chunkOverlap: ollamaInfo?.chunkOverlap,
                 defaultEM: ollamaInfo?.defaultEM,
-                totalFilePerKB: ollamaInfo?.totalFilePerKB
+                totalFilePerKB: ollamaInfo?.totalFilePerKB,
+                noOfRetrievedDocs: ollamaInfo?.noOfRetrievedDocs
               }}>
               <Form.Item
                 name="defaultEM"
@@ -142,6 +159,23 @@ export const RagSettings = () => {
                 />
               </Form.Item>
 
+              <Form.Item
+                name="noOfRetrievedDocs"
+                label={t("rag.ragSettings.noOfRetrievedDocs.label")}
+                rules={[
+                  {
+                    required: true,
+                    message: t("rag.ragSettings.noOfRetrievedDocs.required")
+                  }
+                ]}>
+                <InputNumber
+                  style={{ width: "100%" }}
+                  placeholder={t(
+                    "rag.ragSettings.noOfRetrievedDocs.placeholder"
+                  )}
+                />
+              </Form.Item>
+
               <Form.Item
                 name="totalFilePerKB"
                 label={t("rag.ragSettings.totalFilePerKB.label")}
@@ -153,11 +187,11 @@ export const RagSettings = () => {
                 ]}>
                 <InputNumber
                   style={{ width: "100%" }}
+                  min={1}
                   placeholder={t("rag.ragSettings.totalFilePerKB.placeholder")}
                 />
               </Form.Item>
 
-
               <div className="flex justify-end">
                 <SaveButton disabled={isSaveRAGPending} btnType="submit" />
               </div>

src/hooks/useMessageOption.tsx

@@ -32,6 +32,7 @@ import { useStorage } from "@plasmohq/storage/hook"
 import { useStoreChatModelSettings } from "@/store/model"
 import { getAllDefaultModelSettings } from "@/services/model-settings"
 import { pageAssistModel } from "@/models"
+import { getNoOfRetrievedDocs } from "@/services/app"
 
 export const useMessageOption = () => {
   const {
@@ -680,8 +681,9 @@ export const useMessageOption = () => {
         const response = await questionOllama.invoke(promptForQuestion)
         query = response.content.toString()
       }
+      const docSize = await getNoOfRetrievedDocs()
 
-      const docs = await vectorstore.similaritySearch(query, 4)
+      const docs = await vectorstore.similaritySearch(query, docSize)
       const context = formatDocs(docs)
       const source = docs.map((doc) => {
         return {

src/services/app.ts

@@ -109,3 +109,12 @@ export const getTotalFilePerKB = async (): Promise<number> => {
 export const setTotalFilePerKB = async (totalFilePerKB: number): Promise<void> => {
   await storage.set("totalFilePerKB", totalFilePerKB);
 };
+
+export const getNoOfRetrievedDocs = async (): Promise<number> => {
+  const noOfRetrievedDocs = await storage.get<number>("noOfRetrievedDocs");
+  return noOfRetrievedDocs || 4 
+}
+
+export const setNoOfRetrievedDocs = async (noOfRetrievedDocs: number): Promise<void> => {
+  await storage.set("noOfRetrievedDocs", noOfRetrievedDocs);
+}

src/services/ollama.ts

@@ -2,7 +2,7 @@ import { Storage } from "@plasmohq/storage"
 import { cleanUrl } from "../libs/clean-url"
 import { urlRewriteRuntime } from "../libs/runtime"
 import { getChromeAIModel } from "./chrome"
-import { setTotalFilePerKB } from "./app"
+import { setNoOfRetrievedDocs, setTotalFilePerKB } from "./app"
 
 const storage = new Storage()
 
@@ -326,12 +326,14 @@ export const saveForRag = async (
   model: string,
   chunkSize: number,
   overlap: number,
-  totalFilePerKB: number
+  totalFilePerKB: number,
+  noOfRetrievedDocs: number
 ) => {
   await setDefaultEmbeddingModelForRag(model)
   await setDefaultEmbeddingChunkSize(chunkSize)
   await setDefaultEmbeddingChunkOverlap(overlap)
   await setTotalFilePerKB(totalFilePerKB)
+  await setNoOfRetrievedDocs(noOfRetrievedDocs)
 }
 
 export const getWebSearchPrompt = async () => {