From a135f9f5b6510f6062a21d95201e97907b616402 Mon Sep 17 00:00:00 2001
From: balibabu <cike8899@users.noreply.github.com>
Date: Wed, 29 May 2024 16:19:08 +0800
Subject: [PATCH] feat: add rerank models to the project #724 #162 (#966)

### What problem does this PR solve?

Vector similarity weight is displayed incorrectly #965
feat: add rerank models to the project #724 #162
### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
---
 web/src/components/rerank.tsx                 | 57 +++++++++++++++++++
 .../components/similarity-slider/index.tsx    |  2 +-
 web/src/constants/knowledge.ts                |  1 +
 web/src/hooks/llmHooks.ts                     |  1 +
 web/src/interfaces/database/chat.ts           |  2 +
 web/src/locales/en.ts                         | 17 ++++--
 web/src/locales/zh-traditional.ts             | 14 +++--
 web/src/locales/zh.ts                         | 14 +++--
 .../components/knowledge-testing/index.tsx    |  5 +-
 .../testing-control/index.tsx                 |  9 +++
 .../chat/chat-configuration-modal/hooks.ts    | 11 ++++
 .../chat/chat-configuration-modal/index.tsx   |  6 +-
 .../model-setting.tsx                         |  8 +--
 .../prompt-engine.tsx                         |  3 +-
 .../system-model-setting-modal/index.tsx      |  7 +++
 15 files changed, 132 insertions(+), 25 deletions(-)
 create mode 100644 web/src/components/rerank.tsx

