fix: embedding recall drop-dead halt (#1415)

labring · May 9, 2024 · d4169bf · d4169bf
1 parent afe5039
commit d4169bf
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 8 deletions.
diff --git a/packages/service/common/vectorStore/controller.d.ts b/packages/service/common/vectorStore/controller.d.ts
@@ -13,6 +13,7 @@ export type InsertVectorProps = {
 };
 
 export type EmbeddingRecallProps = {
+  teamId: string;
   datasetIds: string[];
   similarity?: number;
   efSearch?: number;

diff --git a/packages/service/common/vectorStore/pg/controller.ts b/packages/service/common/vectorStore/pg/controller.ts
@@ -129,15 +129,16 @@ export const embeddingRecall = async (
 ): Promise<{
   results: EmbeddingRecallItemType[];
 }> => {
-  const { datasetIds, vectors, limit, similarity = 0, retry = 2, efSearch = 100 } = props;
+  const { teamId, datasetIds, vectors, limit, similarity = 0, retry = 2, efSearch = 100 } = props;
 
   try {
     const results: any = await PgClient.query(
       `BEGIN;
         SET LOCAL hnsw.ef_search = ${efSearch};
         select id, collection_id, vector <#> '[${vectors[0]}]' AS score 
           from ${PgDatasetTableName} 
-          where dataset_id IN (${datasetIds.map((id) => `'${String(id)}'`).join(',')})
+          where team_id='${teamId}' 
+              AND dataset_id IN (${datasetIds.map((id) => `'${String(id)}'`).join(',')})
               AND vector <#> '[${vectors[0]}]' < -${similarity}
           order by score limit ${limit};
         COMMIT;`
@@ -153,10 +154,14 @@ export const embeddingRecall = async (
       }))
     };
   } catch (error) {
+    console.log(error);
     if (retry <= 0) {
       return Promise.reject(error);
     }
-    return embeddingRecall(props);
+    return embeddingRecall({
+      ...props,
+      retry: retry - 1
+    });
   }
 };
 

diff --git a/packages/service/core/dataset/search/controller.ts b/packages/service/core/dataset/search/controller.ts
@@ -59,19 +59,19 @@ export async function searchDatasetData(props: SearchDatasetDataProps) {
   const countRecallLimit = () => {
     if (searchMode === DatasetSearchModeEnum.embedding) {
       return {
-        embeddingLimit: 150,
+        embeddingLimit: 100,
         fullTextLimit: 0
       };
     }
     if (searchMode === DatasetSearchModeEnum.fullTextRecall) {
       return {
         embeddingLimit: 0,
-        fullTextLimit: 150
+        fullTextLimit: 100
       };
     }
     return {
-      embeddingLimit: 100,
-      fullTextLimit: 80
+      embeddingLimit: 80,
+      fullTextLimit: 60
     };
   };
   const embeddingRecall = async ({ query, limit }: { query: string; limit: number }) => {
@@ -82,9 +82,10 @@ export async function searchDatasetData(props: SearchDatasetDataProps) {
     });
 
     const { results } = await recallFromVectorStore({
+      teamId,
+      datasetIds,
       vectors,
       limit,
-      datasetIds,
       efSearch: global.systemEnv?.pgHNSWEfSearch
     });