feat: maxlink and search sql

ajavaer · Jun 5, 2023 · 2f1506b · 2f1506b
1 parent 48fbc74
commit 2f1506b
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 11 deletions.
diff --git a/src/pages/api/openapi/kb/appKbSearch.ts b/src/pages/api/openapi/kb/appKbSearch.ts
@@ -91,16 +91,17 @@ export async function appKbSearch({
   });
 
   // search kb
-  const { rows: searchRes } = await PgClient.select<QuoteItemType>('modelData', {
-    fields: ['id', 'q', 'a', 'source'],
-    where: [
-      `kb_id IN (${model.chat.relatedKbs.map((item) => `'${item}'`).join(',')})`,
-      'AND',
-      `vector <#> '[${promptVector[0]}]' < ${similarity}`
-    ],
-    order: [{ field: 'vector', mode: `<#> '[${promptVector[0]}]'` }],
-    limit: 8
-  });
+  const res: any = await PgClient.query(
+    `BEGIN;
+    SET LOCAL ivfflat.probes = ${process.env.PG_IVFFLAT_PROBE || 100};
+    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 <#> '[${
+      promptVector[0]
+    }]' limit 8;
+    COMMIT;`
+  );
+  const searchRes: QuoteItemType[] = res?.[2]?.rows || [];
 
   // filter same search result
   const idSet = new Set<string>();

diff --git a/src/service/pg.ts b/src/service/pg.ts
@@ -13,7 +13,7 @@ export const connectPg = async () => {
     user: process.env.PG_USER,
     password: process.env.PG_PASSWORD,
     database: process.env.PG_DB_NAME,
-    max: maxLink,
+    max: Math.floor(maxLink * 0.5),
     idleTimeoutMillis: 60000,
     connectionTimeoutMillis: 20000
   });