diff --git a/src/pages/api/openapi/kb/appKbSearch.ts b/src/pages/api/openapi/kb/appKbSearch.ts index da4e16e35..4b4e0bca9 100644 --- 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('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(); diff --git a/src/service/pg.ts b/src/service/pg.ts index ec4d03fe1..abb43a570 100644 --- 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 });