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

This commit is contained in:
Archer
2024-05-09 16:13:06 +08:00
committed by GitHub
parent afe5039cd3
commit d4169bf066
3 changed files with 15 additions and 8 deletions

View File

@@ -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
});
}
};