perf: search kb model

src/service/tools/searchKb.ts

@@ -1,6 +1,8 @@
 import { openaiCreateEmbedding } from '../utils/openai';
 import { PgClient } from '@/service/pg';
-import { ModelDataStatusEnum } from '@/constants/model';
+import { ModelDataStatusEnum, ModelVectorSearchModeEnum } from '@/constants/model';
+import { ModelSchema } from '@/types/mongoSchema';
+import { systemPromptFilter } from '../utils/tools';
 
 /**
  *  use openai embedding search kb
@@ -10,16 +12,22 @@ export const searchKb_openai = async ({
   isPay,
   text,
   similarity,
-  modelId,
+  model,
   userId
 }: {
   apiKey: string;
   isPay: boolean;
   text: string;
-  modelId: string;
+  model: ModelSchema;
   userId: string;
   similarity: number;
-}) => {
+}): Promise<{
+  code: 200 | 201;
+  searchPrompt?: {
+    obj: 'Human' | 'AI' | 'SYSTEM';
+    value: string;
+  };
+}> => {
   // 获取提示词的向量
   const { vector: promptVector } = await openaiCreateEmbedding({
     isPay,
@@ -28,12 +36,12 @@ export const searchKb_openai = async ({
     text
   });
 
-  const vectorSearch = await PgClient.select<{ id: string; q: string; a: string }>('modelData', {
-    fields: ['id', 'q', 'a'],
+  const vectorSearch = await PgClient.select<{ q: string; a: string }>('modelData', {
+    fields: ['q', 'a'],
     where: [
       ['status', ModelDataStatusEnum.ready],
       'AND',
-      ['model_id', modelId],
+      ['model_id', model._id],
       'AND',
       `vector <=> '[${promptVector}]' < ${similarity}`
     ],
@@ -43,5 +51,51 @@ export const searchKb_openai = async ({
 
   const systemPrompts: string[] = vectorSearch.rows.map((item) => `${item.q}\n${item.a}`);
 
-  return { systemPrompts };
+  // filter system prompt
+  if (
+    systemPrompts.length === 0 &&
+    model.chat.searchMode === ModelVectorSearchModeEnum.hightSimilarity
+  ) {
+    return {
+      code: 201,
+      searchPrompt: {
+        obj: 'AI',
+        value: '对不起，你的问题不在知识库中。'
+      }
+    };
+  }
+  /* 高相似度+无上下文，不添加额外知识,仅用系统提示词 */
+  if (systemPrompts.length === 0 && model.chat.searchMode === ModelVectorSearchModeEnum.noContext) {
+    return {
+      code: 200,
+      searchPrompt: model.chat.systemPrompt
+        ? {
+            obj: 'SYSTEM',
+            value: model.chat.systemPrompt
+          }
+        : undefined
+    };
+  }
+
+  // 有匹配情况下，system 添加知识库内容。
+  // 系统提示词过滤，最多 2500 tokens
+  const filterSystemPrompt = systemPromptFilter({
+    model: model.chat.chatModel,
+    prompts: systemPrompts,
+    maxTokens: 2500
+  });
+
+  return {
+    code: 200,
+    searchPrompt: {
+      obj: 'SYSTEM',
+      value: `
+${model.chat.systemPrompt}
+${
+  model.chat.searchMode === ModelVectorSearchModeEnum.hightSimilarity ? `不回答知识库外的内容.` : ''
+}
+知识库内容为: ${filterSystemPrompt}'
+`
+    }
+  };
 };