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

Feature/model data embedding for chat and support status for metric and dimension (#311)

Browse Source 
* (improvement)(semantic) add offline status for metric and dimension

* (improvement)(chat) add metric recall

---------

Co-authored-by: jolunoluo
This commit is contained in:
LXW
2023-11-02 18:44:58 +08:00
committed by GitHub
parent f4e3922f47
commit ad20380283
89 changed files with 1572 additions and 896 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
common/src/main/java/com/tencent/supersonic/common/config/EmbeddingConfig.java Normal file
View File

@@ -0,0 +1,35 @@
package com.tencent.supersonic.common.config;
import lombok.Data;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Configuration;
@Configuration
@Data
public class EmbeddingConfig {
@Value("${embedding.url:}")
private String url;
@Value("${embedding.recognize.path:/preset_query_retrival}")
private String recognizePath;
@Value("${embedding.delete.path:/preset_delete_by_ids}")
private String deletePath;
@Value("${embedding.add.path:/preset_query_add}")
private String addPath;
@Value("${embedding.nResult:1}")
private String nResult;
@Value("${embedding.solvedQuery.recall.path:/solved_query_retrival}")
private String solvedQueryRecallPath;
@Value("${embedding.solvedQuery.add.path:/solved_query_add}")
private String solvedQueryAddPath;
@Value("${embedding.solved.query.nResult:5}")
private String solvedQueryResultNum;
}

9
common/src/main/java/com/tencent/supersonic/common/pojo/DataAddEvent.java
View File

@@ -1,5 +1,6 @@
package com.tencent.supersonic.common.pojo;
import com.tencent.supersonic.common.pojo.enums.TypeEnums;
import org.springframework.context.ApplicationEvent;
public class DataAddEvent extends ApplicationEvent {
@@ -7,9 +8,9 @@ public class DataAddEvent extends ApplicationEvent {
private String name;
private Long modelId;
private Long id;
private String type;
private TypeEnums type;
public DataAddEvent(Object source, String name, Long modelId, Long id, String type) {
public DataAddEvent(Object source, String name, Long modelId, Long id, TypeEnums type) {
super(source);
this.name = name;
this.modelId = modelId;
@@ -41,11 +42,11 @@ public class DataAddEvent extends ApplicationEvent {
this.modelId = modelId;
}
public void setType(String type) {
public void setType(TypeEnums type) {
this.type = type;
}
public String getType() {
public TypeEnums getType() {
return type;
}
}

9
common/src/main/java/com/tencent/supersonic/common/pojo/DataDeleteEvent.java
View File

@@ -1,5 +1,6 @@
package com.tencent.supersonic.common.pojo;
import com.tencent.supersonic.common.pojo.enums.TypeEnums;
import org.springframework.context.ApplicationEvent;
public class DataDeleteEvent extends ApplicationEvent {
@@ -7,9 +8,9 @@ public class DataDeleteEvent extends ApplicationEvent {
private String name;
private Long modelId;
private Long id;
private String type;
private TypeEnums type;
public DataDeleteEvent(Object source, String name, Long modelId, Long id, String type) {
public DataDeleteEvent(Object source, String name, Long modelId, Long id, TypeEnums type) {
super(source);
this.name = name;
this.modelId = modelId;
@@ -41,11 +42,11 @@ public class DataDeleteEvent extends ApplicationEvent {
this.modelId = modelId;
}
public void setType(String type) {
public void setType(TypeEnums type) {
this.type = type;
}
public String getType() {
public TypeEnums getType() {
return type;
}
}

28
common/src/main/java/com/tencent/supersonic/common/pojo/DataEvent.java Normal file
View File

@@ -0,0 +1,28 @@
package com.tencent.supersonic.common.pojo;
import com.tencent.supersonic.common.pojo.enums.EventType;
import org.springframework.context.ApplicationEvent;
import java.util.List;
public class DataEvent extends ApplicationEvent {
private List<DataItem> dataItems;
private EventType eventType;
public DataEvent(Object source, List<DataItem> dataItems, EventType eventType) {
super(source);
this.dataItems = dataItems;
this.eventType = eventType;
}
public List<DataItem> getDataItems() {
return dataItems;
}
public EventType getEventType() {
return eventType;
}
}

26
common/src/main/java/com/tencent/supersonic/common/pojo/DataItem.java Normal file
View File

@@ -0,0 +1,26 @@
package com.tencent.supersonic.common.pojo;
import com.tencent.supersonic.common.pojo.enums.TypeEnums;
import lombok.Builder;
import lombok.Data;
@Data
@Builder
public class DataItem {
private Long id;
private String bizName;
private String name;
private String newName;
private TypeEnums type;
private Long modelId;
public String getNewName() {
return newName == null ? name : newName;
}
}

9
common/src/main/java/com/tencent/supersonic/common/pojo/DataUpdateEvent.java
View File

@@ -1,5 +1,6 @@
package com.tencent.supersonic.common.pojo;
import com.tencent.supersonic.common.pojo.enums.TypeEnums;
import org.springframework.context.ApplicationEvent;
public class DataUpdateEvent extends ApplicationEvent {
@@ -8,9 +9,9 @@ public class DataUpdateEvent extends ApplicationEvent {
private String newName;
private Long modelId;
private Long id;
private String type;
private TypeEnums type;
public DataUpdateEvent(Object source, String name, String newName, Long modelId, Long id, String type) {
public DataUpdateEvent(Object source, String name, String newName, Long modelId, Long id, TypeEnums type) {
super(source);
this.name = name;
this.newName = newName;
@@ -51,11 +52,11 @@ public class DataUpdateEvent extends ApplicationEvent {
this.modelId = modelId;
}
public void setType(String type) {
public void setType(TypeEnums type) {
this.type = type;
}
public String getType() {
public TypeEnums getType() {
return type;
}
}

9
common/src/main/java/com/tencent/supersonic/common/pojo/enums/EventType.java Normal file
View File

@@ -0,0 +1,9 @@
package com.tencent.supersonic.common.pojo.enums;
public enum EventType {
ADD,
UPDATE,
DELETE
}

17
common/src/main/java/com/tencent/supersonic/common/util/embedding/EmbeddingCollection.java Normal file
View File

@@ -0,0 +1,17 @@
package com.tencent.supersonic.common.util.embedding;
import lombok.Data;
import java.util.Map;
@Data
public class EmbeddingCollection {
private String id;
private String name;
private Map<String, String> metaData;
}

20
common/src/main/java/com/tencent/supersonic/common/util/embedding/EmbeddingQuery.java Normal file
View File

@@ -0,0 +1,20 @@
package com.tencent.supersonic.common.util.embedding;
import lombok.Data;
import java.util.List;
import java.util.Map;
@Data
public class EmbeddingQuery {
private String queryId;
private String query;
private Map<String, String> metadata;
private List<Double> queryEmbedding;
}

107
common/src/main/java/com/tencent/supersonic/common/util/embedding/EmbeddingUtils.java Normal file
View File

@@ -0,0 +1,107 @@
package com.tencent.supersonic.common.util.embedding;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.serializer.SerializerFeature;
import com.google.common.collect.Lists;
import com.tencent.supersonic.common.config.EmbeddingConfig;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.util.UriComponentsBuilder;
import java.net.URI;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;
@Slf4j
@Component
public class EmbeddingUtils {
@Autowired
private EmbeddingConfig embeddingConfig;
private RestTemplate restTemplate = new RestTemplate();
public void addCollection(String collectionName) {
List<String> collections = getCollectionList();
if (collections.contains(collectionName)) {
return;
}
String url = String.format("%s/create_collection?collection_name=%s",
embeddingConfig.getUrl(), collectionName);
doRequest(url, null, HttpMethod.GET);
}
public void addQuery(String collectionName, List<EmbeddingQuery> queries) {
if (CollectionUtils.isEmpty(queries)) {
return;
}
String url = String.format("%s/add_query?collection_name=%s",
embeddingConfig.getUrl(), collectionName);
doRequest(url, JSONObject.toJSONString(queries, SerializerFeature.WriteMapNullValue), HttpMethod.POST);
}
public void deleteQuery(String collectionName, List<EmbeddingQuery> queries) {
if (CollectionUtils.isEmpty(queries)) {
return;
}
List<String> queryIds = queries.stream().map(EmbeddingQuery::getQueryId).collect(Collectors.toList());
String url = String.format("%s/delete_query_by_ids?collection_name=%s",
embeddingConfig.getUrl(), collectionName);
doRequest(url, JSONObject.toJSONString(queryIds), HttpMethod.POST);
}
public List<RetrieveQueryResult> retrieveQuery(String collectionName, RetrieveQuery retrieveQuery, int num) {
String url = String.format("%s/retrieve_query?collection_name=%s&n_results=%s",
embeddingConfig.getUrl(), collectionName, num);
ResponseEntity<String> responseEntity = doRequest(url, JSONObject.toJSONString(retrieveQuery,
SerializerFeature.WriteMapNullValue), HttpMethod.POST);
if (!responseEntity.hasBody()) {
return Lists.newArrayList();
}
return JSONObject.parseArray(responseEntity.getBody(), RetrieveQueryResult.class);
}
private List<String> getCollectionList() {
String url = embeddingConfig.getUrl() + "/list_collections";
ResponseEntity<String> responseEntity = doRequest(url, null, HttpMethod.GET);
if (!responseEntity.hasBody()) {
return Lists.newArrayList();
}
List<EmbeddingCollection> embeddingCollections = JSONObject.parseArray(responseEntity.getBody(),
EmbeddingCollection.class);
return embeddingCollections.stream().map(EmbeddingCollection::getName).collect(Collectors.toList());
}
private ResponseEntity doRequest(String url, String jsonBody, HttpMethod httpMethod) {
try {
HttpHeaders headers = new HttpHeaders();
headers.setContentType(MediaType.APPLICATION_JSON);
headers.setLocation(URI.create(url));
URI requestUrl = UriComponentsBuilder
.fromHttpUrl(url).build().encode().toUri();
HttpEntity<String> entity = new HttpEntity<>(headers);
if (jsonBody != null) {
log.info("[embedding] request body :{}", jsonBody);
entity = new HttpEntity<>(jsonBody, headers);
}
ResponseEntity<String> responseEntity = restTemplate.exchange(requestUrl,
httpMethod, entity, new ParameterizedTypeReference<String>() {});
log.info("[embedding] url :{} result body:{}", url, responseEntity);
return responseEntity;
} catch (Throwable e) {
log.warn("connect to embedding service failed, url:{}", url);
}
return ResponseEntity.of(Optional.empty());
}
}

18
common/src/main/java/com/tencent/supersonic/common/util/embedding/Retrieval.java Normal file
View File

@@ -0,0 +1,18 @@
package com.tencent.supersonic.common.util.embedding;
import lombok.Data;
import java.util.Map;
@Data
public class Retrieval {
private Long id;
private double distance;
private String query;
private Map<String, String> metadata;
}

20
common/src/main/java/com/tencent/supersonic/common/util/embedding/RetrieveQuery.java Normal file
View File

@@ -0,0 +1,20 @@
package com.tencent.supersonic.common.util.embedding;
import lombok.Builder;
import lombok.Data;
import java.util.List;
import java.util.Map;
@Data
@Builder
public class RetrieveQuery {
private List<String> queryTextsList;
private Map<String, String> filterCondition;
private List<List<Double>> queryEmbeddings;
}

15
common/src/main/java/com/tencent/supersonic/common/util/embedding/RetrieveQueryResult.java Normal file
View File

@@ -0,0 +1,15 @@
package com.tencent.supersonic.common.util.embedding;
import lombok.Data;
import java.util.List;
@Data
public class RetrieveQueryResult {
private String query;
private List<Retrieval> retrieval;
}