feat: admin set pgIvP

admin/service/schema.js

@@ -105,6 +105,10 @@ const SystemSchema = new mongoose.Schema({
     type: Number,
     default: 10
   },
+  pgIvfflatProbe: {
+    type: Number,
+    default: 10
+  },
   sensitiveCheck: {
     type: Boolean,
     default: false

admin/src/fields.ts

@@ -45,5 +45,6 @@ export const SystemFields = [
   createTextField('gpt4Key', { label: 'gpt4Key' }),
   createTextField('vectorMaxProcess', { label: '向量最大进程' }),
   createTextField('qaMaxProcess', { label: 'qa最大进程' }),
+  createTextField('pgIvfflatProbe', { label: 'pg 探针数量' }),
   createTextField('sensitiveCheck', { label: '敏感词校验(true,false)' })
 ];

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

@@ -98,6 +98,7 @@ export async function appKbSearch({
   // search kb
   const res: any = await PgClient.query(
     `BEGIN;
+    SET LOCAL ivfflat.probes = ${global.systemEnv.pgIvfflatProbe || 10};
     select id,q,a,source from modelData where kb_id IN (${model.chat.relatedKbs
       .map((item) => `'${item}'`)
       .join(',')}) AND vector <#> '[${promptVector[0]}]' < -${similarity} order by vector <#> '[${
@@ -106,7 +107,7 @@ export async function appKbSearch({
     COMMIT;`
   );
 
-  const searchRes: QuoteItemType[] = res?.[1]?.rows || [];
+  const searchRes: QuoteItemType[] = res?.[2]?.rows || [];
 
   // filter same search result
   const idSet = new Set<string>();

client/src/service/models/system.ts

@@ -21,6 +21,10 @@ const SystemSchema = new Schema({
     type: Number,
     default: 10
   },
+  pgIvfflatProbe: {
+    type: Number,
+    default: 10
+  },
   sensitiveCheck: {
     type: Boolean,
     default: false

client/src/service/mongo.ts

@@ -20,6 +20,7 @@ export async function connectToDatabase(): Promise<void> {
     gpt4Key: process.env.GPT4KEY || '',
     vectorMaxProcess: 10,
     qaMaxProcess: 10,
+    pgIvfflatProbe: 10,
     sensitiveCheck: false
   };
   // proxy obj

client/src/types/index.d.ts

@@ -28,6 +28,7 @@ declare global {
     gpt4Key: string;
     vectorMaxProcess: number;
     qaMaxProcess: number;
+    pgIvfflatProbe: number;
     sensitiveCheck: boolean;
   };