(improvement)(chat) Upgrade and optimize the embedding metastore. (#1198)

This commit is contained in:
lexluo09
2024-06-23 21:46:10 +08:00
committed by GitHub
parent 2ae94fb38c
commit 4d6cbf31f7
46 changed files with 3788 additions and 498 deletions

View File

@@ -4,9 +4,9 @@ import com.tencent.supersonic.common.config.EmbeddingConfig;
import com.tencent.supersonic.common.pojo.DataEvent;
import com.tencent.supersonic.common.pojo.DataItem;
import com.tencent.supersonic.common.pojo.enums.EventType;
import dev.langchain4j.store.embedding.ComponentFactory;
import com.tencent.supersonic.common.service.EmbeddingService;
import dev.langchain4j.store.embedding.EmbeddingQuery;
import dev.langchain4j.store.embedding.S2EmbeddingStore;
import java.util.List;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
@@ -15,8 +15,6 @@ import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;
import java.util.List;
@Component
@Slf4j
public class MetaEmbeddingListener implements ApplicationListener<DataEvent> {
@@ -24,7 +22,8 @@ public class MetaEmbeddingListener implements ApplicationListener<DataEvent> {
@Autowired
private EmbeddingConfig embeddingConfig;
private S2EmbeddingStore s2EmbeddingStore = ComponentFactory.getS2EmbeddingStore();
@Autowired
private EmbeddingService embeddingService;
@Value("${embedding.operation.sleep.time:3000}")
private Integer embeddingOperationSleepTime;
@@ -41,14 +40,14 @@ public class MetaEmbeddingListener implements ApplicationListener<DataEvent> {
return;
}
sleep();
s2EmbeddingStore.addCollection(embeddingConfig.getMetaCollectionName());
embeddingService.addCollection(embeddingConfig.getMetaCollectionName());
if (event.getEventType().equals(EventType.ADD)) {
s2EmbeddingStore.addQuery(embeddingConfig.getMetaCollectionName(), embeddingQueries);
embeddingService.addQuery(embeddingConfig.getMetaCollectionName(), embeddingQueries);
} else if (event.getEventType().equals(EventType.DELETE)) {
s2EmbeddingStore.deleteQuery(embeddingConfig.getMetaCollectionName(), embeddingQueries);
embeddingService.deleteQuery(embeddingConfig.getMetaCollectionName(), embeddingQueries);
} else if (event.getEventType().equals(EventType.UPDATE)) {
s2EmbeddingStore.deleteQuery(embeddingConfig.getMetaCollectionName(), embeddingQueries);
s2EmbeddingStore.addQuery(embeddingConfig.getMetaCollectionName(), embeddingQueries);
embeddingService.deleteQuery(embeddingConfig.getMetaCollectionName(), embeddingQueries);
embeddingService.addQuery(embeddingConfig.getMetaCollectionName(), embeddingQueries);
}
}

View File

@@ -2,25 +2,23 @@ package com.tencent.supersonic.headless.server.schedule;
import com.tencent.supersonic.common.config.EmbeddingConfig;
import com.tencent.supersonic.common.pojo.DataItem;
import dev.langchain4j.store.embedding.ComponentFactory;
import dev.langchain4j.store.embedding.EmbeddingQuery;
import dev.langchain4j.store.embedding.InMemoryS2EmbeddingStore;
import dev.langchain4j.store.embedding.S2EmbeddingStore;
import com.tencent.supersonic.common.service.EmbeddingService;
import com.tencent.supersonic.headless.server.service.DimensionService;
import com.tencent.supersonic.headless.server.service.MetricService;
import dev.langchain4j.store.embedding.EmbeddingQuery;
import java.util.List;
import javax.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;
import javax.annotation.PreDestroy;
import java.util.List;
@Component
@Slf4j
public class EmbeddingTask {
private S2EmbeddingStore s2EmbeddingStore = ComponentFactory.getS2EmbeddingStore();
@Autowired
private EmbeddingService embeddingService;
@Autowired
private EmbeddingConfig embeddingConfig;
@Autowired
@@ -31,23 +29,22 @@ public class EmbeddingTask {
@PreDestroy
public void onShutdown() {
embeddingStorePersistentToFile();
// embeddingStorePersistentToFile();
}
private void embeddingStorePersistentToFile() {
if (s2EmbeddingStore instanceof InMemoryS2EmbeddingStore) {
log.info("start persistentToFile");
((InMemoryS2EmbeddingStore) s2EmbeddingStore).persistentToFile();
log.info("end persistentToFile");
}
}
// private void embeddingStorePersistentToFile() {
// if (embeddingService instanceof InMemoryEmbeddingService) {
// log.info("start persistentToFile");
// ((InMemoryEmbeddingService) embeddingService).persistentToFile();
// log.info("end persistentToFile");
// }
// }
@Scheduled(cron = "${inMemoryEmbeddingStore.persistent.cron:0 0 * * * ?}")
public void executeTask() {
embeddingStorePersistentToFile();
// embeddingStorePersistentToFile();
}
/***
* reload meta embedding
*/
@@ -57,11 +54,11 @@ public class EmbeddingTask {
try {
List<DataItem> metricDataItems = metricService.getDataEvent().getDataItems();
s2EmbeddingStore.addQuery(embeddingConfig.getMetaCollectionName(),
embeddingService.addQuery(embeddingConfig.getMetaCollectionName(),
EmbeddingQuery.convertToEmbedding(metricDataItems));
List<DataItem> dimensionDataItems = dimensionService.getDataEvent().getDataItems();
s2EmbeddingStore.addQuery(embeddingConfig.getMetaCollectionName(),
embeddingService.addQuery(embeddingConfig.getMetaCollectionName(),
EmbeddingQuery.convertToEmbedding(dimensionDataItems));
} catch (Exception e) {
log.error("reload.meta.embedding error", e);