diff --git a/packages/service/common/vectorStore/controller.d.ts b/packages/service/common/vectorStore/controller.d.ts index 442b9012b..311544de0 100644 --- 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 index 828a3b227..f6d115722 100644 --- a/packages/service/common/vectorStore/pg/controller.ts +++ b/packages/service/common/vectorStore/pg/controller.ts @@ -129,7 +129,7 @@ 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( @@ -137,7 +137,8 @@ export const embeddingRecall = async ( 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 index 87b89c55f..ecefb7aff 100644 --- 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 });