feat: chat quote

src/pages/api/openapi/kb/appKbSearch.ts

@@ -0,0 +1,224 @@
+import type { NextApiRequest, NextApiResponse } from 'next';
+import { jsonRes } from '@/service/response';
+import { authUser } from '@/service/utils/auth';
+import { PgClient } from '@/service/pg';
+import { withNextCors } from '@/service/utils/tools';
+import type { ChatItemSimpleType } from '@/types/chat';
+import type { ModelSchema } from '@/types/mongoSchema';
+import { ModelVectorSearchModeEnum } from '@/constants/model';
+import { authModel } from '@/service/utils/auth';
+import { ChatModelMap } from '@/constants/model';
+import { ChatRoleEnum } from '@/constants/chat';
+import { openaiEmbedding } from '../plugin/openaiEmbedding';
+import { ModelDataStatusEnum } from '@/constants/model';
+import { modelToolMap } from '@/utils/plugin';
+
+export type QuoteItemType = { id: string; q: string; a: string };
+type Props = {
+  prompts: ChatItemSimpleType[];
+  similarity: number;
+  appId: string;
+};
+type Response = {
+  code: 200 | 201;
+  rawSearch: QuoteItemType[];
+  searchPrompts: {
+    obj: ChatRoleEnum;
+    value: string;
+  }[];
+};
+
+export default withNextCors(async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
+  try {
+    const { userId } = await authUser({ req });
+
+    if (!userId) {
+      throw new Error('userId is empty');
+    }
+
+    const { prompts, similarity, appId } = req.body as Props;
+
+    if (!similarity || !Array.isArray(prompts) || !appId) {
+      throw new Error('params is error');
+    }
+
+    // auth model
+    const { model } = await authModel({
+      modelId: appId,
+      userId
+    });
+
+    const result = await appKbSearch({
+      userId,
+      prompts,
+      similarity,
+      model
+    });
+
+    jsonRes<Response>(res, {
+      data: result
+    });
+  } catch (err) {
+    console.log(err);
+    jsonRes(res, {
+      code: 500,
+      error: err
+    });
+  }
+});
+
+export async function appKbSearch({
+  model,
+  userId,
+  prompts,
+  similarity
+}: {
+  userId: string;
+  prompts: ChatItemSimpleType[];
+  similarity: number;
+  model: ModelSchema;
+}): Promise<Response> {
+  const modelConstantsData = ChatModelMap[model.chat.chatModel];
+
+  // search two times.
+  const userPrompts = prompts.filter((item) => item.obj === 'Human');
+
+  const input: string[] = [
+    userPrompts[userPrompts.length - 1].value,
+    userPrompts[userPrompts.length - 2]?.value
+  ].filter((item) => item);
+
+  // get vector
+  const promptVectors = await openaiEmbedding({
+    userId,
+    input
+  });
+
+  // search kb
+  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',
+          `kb_id IN (${model.chat.relatedKbs.map((item) => `'${item}'`).join(',')})`,
+          'AND',
+          `vector <=> '[${promptVector}]' < ${similarity}`
+        ],
+        order: [{ field: 'vector', mode: `<=> '[${promptVector}]'` }],
+        limit: promptVectors.length === 1 ? 15 : 10
+      }).then((res) => res.rows)
+    )
+  );
+
+  // filter same search result
+  const idSet = new Set<string>();
+  const filterSearch = searchRes.map((search) =>
+    search.filter((item) => {
+      if (idSet.has(item.id)) {
+        return false;
+      }
+      idSet.add(item.id);
+      return true;
+    })
+  );
+
+  // slice search result by rate.
+  const sliceRateMap: Record<number, number[]> = {
+    1: [1],
+    2: [0.7, 0.3]
+  };
+  const sliceRate = sliceRateMap[searchRes.length] || sliceRateMap[0];
+  // 计算固定提示词的 token 数量
+  const fixedPrompts = [
+    // user system prompt
+    ...(model.chat.systemPrompt
+      ? [
+          {
+            obj: ChatRoleEnum.System,
+            value: model.chat.systemPrompt
+          }
+        ]
+      : model.chat.searchMode === ModelVectorSearchModeEnum.noContext
+      ? [
+          {
+            obj: ChatRoleEnum.System,
+            value: `知识库是关于"${model.name}"的内容,根据知识库内容回答问题.`
+          }
+        ]
+      : [
+          {
+            obj: ChatRoleEnum.System,
+            value: `玩一个问答游戏,规则为:
+1.你完全忘记你已有的知识
+2.你只回答关于"${model.name}"的问题
+3.你只从知识库中选择内容进行回答
+4.如果问题不在知识库中,你会回答:"我不知道。"
+请务必遵守规则`
+          }
+        ])
+  ];
+  const fixedSystemTokens = modelToolMap[model.chat.chatModel].countTokens({
+    messages: fixedPrompts
+  });
+  const maxTokens = modelConstantsData.systemMaxToken - fixedSystemTokens;
+  const sliceResult = sliceRate.map((rate, i) =>
+    modelToolMap[model.chat.chatModel]
+      .tokenSlice({
+        maxToken: Math.round(maxTokens * rate),
+        messages: filterSearch[i].map((item) => ({
+          obj: ChatRoleEnum.System,
+          value: `${item.q}\n${item.a}`
+        }))
+      })
+      .map((item) => item.value)
+  );
+
+  // slice filterSearch
+  const sliceSearch = filterSearch.map((item, i) => item.slice(0, sliceResult[i].length)).flat();
+
+  //  system prompt
+  const systemPrompt = sliceResult.flat().join('\n').trim();
+
+  /* 高相似度+不回复 */
+  if (!systemPrompt && model.chat.searchMode === ModelVectorSearchModeEnum.hightSimilarity) {
+    return {
+      code: 201,
+      rawSearch: [],
+      searchPrompts: [
+        {
+          obj: ChatRoleEnum.System,
+          value: '对不起，你的问题不在知识库中。'
+        }
+      ]
+    };
+  }
+  /* 高相似度+无上下文，不添加额外知识,仅用系统提示词 */
+  if (!systemPrompt && model.chat.searchMode === ModelVectorSearchModeEnum.noContext) {
+    return {
+      code: 200,
+      rawSearch: [],
+      searchPrompts: model.chat.systemPrompt
+        ? [
+            {
+              obj: ChatRoleEnum.System,
+              value: model.chat.systemPrompt
+            }
+          ]
+        : []
+    };
+  }
+
+  return {
+    code: 200,
+    rawSearch: sliceSearch,
+    searchPrompts: [
+      {
+        obj: ChatRoleEnum.System,
+        value: `知识库:${systemPrompt}`
+      },
+      ...fixedPrompts
+    ]
+  };
+}