diff --git a/web/src/components/rerank.tsx b/web/src/components/rerank.tsx
new file mode 100644
index 000000000..e78fa3aff
--- /dev/null
+++ b/web/src/components/rerank.tsx
@@ -0,0 +1,57 @@
+import { LlmModelType } from '@/constants/knowledge';
+import { useTranslate } from '@/hooks/commonHooks';
+import { useSelectLlmOptionsByModelType } from '@/hooks/llmHooks';
+import { Form, Select, Slider } from 'antd';
+
+type FieldType = {
+  rerank_id?: string;
+  top_k?: number;
+};
+
+export const RerankItem = () => {
+  const { t } = useTranslate('knowledgeDetails');
+  const allOptions = useSelectLlmOptionsByModelType();
+
+  return (
+    <Form.Item
+      label={t('rerankModel')}
+      name={'rerank_id'}
+      tooltip={t('rerankTip')}
+    >
+      <Select
+        options={allOptions[LlmModelType.Rerank]}
+        allowClear
+        placeholder={t('rerankPlaceholder')}
+      />
+    </Form.Item>
+  );
+};
+
+const Rerank = () => {
+  const { t } = useTranslate('knowledgeDetails');
+
+  return (
+    <>
+      <RerankItem></RerankItem>
+      <Form.Item noStyle dependencies={['rerank_id']}>
+        {({ getFieldValue }) => {
+          const rerankId = getFieldValue('rerank_id');
+          return (
+            rerankId && (
+              <Form.Item<FieldType>
+                label={t('topK')}
+                name={'top_k'}
+                initialValue={1024}
+                tooltip={t('topKTip')}
+              >
+                <Slider max={2048} min={1} />
+              </Form.Item>
+            )
+          );
+        }}
+      </Form.Item>
+    </>
+  );
+};
+
+export default Rerank;
diff --git a/web/src/components/similarity-slider/index.tsx b/web/src/components/similarity-slider/index.tsx
index 5df5ceda3..588b83451 100644
--- a/web/src/components/similarity-slider/index.tsx
+++ b/web/src/components/similarity-slider/index.tsx
@@ -26,7 +26,7 @@ const SimilaritySlider = ({ isTooltipShown = false }: IProps) => {
       <Form.Item<FieldType>
         label={t('vectorSimilarityWeight')}
         name={'vector_similarity_weight'}
-        initialValue={0.3}
+        initialValue={1 - 0.3}
         tooltip={isTooltipShown && t('vectorSimilarityWeightTip')}
       >
         <Slider max={1} step={0.01} />
diff --git a/web/src/constants/knowledge.ts b/web/src/constants/knowledge.ts
index 217f78b54..586a15bcc 100644
--- a/web/src/constants/knowledge.ts
+++ b/web/src/constants/knowledge.ts
@@ -47,6 +47,7 @@ export enum LlmModelType {
   Chat = 'chat',
   Image2text = 'image2text',
   Speech2text = 'speech2text',
+  Rerank = 'rerank',
 }
 
 export enum KnowledgeSearchParams {
diff --git a/web/src/hooks/llmHooks.ts b/web/src/hooks/llmHooks.ts
index cf872765a..78e8d947a 100644
--- a/web/src/hooks/llmHooks.ts
+++ b/web/src/hooks/llmHooks.ts
@@ -92,6 +92,7 @@ export const useSelectLlmOptionsByModelType = () => {
     [LlmModelType.Speech2text]: groupOptionsByModelType(
       LlmModelType.Speech2text,
     ),
+    [LlmModelType.Rerank]: groupOptionsByModelType(LlmModelType.Rerank),
   };
 };
 
diff --git a/web/src/interfaces/database/chat.ts b/web/src/interfaces/database/chat.ts
index 5e87a92de..e96bb7c8b 100644
--- a/web/src/interfaces/database/chat.ts
+++ b/web/src/interfaces/database/chat.ts
@@ -47,6 +47,8 @@ export interface IDialog {
   tenant_id: string;
   update_date: string;
   update_time: number;
+  vector_similarity_weight: number;
+  similarity_threshold: number;
 }
 
 export interface IConversation {
diff --git a/web/src/locales/en.ts b/web/src/locales/en.ts
index 0e345ea63..6a4254f6e 100644
--- a/web/src/locales/en.ts
+++ b/web/src/locales/en.ts
@@ -93,15 +93,12 @@ export default {
       progressMsg: 'Progress Msg',
       testingDescription:
         'Final step! After success, leave the rest to Infiniflow AI.',
-      topK: 'Top K',
-      topKTip:
-        "For the computaion cost, not all the retrieved chunk will be computed vector cosine similarity with query. The bigger the 'Top K' is, the higher the recall rate is, the slower the retrieval speed is.",
       similarityThreshold: 'Similarity threshold',
       similarityThresholdTip:
         "We use hybrid similarity score to evaluate distance between two lines of text. It's weighted keywords similarity and vector cosine similarity. If the similarity between query and chunk is less than this threshold, the chunk will be filtered out.",
-      vectorSimilarityWeight: 'Vector similarity weight',
+      vectorSimilarityWeight: 'Keywords similarity weight',
       vectorSimilarityWeightTip:
-        "We use hybrid similarity score to evaluate distance between two lines of text. It's weighted keywords similarity and vector cosine similarity. The sum of both weights is 1.0.",
+        " We use hybrid similarity score to evaluate distance between two lines of text. It's weighted keywords similarity and vector cosine similarity or rerank score(0~1). The sum of both weights is 1.0.",
       testText: 'Test text',
       testTextPlaceholder: 'Please input your question!',
       testingLabel: 'Testing',
@@ -143,6 +140,11 @@ export default {
       chunk: 'Chunk',
       bulk: 'Bulk',
       cancel: 'Cancel',
+      rerankModel: 'Rerank Model',
+      rerankPlaceholder: 'Please select',
+      rerankTip: `If it's empty. It uses embeddings of query and chunks to compuste vector cosine similarity. Otherwise, it uses rerank score in place of  vector cosine similarity.`,
+      topK: 'Top-K',
+      topKTip: `K chunks will be fed into rerank models.`,
     },
     knowledgeConfiguration: {
       titleDescription:
@@ -465,6 +467,8 @@ The above is the content you need to summarize.`,
       sequence2txtModel: 'Sequence2txt model',
       sequence2txtModelTip:
         'The default ASR model all the newly created knowledgebase will use. Use this model to translate voices to corresponding text.',
+      rerankModel: 'Rerank Model',
+      rerankModelTip: `The default rerank model is used to rerank chunks retrieved by users' questions.`,
       workspace: 'Workspace',
       upgrade: 'Upgrade',
       addLlmTitle: 'Add LLM',
@@ -477,7 +481,8 @@ The above is the content you need to summarize.`,
       baseUrlNameMessage: 'Please input your base url!',
       vision: 'Does it support Vision?',
       ollamaLink: 'How to integrate {{name}}',
-      volcModelNameMessage: 'Please input your model name! Format: {"ModelName":"EndpointID"}',
+      volcModelNameMessage:
+        'Please input your model name! Format: {"ModelName":"EndpointID"}',
       addVolcEngineAK: 'VOLC ACCESS_KEY',
       volcAKMessage: 'Please input your VOLC_ACCESS_KEY',
       addVolcEngineSK: 'VOLC SECRET_KEY',
diff --git a/web/src/locales/zh-traditional.ts b/web/src/locales/zh-traditional.ts
index 1369318ac..c22add8a2 100644
--- a/web/src/locales/zh-traditional.ts
+++ b/web/src/locales/zh-traditional.ts
@@ -91,15 +91,12 @@ export default {
       processDuration: '過程持續時間',
       progressMsg: '進度消息',
       testingDescription: '最後一步！成功後，剩下的就交給Infiniflow AI吧。',
-      topK: 'top k',
-      topKTip:
-        '對於計算成本，並非所有檢索到的塊都會計算與查詢的向量餘弦相似度。Top K越大，召回率越高，檢索速度越慢。',
       similarityThreshold: '相似度閾值',
       similarityThresholdTip:
         '我們使用混合相似度得分來評估兩行文本之間的距離。它是加權關鍵詞相似度和向量餘弦相似度。如果查詢和塊之間的相似度小於此閾值，則該塊將被過濾掉。',
-      vectorSimilarityWeight: '向量相似度權重',
+      vectorSimilarityWeight: '關鍵字相似度權重',
       vectorSimilarityWeightTip:
-        '我們使用混合相似度得分來評估兩行文本之間的距離。它是加權關鍵詞相似度和向量餘弦相似度。兩個權重之和為 1.0。',
+        '我們使用混合相似性評分來評估兩行文本之間的距離。它是加權關鍵字相似性和矢量餘弦相似性或rerank得分（0〜1）。兩個權重的總和為1.0。',
       testText: '測試文本',
       testTextPlaceholder: '請輸入您的問題！',
       testingLabel: '測試',
@@ -139,6 +136,11 @@ export default {
       chunk: '解析塊',
       bulk: '批量',
       cancel: '取消',
+      rerankModel: 'rerank模型',
+      rerankPlaceholder: '請選擇',
+      rerankTip: `如果是空的。它使用查詢和塊的嵌入來構成矢量餘弦相似性。否則，它使用rerank評分代替矢量餘弦相似性。`,
+      topK: 'Top-K',
+      topKTip: `K塊將被送入Rerank型號。`,
     },
     knowledgeConfiguration: {
       titleDescription: '在這裡更新您的知識庫詳細信息，尤其是解析方法。',
@@ -429,6 +431,8 @@ export default {
       sequence2txtModel: 'sequence2Txt模型',
       sequence2txtModelTip:
         '所有新創建的知識庫都將使用默認的 ASR 模型。使用此模型將語音翻譯為相應的文本。',
+      rerankModel: 'rerank模型',
+      rerankModelTip: `默認的重讀模型用於用戶問題檢索到重讀塊。`,
       workspace: '工作空間',
       upgrade: '升級',
       addLlmTitle: '添加Llm',
diff --git a/web/src/locales/zh.ts b/web/src/locales/zh.ts
index a32e6e4b2..040f5e1e7 100644
--- a/web/src/locales/zh.ts
+++ b/web/src/locales/zh.ts
@@ -91,15 +91,12 @@ export default {
       processDuration: '过程持续时间',
       progressMsg: '进度消息',
       testingDescription: '最后一步！ 成功后，剩下的就交给Infiniflow AI吧。',
-      topK: 'Top K',
-      topKTip:
-        '对于计算成本，并非所有检索到的块都会计算与查询的向量余弦相似度。 Top K越大，召回率越高，检索速度越慢。',
       similarityThreshold: '相似度阈值',
       similarityThresholdTip:
         '我们使用混合相似度得分来评估两行文本之间的距离。 它是加权关键词相似度和向量余弦相似度。 如果查询和块之间的相似度小于此阈值，则该块将被过滤掉。',
-      vectorSimilarityWeight: '向量相似度权重',
+      vectorSimilarityWeight: '关键字相似度权重',
       vectorSimilarityWeightTip:
-        '我们使用混合相似度得分来评估两行文本之间的距离。 它是加权关键词相似度和向量余弦相似度。 两个权重之和为 1.0。',
+        '我们使用混合相似性评分来评估两行文本之间的距离。它是加权关键字相似性和矢量余弦相似性或rerank得分（0〜1）。两个权重的总和为1.0。',
       testText: '测试文本',
       testTextPlaceholder: '请输入您的问题！',
       testingLabel: '测试',
@@ -140,6 +137,11 @@ export default {
       chunk: '解析块',
       bulk: '批量',
       cancel: '取消',
+      rerankModel: 'Rerank模型',
+      rerankPlaceholder: '请选择',
+      rerankTip: `如果是空的。它使用查询和块的嵌入来构成矢量余弦相似性。否则，它使用rerank评分代替矢量余弦相似性。`,
+      topK: 'Top-K',
+      topKTip: `K块将被送入Rerank型号。`,
     },
     knowledgeConfiguration: {
       titleDescription: '在这里更新您的知识库详细信息，尤其是解析方法。',
@@ -446,6 +448,8 @@ export default {
       sequence2txtModel: 'Sequence2txt模型',
       sequence2txtModelTip:
         '所有新创建的知识库都将使用默认的 ASR 模型。 使用此模型将语音翻译为相应的文本。',
+      rerankModel: 'Rerank模型',
+      rerankModelTip: `默认的重读模型用于用户问题检索到重读块。`,
       workspace: '工作空间',
       upgrade: '升级',
       addLlmTitle: '添加 LLM',
diff --git a/web/src/pages/add-knowledge/components/knowledge-testing/index.tsx b/web/src/pages/add-knowledge/components/knowledge-testing/index.tsx
index 47168611f..478e99410 100644
--- a/web/src/pages/add-knowledge/components/knowledge-testing/index.tsx
+++ b/web/src/pages/add-knowledge/components/knowledge-testing/index.tsx
@@ -15,7 +15,10 @@ const KnowledgeTesting = () => {
 
   const handleTesting = async () => {
     const values = await form.validateFields();
-    testChunk(values);
+    testChunk({
+      ...values,
+      vector_similarity_weight: 1 - values.vector_similarity_weight,
+    });
   };
 
   useEffect(() => {
diff --git a/web/src/pages/add-knowledge/components/knowledge-testing/testing-control/index.tsx b/web/src/pages/add-knowledge/components/knowledge-testing/testing-control/index.tsx
index 2fd9c70ca..441a7402e 100644
--- a/web/src/pages/add-knowledge/components/knowledge-testing/testing-control/index.tsx
+++ b/web/src/pages/add-knowledge/components/knowledge-testing/testing-control/index.tsx
@@ -2,8 +2,11 @@ import SimilaritySlider from '@/components/similarity-slider';
 import { Button, Card, Divider, Flex, Form, Input } from 'antd';
 import { FormInstance } from 'antd/lib';
 
+import Rerank from '@/components/rerank';
 import { useTranslate } from '@/hooks/commonHooks';
+import { useFetchLlmList } from '@/hooks/llmHooks';
 import { useOneNamespaceEffectsLoading } from '@/hooks/storeHooks';
+import { useEffect } from 'react';
 import styles from './index.less';
 
 type FieldType = {
@@ -23,6 +26,11 @@ const TestingControl = ({ form, handleTesting }: IProps) => {
     'testDocumentChunk',
   ]);
   const { t } = useTranslate('knowledgeDetails');
+  const fetchLlmList = useFetchLlmList();
+
+  useEffect(() => {
+    fetchLlmList();
+  }, [fetchLlmList]);
 
   const buttonDisabled =
     !question || (typeof question === 'string' && question.trim() === '');
@@ -37,6 +45,7 @@ const TestingControl = ({ form, handleTesting }: IProps) => {
       <section>
         <Form name="testing" layout="vertical" form={form}>
           <SimilaritySlider isTooltipShown></SimilaritySlider>
+          <Rerank></Rerank>
           <Card size="small" title={t('testText')}>
             <Form.Item<FieldType>
               name={'question'}
diff --git a/web/src/pages/chat/chat-configuration-modal/hooks.ts b/web/src/pages/chat/chat-configuration-modal/hooks.ts
index 5a2ecf8da..28db2b52b 100644
--- a/web/src/pages/chat/chat-configuration-modal/hooks.ts
+++ b/web/src/pages/chat/chat-configuration-modal/hooks.ts
@@ -1,3 +1,4 @@
+import { useFetchLlmList } from '@/hooks/llmHooks';
 import {
   useFetchTenantInfo,
   useSelectTenantInfo,
@@ -16,3 +17,13 @@ export const useFetchModelId = (visible: boolean) => {
 
   return tenantInfo?.llm_id ?? '';
 };
+
+export const useFetchLlmModelOnVisible = (visible: boolean) => {
+  const fetchLlmList = useFetchLlmList();
+
+  useEffect(() => {
+    if (visible) {
+      fetchLlmList();
+    }
+  }, [fetchLlmList, visible]);
+};
diff --git a/web/src/pages/chat/chat-configuration-modal/index.tsx b/web/src/pages/chat/chat-configuration-modal/index.tsx
index d025ba160..fbce78f4a 100644
--- a/web/src/pages/chat/chat-configuration-modal/index.tsx
+++ b/web/src/pages/chat/chat-configuration-modal/index.tsx
@@ -14,7 +14,7 @@ import { variableEnabledFieldMap } from '../constants';
 import { IPromptConfigParameters } from '../interface';
 import { excludeUnEnabledVariables } from '../utils';
 import AssistantSetting from './assistant-setting';
-import { useFetchModelId } from './hooks';
+import { useFetchLlmModelOnVisible, useFetchModelId } from './hooks';
 import ModelSetting from './model-setting';
 import PromptEngine from './prompt-engine';
 
@@ -92,6 +92,7 @@ const ChatConfigurationModal = ({
     const finalValues = {
       dialog_id: initialDialog.id,
       ...nextValues,
+      vector_similarity_weight: 1 - nextValues.vector_similarity_weight,
       prompt_config: {
         ...nextValues.prompt_config,
         parameters: promptEngineRef.current,
@@ -115,6 +116,8 @@ const ChatConfigurationModal = ({
     form.resetFields();
   };
 
+  useFetchLlmModelOnVisible(visible);
+
   const title = (
     <Flex gap={16}>
       <ChatConfigurationAtom></ChatConfigurationAtom>
@@ -142,6 +145,7 @@ const ChatConfigurationModal = ({
           settledModelVariableMap[ModelVariableType.Precise],
         icon: fileList,
         llm_id: initialDialog.llm_id ?? modelId,
+        vector_similarity_weight: 1 - initialDialog.vector_similarity_weight,
       });
     }
   }, [initialDialog, form, visible, modelId]);
diff --git a/web/src/pages/chat/chat-configuration-modal/model-setting.tsx b/web/src/pages/chat/chat-configuration-modal/model-setting.tsx
index 9d3471a1e..8f4502a11 100644
--- a/web/src/pages/chat/chat-configuration-modal/model-setting.tsx
+++ b/web/src/pages/chat/chat-configuration-modal/model-setting.tsx
@@ -10,7 +10,7 @@ import { useEffect } from 'react';
 import { ISegmentedContentProps } from '../interface';
 
 import { useTranslate } from '@/hooks/commonHooks';
-import { useFetchLlmList, useSelectLlmOptions } from '@/hooks/llmHooks';
+import { useSelectLlmOptionsByModelType } from '@/hooks/llmHooks';
 import { Variable } from '@/interfaces/database/chat';
 import { variableEnabledFieldMap } from '../constants';
 import styles from './index.less';
@@ -30,7 +30,7 @@ const ModelSetting = ({
     value: x,
   }));
 
-  const modelOptions = useSelectLlmOptions();
+  const modelOptions = useSelectLlmOptionsByModelType();
 
   const handleParametersChange = (value: ModelVariableType) => {
     const variable = settledModelVariableMap[value];
@@ -56,8 +56,6 @@ const ModelSetting = ({
     }
   }, [form, initialLlmSetting, visible]);
 
-  useFetchLlmList(LlmModelType.Chat);
-
   return (
     <section
       className={classNames({
@@ -70,7 +68,7 @@ const ModelSetting = ({
         tooltip={t('modelTip')}
         rules={[{ required: true, message: t('modelMessage') }]}
       >
-        <Select options={modelOptions} showSearch />
+        <Select options={modelOptions[LlmModelType.Chat]} showSearch />
       </Form.Item>
       <Divider></Divider>
       <Form.Item
diff --git a/web/src/pages/chat/chat-configuration-modal/prompt-engine.tsx b/web/src/pages/chat/chat-configuration-modal/prompt-engine.tsx
index 15af1b62c..4f1b89312 100644
--- a/web/src/pages/chat/chat-configuration-modal/prompt-engine.tsx
+++ b/web/src/pages/chat/chat-configuration-modal/prompt-engine.tsx
@@ -29,6 +29,7 @@ import {
 } from '../interface';
 import { EditableCell, EditableRow } from './editable-cell';
 
+import Rerank from '@/components/rerank';
 import { useTranslate } from '@/hooks/commonHooks';
 import { useSelectPromptConfigParameters } from '../hooks';
 import styles from './index.less';
@@ -172,7 +173,7 @@ const PromptEngine = (
       >
         <Slider max={30} />
       </Form.Item>
-
+      <Rerank></Rerank>
       <section className={classNames(styles.variableContainer)}>
         <Row align={'middle'} justify="end">
           <Col span={7} className={styles.variableAlign}>
diff --git a/web/src/pages/user-setting/setting-model/system-model-setting-modal/index.tsx b/web/src/pages/user-setting/setting-model/system-model-setting-modal/index.tsx
index e3a8149dd..1589c4ad2 100644
--- a/web/src/pages/user-setting/setting-model/system-model-setting-modal/index.tsx
+++ b/web/src/pages/user-setting/setting-model/system-model-setting-modal/index.tsx
@@ -76,6 +76,13 @@ const SystemModelSettingModal = ({
         >
           <Select options={allOptions[LlmModelType.Speech2text]} />
         </Form.Item>
+        <Form.Item
+          label={t('rerankModel')}
+          name="rerank_id"
+          tooltip={t('rerankModelTip')}
+        >
+          <Select options={allOptions[LlmModelType.Rerank]} />
+        </Form.Item>
       </Form>
     </Modal>
   );