perf: kb-add last question to search

17a42ac0 · archer · e384893a · 17a42ac0 · 17a42ac0 · 17a42ac0
Commit 17a42ac0 authored May 03, 2023 by archer
7 changed files
--- a/public/docs/versionIntro.md
+++ b/public/docs/versionIntro.md
 ### Fast GPT V3.1
+- 优化 - 知识库搜索，会将上一个问题并入搜索范围。
 - 优化 - 模型结构设计，不再区分知识库和对话模型，而是通过开关的形式，手动选择手否需要进行知识库搜索。
 - 新增 - 模型共享市场，可以使用其他用户分享的模型。
 - 新增 - 邀请好友注册功能。
--- a/src/pages/api/chat/chat.ts
+++ b/src/pages/api/chat/chat.ts
@@ -58,7 +58,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
      const { code, searchPrompt } = await searchKb({
        userApiKey,
        systemApiKey,
-        text: prompt.value,
+        prompts,
        similarity: ModelVectorSearchModeMap[model.chat.searchMode]?.similarity,
        model,
        userId

--- a/src/pages/api/openapi/chat/chat.ts
+++ b/src/pages/api/openapi/chat/chat.ts
@@ -66,7 +66,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
      const { code, searchPrompt } = await searchKb({
        systemApiKey: apiKey,
-        text: prompts[prompts.length - 1].value,
+        prompts,
        similarity,
        model,
        userId

--- a/src/pages/api/openapi/chat/lafGpt.ts
+++ b/src/pages/api/openapi/chat/lafGpt.ts
@@ -118,7 +118,7 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
    const { searchPrompt } = await searchKb({
      systemApiKey: apiKey,
      similarity: ModelVectorSearchModeMap[model.chat.searchMode]?.similarity,
-      text: prompt.value,
+      prompts,
      model,
      userId
    });

--- a/src/service/events/generateVector.ts
+++ b/src/service/events/generateVector.ts
@@ -60,8 +60,8 @@ export async function generateVector(next = false): Promise<any> {
    }
    // 生成词向量
-    const { vector } = await openaiCreateEmbedding({
+    const { vectors } = await openaiCreateEmbedding({
-      text: dataItem.q,
+      textArr: [dataItem.q],
      userId: dataItem.userId,
      userApiKey,
      systemApiKey
@@ -70,7 +70,7 @@ export async function generateVector(next = false): Promise<any> {
    // 更新 pg 向量和状态数据
    await PgClient.update('modelData', {
      values: [
-        { key: 'vector', value: `[${vector}]` },
+        { key: 'vector', value: `[${vectors[0]}]` },
        { key: 'status', value: `ready` }
      ],
      where: [['id', dataId]]

--- a/src/service/plugins/searchKb.ts
+++ b/src/service/plugins/searchKb.ts
@@ -4,6 +4,7 @@ import { ModelSchema } from '@/types/mongoSchema';
 import { openaiCreateEmbedding } from '../utils/chat/openai';
 import { ChatRoleEnum } from '@/constants/chat';
 import { sliceTextByToken } from '@/utils/chat';
+import { ChatItemSimpleType } from '@/types/chat';
 /**
 *  use openai embedding search kb
@@ -11,14 +12,14 @@ import { sliceTextByToken } from '@/utils/chat';
 export const searchKb = async ({
  userApiKey,
  systemApiKey,
-  text,
+  prompts,
  similarity = 0.2,
  model,
  userId
 }: {
  userApiKey?: string;
  systemApiKey: string;
-  text: string;
+  prompts: ChatItemSimpleType[];
  model: ModelSchema;
  userId: string;
  similarity?: number;
@@ -29,30 +30,56 @@ export const searchKb = async ({
    value: string;
  };
 }> => {
-  const modelConstantsData = ChatModelMap[model.chat.chatModel];
+  async function search(textArr: string[] = []) {
+    // 获取提示词的向量
+    const { vectors: promptVectors } = await openaiCreateEmbedding({
+      userApiKey,
+      systemApiKey,
+      userId,
+      textArr
+    });
+    const searchRes = await Promise.all(
+      promptVectors.map((promptVector) =>
+        PgClient.select<{ id: string; q: string; a: string }>('modelData', {
+          fields: ['id', 'q', 'a'],
+          where: [
+            ['status', ModelDataStatusEnum.ready],
+            'AND',
+            ['model_id', model._id],
+            'AND',
+            `vector <=> '[${promptVector}]' < ${similarity}`
+          ],
+          order: [{ field: 'vector', mode: `<=> '[${promptVector}]'` }],
+          limit: 20
+        }).then((res) => res.rows)
+      )
+    );
-  // 获取提示词的向量
+    // Remove repeat record
-  const { vector: promptVector } = await openaiCreateEmbedding({
+    const idSet = new Set<string>();
-    userApiKey,
+    const filterSearch = searchRes.map((search) =>
-    systemApiKey,
+      search.filter((item) => {
-    userId,
+        if (idSet.has(item.id)) {
-    text
+          return false;
-  });
+        }
+        idSet.add(item.id);
+        return true;
+      })
+    );
+    return filterSearch.map((item) => item.map((item) => `${item.q}\n${item.a}`).join('\n'));
+  }
+  const modelConstantsData = ChatModelMap[model.chat.chatModel];
-  const vectorSearch = await PgClient.select<{ q: string; a: string }>('modelData', {
+  // search three times
-    fields: ['q', 'a'],
+  const userPrompts = prompts.filter((item) => item.obj === 'Human');
-    where: [
-      ['status', ModelDataStatusEnum.ready],
-      'AND',
-      ['model_id', model._id],
-      'AND',
-      `vector <=> '[${promptVector}]' < ${similarity}`
-    ],
-    order: [{ field: 'vector', mode: `<=> '[${promptVector}]'` }],
-    limit: 20
-  });
-  const systemPrompts: string[] = vectorSearch.rows.map((item) => `${item.q}\n${item.a}`);
+  const searchArr: string[] = [
+    userPrompts[userPrompts.length - 1].value,
+    userPrompts[userPrompts.length - 2]?.value
+  ].filter((item) => item);
+  const systemPrompts = await search(searchArr);
  // filter system prompt
  if (
@@ -80,13 +107,24 @@ export const searchKb = async ({
    };
  }
-  // 有匹配情况下，system 添加知识库内容。
+  /* 有匹配情况下，system 添加知识库内容。 */
-  // 系统提示词过滤，最多 65% tokens
-  const filterSystemPrompt = sliceTextByToken({
+  // filter system prompts. max 70% tokens
-    model: model.chat.chatModel,
+  const filterRateMap: Record<number, number[]> = {
-    text: systemPrompts.join('\n'),
+    1: [0.7],
-    length: Math.floor(modelConstantsData.contextMaxToken * 0.65)
+    2: [0.5, 0.2]
-  });
+  };
+  const filterRate = filterRateMap[systemPrompts.length] || filterRateMap[0];
+  const filterSystemPrompt = filterRate
+    .map((rate, i) =>
+      sliceTextByToken({
+        model: model.chat.chatModel,
+        text: systemPrompts[i],
+        length: Math.floor(modelConstantsData.contextMaxToken * rate)
+      })
+    )
+    .join('\n');
  return {
    code: 200,

--- a/src/service/utils/chat/openai.ts
+++ b/src/service/utils/chat/openai.ts
@@ -22,12 +22,12 @@ export const openaiCreateEmbedding = async ({
  userApiKey,
  systemApiKey,
  userId,
-  text
+  textArr
 }: {
  userApiKey?: string;
  systemApiKey: string;
  userId: string;
-  text: string;
+  textArr: string[];
 }) => {
  // 获取 chatAPI
  const chatAPI = getOpenAIApi(userApiKey || systemApiKey);
@@ -37,7 +37,7 @@ export const openaiCreateEmbedding = async ({
    .createEmbedding(
      {
        model: embeddingModel,
-        input: text
+        input: textArr
      },
      {
        timeout: 60000,
@@ -46,18 +46,18 @@ export const openaiCreateEmbedding = async ({
    )
    .then((res) => ({
      tokenLen: res.data.usage.total_tokens || 0,
-      vector: res.data.data?.[0]?.embedding || []
+      vectors: res.data.data.map((item) => item.embedding)
    }));
  pushGenerateVectorBill({
    isPay: !userApiKey,
    userId,
-    text,
+    text: textArr.join(''),
    tokenLen: res.tokenLen
  });
  return {
-    vector: res.vector,
+    vectors: res.vectors,
    chatAPI
  };
 };