topdog/supersonic
1
0
Fork 0
mirror of https://github.com/tencentmusic/supersonic.git synced 2025-12-10 11:07:06 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

(improvement)(common) Support deleting embedding data in InMemoryEmbeddingStore. (#1360)

Browse Source
This commit is contained in:
lexluo09
2024-07-05 23:29:22 +08:00
committed by GitHub
parent a1ab7ac1c1
commit d39db734c4
1 changed files with 17 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
common/src/main/java/com/tencent/supersonic/common/service/impl/EmbeddingServiceImpl.java
View File

@@ -16,7 +16,9 @@ import dev.langchain4j.store.embedding.RetrieveQuery;
import dev.langchain4j.store.embedding.RetrieveQueryResult;
import dev.langchain4j.store.embedding.TextSegmentConvert;
import dev.langchain4j.store.embedding.filter.Filter;
import dev.langchain4j.store.embedding.filter.MetadataFilterBuilder;
import dev.langchain4j.store.embedding.filter.comparison.IsEqualTo;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
@@ -95,6 +97,21 @@ public class EmbeddingServiceImpl implements EmbeddingService {
@Override
public void deleteQuery(String collectionName, List<TextSegment> queries) {
//Not supported yet in Milvus and Chroma
EmbeddingStore embeddingStore = embeddingStoreFactory.create(collectionName);
try {
if (embeddingStore instanceof InMemoryEmbeddingStore) {
InMemoryEmbeddingStore inMemoryEmbeddingStore = (InMemoryEmbeddingStore) embeddingStore;
List<String> queryIds = queries.stream().map(textSegment -> TextSegmentConvert.getQueryId(textSegment))
.filter(Objects::nonNull).collect(Collectors.toList());
if (CollectionUtils.isNotEmpty(queryIds)) {
MetadataFilterBuilder filterBuilder = new MetadataFilterBuilder(TextSegmentConvert.QUERY_ID);
Filter filter = filterBuilder.isIn(queryIds);
inMemoryEmbeddingStore.removeAll(filter);
}
}
} catch (Exception e) {
log.error("deleteQuery error,collectionName:{},queries:{}", collectionName, queries);
}
}
@Override