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)(chat) Vector retrieval supports filtering by modelId collection during query. (#1576)

Browse Source
This commit is contained in:
lexluo09
2024-08-16 21:31:07 +08:00
committed by GitHub
parent 6aff51d394
commit 7150f19def
4 changed files with 73 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
chat/server/src/main/java/com/tencent/supersonic/chat/server/processor/execute/MetricRecommendProcessor.java
View File

@@ -1,20 +1,19 @@
package com.tencent.supersonic.chat.server.processor.execute;
import com.alibaba.fastjson.JSONObject;
import com.tencent.supersonic.chat.api.pojo.response.QueryResult;
import com.tencent.supersonic.chat.server.pojo.ExecuteContext;
import com.tencent.supersonic.common.pojo.Constants;
import com.tencent.supersonic.common.pojo.enums.DictWordType;
import com.tencent.supersonic.common.pojo.enums.QueryType;
import com.tencent.supersonic.common.util.ContextUtils;
import dev.langchain4j.store.embedding.Retrieval;
import dev.langchain4j.store.embedding.RetrieveQuery;
import dev.langchain4j.store.embedding.RetrieveQueryResult;
import com.tencent.supersonic.headless.api.pojo.SchemaElement;
import com.tencent.supersonic.headless.api.pojo.SchemaElementType;
import com.tencent.supersonic.headless.api.pojo.SemanticParseInfo;
import com.tencent.supersonic.chat.api.pojo.response.QueryResult;
import com.tencent.supersonic.headless.chat.knowledge.MetaEmbeddingService;
import java.util.Objects;
import dev.langchain4j.store.embedding.Retrieval;
import dev.langchain4j.store.embedding.RetrieveQuery;
import dev.langchain4j.store.embedding.RetrieveQueryResult;
import org.springframework.util.CollectionUtils;
import java.util.Collections;
@@ -23,6 +22,7 @@ import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;
@@ -45,7 +45,7 @@ public class MetricRecommendProcessor implements ExecuteResultProcessor {
return;
}
List<String> metricNames = Collections.singletonList(parseInfo.getMetrics().iterator().next().getName());
Map<String, String> filterCondition = new HashMap<>();
Map<String, Object> filterCondition = new HashMap<>();
filterCondition.put("modelId", parseInfo.getMetrics().iterator().next().getDataSetId().toString());
filterCondition.put("type", SchemaElementType.METRIC.name());
RetrieveQuery retrieveQuery = RetrieveQuery.builder().queryTextsList(metricNames)

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

@@ -75,7 +75,7 @@ public class EmbeddingServiceImpl implements EmbeddingService {
if (cachedResult != null) {
return cachedResult;
}
Map<String, String> filterCondition = new HashMap<>();
Map<String, Object> filterCondition = new HashMap<>();
filterCondition.put(TextSegmentConvert.QUERY_ID, queryId);
Filter filter = createCombinedFilter(filterCondition);
EmbeddingSearchRequest request = EmbeddingSearchRequest.builder()
@@ -115,7 +115,7 @@ public class EmbeddingServiceImpl implements EmbeddingService {
public List<RetrieveQueryResult> retrieveQuery(String collectionName, RetrieveQuery retrieveQuery, int num) {
EmbeddingStore embeddingStore = EmbeddingStoreFactoryProvider.getFactory().create(collectionName);
EmbeddingModel embeddingModel = ModelProvider.getEmbeddingModel();
Map<String, String> filterCondition = retrieveQuery.getFilterCondition();
Map<String, Object> filterCondition = retrieveQuery.getFilterCondition();
return retrieveQuery.getQueryTextsList().stream()
.map(queryText -> retrieveSingleQuery(queryText, embeddingModel, embeddingStore, filterCondition, num))
.collect(Collectors.toList());
@@ -124,7 +124,7 @@ public class EmbeddingServiceImpl implements EmbeddingService {
private RetrieveQueryResult retrieveSingleQuery(String queryText,
EmbeddingModel embeddingModel,
EmbeddingStore embeddingStore,
Map<String, String> filterCondition,
Map<String, Object> filterCondition,
int num) {
Embedding embeddedText = embeddingModel.embed(queryText).content();
Filter filter = createCombinedFilter(filterCondition);
@@ -159,14 +159,27 @@ public class EmbeddingServiceImpl implements EmbeddingService {
return retrieval;
}
private static Filter createCombinedFilter(Map<String, String> map) {
public static Filter createCombinedFilter(Map<String, Object> map) {
if (MapUtils.isEmpty(map)) {
return null;
}
Filter result = null;
for (Map.Entry<String, String> entry : map.entrySet()) {
IsEqualTo isEqualTo = new IsEqualTo(entry.getKey(), entry.getValue());
result = (result == null) ? isEqualTo : Filter.and(result, isEqualTo);
for (Map.Entry<String, Object> entry : map.entrySet()) {
String key = entry.getKey();
Object value = entry.getValue();
Filter orFilter = null;
if (value instanceof List) {
for (String val : (List<String>) value) {
IsEqualTo isEqualTo = new IsEqualTo(key, val);
orFilter = (orFilter == null) ? isEqualTo : Filter.or(orFilter, isEqualTo);
}
} else if (value instanceof String) {
orFilter = new IsEqualTo(key, value);
}
if (orFilter != null) {
result = (result == null) ? orFilter : Filter.and(result, orFilter);
}
}
return result;
}

2
common/src/main/java/dev/langchain4j/store/embedding/RetrieveQuery.java
View File

@@ -12,7 +12,7 @@ public class RetrieveQuery {
private List<String> queryTextsList;
private Map<String, String> filterCondition;
private Map<String, Object> filterCondition;
private List<List<Double>> queryEmbeddings;

98
headless/chat/src/main/java/com/tencent/supersonic/headless/chat/knowledge/MetaEmbeddingService.java
View File

@@ -8,20 +8,20 @@ import com.tencent.supersonic.headless.chat.knowledge.helper.NatureHelper;
import dev.langchain4j.store.embedding.Retrieval;
import dev.langchain4j.store.embedding.RetrieveQuery;
import dev.langchain4j.store.embedding.RetrieveQueryResult;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;
@Service
@Slf4j
public class MetaEmbeddingService {
@@ -32,14 +32,16 @@ public class MetaEmbeddingService {
private EmbeddingConfig embeddingConfig;
public List<RetrieveQueryResult> retrieveQuery(RetrieveQuery retrieveQuery, int num,
Map<Long, List<Long>> modelIdToDataSetIds, Set<Long> detectDataSetIds) {
Map<Long, List<Long>> modelIdToDataSetIds,
Set<Long> detectDataSetIds) {
// dataSetIds->modelIds
Set<Long> allModels = NatureHelper.getModelIds(modelIdToDataSetIds, detectDataSetIds);
if (CollectionUtils.isNotEmpty(allModels) && allModels.size() == 1) {
Map<String, String> filterCondition = new HashMap<>();
String modelId = allModels.stream().findFirst().get().toString();
filterCondition.put("modelId", modelId + DictWordType.NATURE_SPILT);
if (CollectionUtils.isNotEmpty(allModels)) {
Map<String, Object> filterCondition = new HashMap<>();
filterCondition.put("modelId", allModels.stream()
.map(modelId -> modelId + DictWordType.NATURE_SPILT)
.collect(Collectors.toList()));
retrieveQuery.setFilterCondition(filterCondition);
}
@@ -48,46 +50,38 @@ public class MetaEmbeddingService {
if (CollectionUtils.isEmpty(resultList)) {
return new ArrayList<>();
}
//filter by modelId
if (CollectionUtils.isEmpty(allModels)) {
return resultList;
}
// Filter and process query results.
return resultList.stream()
.map(retrieveQueryResult -> {
List<Retrieval> retrievals = retrieveQueryResult.getRetrieval();
if (CollectionUtils.isEmpty(retrievals)) {
return retrieveQueryResult;
}
//filter by modelId
retrievals.removeIf(retrieval -> {
Long modelId = Retrieval.getLongId(retrieval.getMetadata().get("modelId"));
if (Objects.isNull(modelId)) {
return CollectionUtils.isEmpty(allModels);
}
return !allModels.contains(modelId);
});
//add dataSetId
retrievals = retrievals.stream().flatMap(retrieval -> {
Long modelId = Retrieval.getLongId(retrieval.getMetadata().get("modelId"));
List<Long> dataSetIdsByModelId = modelIdToDataSetIds.get(modelId);
if (!CollectionUtils.isEmpty(dataSetIdsByModelId)) {
Set<Retrieval> result = new HashSet<>();
for (Long dataSetId : dataSetIdsByModelId) {
Retrieval retrievalNew = new Retrieval();
BeanUtils.copyProperties(retrieval, retrievalNew);
retrievalNew.getMetadata().putIfAbsent("dataSetId", dataSetId + Constants.UNDERLINE);
result.add(retrievalNew);
}
return result.stream();
}
Set<Retrieval> result = new HashSet<>();
result.add(retrieval);
return result.stream();
}).collect(Collectors.toList());
retrieveQueryResult.setRetrieval(retrievals);
return retrieveQueryResult;
})
.filter(retrieveQueryResult -> CollectionUtils.isNotEmpty(retrieveQueryResult.getRetrieval()))
.map(result -> getRetrieveQueryResult(modelIdToDataSetIds, result))
.filter(result -> CollectionUtils.isNotEmpty(result.getRetrieval()))
.collect(Collectors.toList());
}
private static RetrieveQueryResult getRetrieveQueryResult(Map<Long,
List<Long>> modelIdToDataSetIds, RetrieveQueryResult result) {
List<Retrieval> retrievals = result.getRetrieval();
if (CollectionUtils.isEmpty(retrievals)) {
return result;
}
// Process each Retrieval object.
List<Retrieval> updatedRetrievals = retrievals.stream()
.flatMap(retrieval -> {
Long modelId = Retrieval.getLongId(retrieval.getMetadata().get("modelId"));
List<Long> dataSetIds = modelIdToDataSetIds.get(modelId);
if (CollectionUtils.isEmpty(dataSetIds)) {
return Stream.of(retrieval);
}
return dataSetIds.stream().map(dataSetId -> {
Retrieval newRetrieval = new Retrieval();
BeanUtils.copyProperties(retrieval, newRetrieval);
newRetrieval.getMetadata().putIfAbsent("dataSetId", dataSetId + Constants.UNDERLINE);
return newRetrieval;
});
})
.collect(Collectors.toList());
result.setRetrieval(updatedRetrievals);
return result;
}
}