iod/page-assist
7
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Add Number of Retrieved Documents on RAG

Browse Source
This commit is contained in:
n4ze3m 2024-08-20 18:32:01 +05:30
parent 564b137e66
commit 35b8579028
14 changed files with 133 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
src/assets/locale/en/settings.json
View File

@ -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": {

10
src/assets/locale/es/settings.json
View File

@ -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",

10
src/assets/locale/fa/settings.json
View File

@ -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",

10
src/assets/locale/fr/settings.json
View File

@ -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",

10
src/assets/locale/it/settings.json
View File

@ -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",

5
src/assets/locale/ja-JP/settings.json
View File

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

5
src/assets/locale/ml/settings.json
View File

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

5
src/assets/locale/pt-BR/settings.json
View File

@ -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": {

8
src/assets/locale/ru/settings.json
View File

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

5
src/assets/locale/zh/settings.json
View File

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

52
src/components/Option/Settings/rag.tsx
View File

@ -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] =
await Promise.all([
const [
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>

4
src/hooks/useMessageOption.tsx
View File

@ -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 {

9
src/services/app.ts
View File

@ -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);
}

6
src/services/ollama.ts
View File

@ -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 () => {