topdog/supersonic
1
0
Fork 0
mirror of https://github.com/tencentmusic/supersonic.git synced 2026-04-18 04:14:21 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

[improvement][common]Add pgvector vector library adaptation. (#1800)

Browse Source
This commit is contained in:
Zhengyang Jia
2024-10-15 09:16:32 +08:00
committed by GitHub
parent 0b71390fde
commit 1ef642d0dd
10 changed files with 169 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
common/src/main/java/com/tencent/supersonic/common/config/EmbeddingStoreParameterConfig.java
View File

@@ -19,7 +19,7 @@ public class EmbeddingStoreParameterConfig extends ParameterConfig {
public static final Parameter EMBEDDING_STORE_PROVIDER = new Parameter(
"s2.embedding.store.provider", EmbeddingStoreType.IN_MEMORY.name(), "向量库类型",
"目前支持三种类型：IN_MEMORY、MILVUS、CHROMA", "list", MODULE_NAME, getCandidateValues());
"目前支持四种类型：IN_MEMORY、MILVUS、CHROMA、PGVECTOR", "list", MODULE_NAME, getCandidateValues());
public static final Parameter EMBEDDING_STORE_BASE_URL =
new Parameter("s2.embedding.store.base.url", "", "BaseUrl", "", "string", MODULE_NAME,
@@ -44,9 +44,18 @@ public class EmbeddingStoreParameterConfig extends ParameterConfig {
new Parameter("s2.embedding.store.databaseName", "", "DatabaseName", "", "string",
MODULE_NAME, null, getDatabaseNameDependency());
public static final Parameter EMBEDDING_STORE_POST =
new Parameter("s2.embedding.store.post", "", "端口", "", "number", MODULE_NAME, null,
getPostDependency());
public static final Parameter EMBEDDING_STORE_USER =
new Parameter("s2.embedding.store.user", "", "用户名", "", "string", MODULE_NAME, null,
getUserDependency());
@Override
public List<Parameter> getSysParameters() {
return Lists.newArrayList(EMBEDDING_STORE_PROVIDER, EMBEDDING_STORE_BASE_URL,
EMBEDDING_STORE_POST, EMBEDDING_STORE_USER,
EMBEDDING_STORE_API_KEY, EMBEDDING_STORE_DATABASE_NAME,
EMBEDDING_STORE_PERSIST_PATH, EMBEDDING_STORE_TIMEOUT, EMBEDDING_STORE_DIMENSION);
}
@@ -62,28 +71,38 @@ public class EmbeddingStoreParameterConfig extends ParameterConfig {
if (StringUtils.isNumeric(getParameterValue(EMBEDDING_STORE_DIMENSION))) {
dimension = Integer.valueOf(getParameterValue(EMBEDDING_STORE_DIMENSION));
}
Integer port = null;
if (StringUtils.isNumeric(getParameterValue(EMBEDDING_STORE_POST))) {
port = Integer.valueOf(getParameterValue(EMBEDDING_STORE_POST));
}
String user = getParameterValue(EMBEDDING_STORE_USER);
return EmbeddingStoreConfig.builder().provider(provider).baseUrl(baseUrl).apiKey(apiKey)
.persistPath(persistPath).databaseName(databaseName).timeOut(Long.valueOf(timeOut))
.dimension(dimension).build();
.dimension(dimension).post(port).user(user).build();
}
private static ArrayList<String> getCandidateValues() {
return Lists.newArrayList(EmbeddingStoreType.IN_MEMORY.name(),
EmbeddingStoreType.MILVUS.name(), EmbeddingStoreType.CHROMA.name());
EmbeddingStoreType.MILVUS.name(),
EmbeddingStoreType.CHROMA.name(),
EmbeddingStoreType.PGVECTOR.name());
}
private static List<Parameter.Dependency> getBaseUrlDependency() {
return getDependency(EMBEDDING_STORE_PROVIDER.getName(),
Lists.newArrayList(EmbeddingStoreType.MILVUS.name(),
EmbeddingStoreType.CHROMA.name()),
EmbeddingStoreType.CHROMA.name(),
EmbeddingStoreType.PGVECTOR.name()),
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "http://localhost:19530",
EmbeddingStoreType.CHROMA.name(), "http://localhost:8000"));
EmbeddingStoreType.CHROMA.name(), "http://localhost:8000",
EmbeddingStoreType.PGVECTOR.name(), "127.0.0.1"));
}
private static List<Parameter.Dependency> getApiKeyDependency() {
return getDependency(EMBEDDING_STORE_PROVIDER.getName(),
Lists.newArrayList(EmbeddingStoreType.MILVUS.name()),
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), DEMO));
Lists.newArrayList(EmbeddingStoreType.MILVUS.name(), EmbeddingStoreType.PGVECTOR.name()),
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), DEMO,
EmbeddingStoreType.PGVECTOR.name(), DEMO));
}
private static List<Parameter.Dependency> getPathDependency() {
@@ -94,13 +113,29 @@ public class EmbeddingStoreParameterConfig extends ParameterConfig {
private static List<Parameter.Dependency> getDimensionDependency() {
return getDependency(EMBEDDING_STORE_PROVIDER.getName(),
Lists.newArrayList(EmbeddingStoreType.MILVUS.name()),
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "384"));
Lists.newArrayList(EmbeddingStoreType.MILVUS.name(), EmbeddingStoreType.PGVECTOR.name()),
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "384",
EmbeddingStoreType.PGVECTOR.name(), "768"));
}
private static List<Parameter.Dependency> getDatabaseNameDependency() {
return getDependency(EMBEDDING_STORE_PROVIDER.getName(),
Lists.newArrayList(EmbeddingStoreType.MILVUS.name()),
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), ""));
Lists.newArrayList(EmbeddingStoreType.MILVUS.name(), EmbeddingStoreType.PGVECTOR.name()),
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "",
EmbeddingStoreType.PGVECTOR.name(), "postgres"));
}
private static List<Parameter.Dependency> getPostDependency() {
return getDependency(
EMBEDDING_STORE_PROVIDER.getName(),
Lists.newArrayList(EmbeddingStoreType.PGVECTOR.name()),
ImmutableMap.of(EmbeddingStoreType.PGVECTOR.name(), "54333"));
}
private static List<Parameter.Dependency> getUserDependency() {
return getDependency(
EMBEDDING_STORE_PROVIDER.getName(),
Lists.newArrayList(EmbeddingStoreType.PGVECTOR.name()),
ImmutableMap.of(EmbeddingStoreType.PGVECTOR.name(), "pgvector"));
}
}

2
common/src/main/java/com/tencent/supersonic/common/pojo/EmbeddingStoreConfig.java
View File

@@ -22,4 +22,6 @@ public class EmbeddingStoreConfig implements Serializable {
private Long timeOut = 60L;
private Integer dimension;
private String databaseName;
private Integer post;
private String user;
}