(fix)(chat) 修复做embedding检索时，出现重复数据的问题 (#1643)

tencentmusic · Sep 11, 2024 · 51e7a9b · 51e7a9b
1 parent 989ae3c
commit 51e7a9b
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/common/src/main/java/com/tencent/supersonic/common/service/impl/EmbeddingServiceImpl.java b/common/src/main/java/com/tencent/supersonic/common/service/impl/EmbeddingServiceImpl.java
@@ -59,6 +59,7 @@ public void addQuery(String collectionName, List<TextSegment> queries) {
                     continue;
                 }
                 embeddingStore.add(embedding, query);
+                cache.put(TextSegmentConvert.getQueryId(query), true);
             } catch (Exception e) {
                 log.error(
                         "embeddingModel embed error question: {}, embeddingStore: {}",
@@ -117,6 +118,9 @@ public void deleteQuery(String collectionName, List<TextSegment> queries) {
                             new MetadataFilterBuilder(TextSegmentConvert.QUERY_ID);
                     Filter filter = filterBuilder.isIn(queryIds);
                     inMemoryEmbeddingStore.removeAll(filter);
+                    for (String queryId : queryIds) {
+                        cache.put(queryId, false);
+                    }
                 }
             } else {
                 throw new RuntimeException("Not supported yet.");