From a837552b56b5406f7481f23d32fe1914ece62c37 Mon Sep 17 00:00:00 2001
From: archer <545436317@qq.com>
Date: Tue, 9 May 2023 16:32:02 +0800
Subject: [PATCH] perf: search prompt

---
 src/pages/api/chat/chat.ts                    |  5 ++-
 src/pages/chat/index.tsx                      |  2 +-
 .../detail/components/SelectFileModal.tsx     |  8 +---
 src/service/plugins/searchKb.ts               | 37 ++++++++++---------
 4 files changed, 25 insertions(+), 27 deletions(-)

diff --git a/src/pages/api/chat/chat.ts b/src/pages/api/chat/chat.ts
index 640ea0fd3..20112cce0 100644
--- a/src/pages/api/chat/chat.ts
+++ b/src/pages/api/chat/chat.ts
@@ -55,7 +55,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
 
     // 使用了知识库搜索
     if (model.chat.useKb) {
-      const { code, searchPrompt } = await searchKb({
+      const { code, searchPrompt, aiPrompt } = await searchKb({
         userOpenAiKey,
         prompts,
         similarity: ModelVectorSearchModeMap[model.chat.searchMode]?.similarity,
@@ -68,6 +68,9 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
         return res.send(searchPrompt?.value);
       }
 
+      if (aiPrompt) {
+        prompts.splice(prompts.length - 1, 0, aiPrompt);
+      }
       searchPrompt && prompts.unshift(searchPrompt);
     } else {
       // 没有用知识库搜索，仅用系统提示词
diff --git a/src/pages/chat/index.tsx b/src/pages/chat/index.tsx
index f6caa882a..ff4af3563 100644
--- a/src/pages/chat/index.tsx
+++ b/src/pages/chat/index.tsx
@@ -508,7 +508,7 @@ const Chat = ({
       isLeavePage.current = true;
       controller.current?.abort();
     };
-  }, []);
+  }, [modelId, chatId]);
 
   return (
     <Flex
diff --git a/src/pages/model/components/detail/components/SelectFileModal.tsx b/src/pages/model/components/detail/components/SelectFileModal.tsx
index 8bd5acce3..56c53ed17 100644
--- a/src/pages/model/components/detail/components/SelectFileModal.tsx
+++ b/src/pages/model/components/detail/components/SelectFileModal.tsx
@@ -61,13 +61,7 @@ const SelectFileModal = ({
   const { openConfirm, ConfirmChild } = useConfirm({
     content: `确认导入该文件，需要一定时间进行拆解，该任务无法终止！如果余额不足，未完成的任务会被直接清除。一共 ${
       splitRes.chunks.length
-    } 组。${
-      splitRes.tokens
-        ? `大约 ${splitRes.tokens} 个tokens, 约 ${formatPrice(
-            splitRes.tokens * modeMap[mode].price
-          )} 元`
-        : ''
-    }`
+    } 组。${splitRes.tokens ? `大约 ${splitRes.tokens} 个tokens。` : ''}`
   });
 
   const onSelectFile = useCallback(
diff --git a/src/service/plugins/searchKb.ts b/src/service/plugins/searchKb.ts
index 692281185..7e9ef3405 100644
--- a/src/service/plugins/searchKb.ts
+++ b/src/service/plugins/searchKb.ts
@@ -24,7 +24,11 @@ export const searchKb = async ({
 }): Promise<{
   code: 200 | 201;
   searchPrompt?: {
-    obj: `${ChatRoleEnum}`;
+    obj: `${ChatRoleEnum.System}`;
+    value: string;
+  };
+  aiPrompt?: {
+    obj: `${ChatRoleEnum.AI}`;
     value: string;
   };
 }> => {
@@ -85,24 +89,11 @@ export const searchKb = async ({
   };
   const filterRate = filterRateMap[systemPrompts.length] || filterRateMap[0];
 
-  // count fixed system prompt
-  const fixedSystemPrompt = `
-${model.chat.systemPrompt}
-${
-  model.chat.searchMode === ModelVectorSearchModeEnum.hightSimilarity ? '不回答知识库外的内容.' : ''
-}
-知识库内容为:`;
-  const fixedSystemTokens = modelToolMap[model.chat.chatModel].countTokens({
-    messages: [{ obj: 'System', value: fixedSystemPrompt }]
-  });
-
-  const maxTokens = modelConstantsData.systemMaxToken - fixedSystemTokens;
-
   const filterSystemPrompt = filterRate
     .map((rate, i) =>
       modelToolMap[model.chat.chatModel].sliceText({
         text: systemPrompts[i],
-        length: Math.floor(maxTokens * rate)
+        length: Math.floor(modelConstantsData.systemMaxToken * rate)
       })
     )
     .join('\n');
@@ -112,7 +103,7 @@ ${
     return {
       code: 201,
       searchPrompt: {
-        obj: ChatRoleEnum.AI,
+        obj: ChatRoleEnum.System,
         value: '对不起，你的问题不在知识库中。'
       }
     };
@@ -135,7 +126,17 @@ ${
     code: 200,
     searchPrompt: {
       obj: ChatRoleEnum.System,
-      value: `${fixedSystemPrompt}'${filterSystemPrompt}'`
-    }
+      value: `知识库:'${filterSystemPrompt}'`
+    },
+    aiPrompt: ModelVectorSearchModeEnum.hightSimilarity
+      ? {
+          obj: 'AI',
+          value: `我来玩一个问答游戏,规则为:
+1.我完全忘记我已有的知识
+2.我只能回答关于"${model.chat.systemPrompt || model.name}"的问题
+3.我只能从知识库中选择内容进行回答
+4.如果问题不在知识库中,我会回答"我不知道。"`
+        }
+      : undefined
   };
 };