(improvement)(chat) Support connecting to Milvus service in the cloud. (#1588)

This commit is contained in:
lexluo09
2024-08-21 16:43:51 +08:00
committed by GitHub
parent 7e013ca36a
commit afa82bf98d
3 changed files with 18 additions and 3 deletions

View File

@@ -45,12 +45,17 @@ public class EmbeddingStoreParameterConfig extends ParameterConfig {
new Parameter("s2.embedding.store.dimension", "", new Parameter("s2.embedding.store.dimension", "",
"纬度", "", "number", "纬度", "", "number",
"向量库配置", null, getDimensionDependency()); "向量库配置", null, getDimensionDependency());
public static final Parameter EMBEDDING_STORE_DATABASE_NAME =
new Parameter("s2.embedding.store.databaseName", "",
"DatabaseName", "", "string",
"向量库配置", null, getDatabaseNameDependency());
@Override @Override
public List<Parameter> getSysParameters() { public List<Parameter> getSysParameters() {
return Lists.newArrayList( return Lists.newArrayList(
EMBEDDING_STORE_PROVIDER, EMBEDDING_STORE_BASE_URL, EMBEDDING_STORE_API_KEY, EMBEDDING_STORE_PROVIDER, EMBEDDING_STORE_BASE_URL, EMBEDDING_STORE_API_KEY,
EMBEDDING_STORE_PERSIST_PATH, EMBEDDING_STORE_TIMEOUT, EMBEDDING_STORE_DIMENSION EMBEDDING_STORE_DATABASE_NAME, EMBEDDING_STORE_PERSIST_PATH,
EMBEDDING_STORE_TIMEOUT, EMBEDDING_STORE_DIMENSION
); );
} }
@@ -60,12 +65,13 @@ public class EmbeddingStoreParameterConfig extends ParameterConfig {
String apiKey = getParameterValue(EMBEDDING_STORE_API_KEY); String apiKey = getParameterValue(EMBEDDING_STORE_API_KEY);
String persistPath = getParameterValue(EMBEDDING_STORE_PERSIST_PATH); String persistPath = getParameterValue(EMBEDDING_STORE_PERSIST_PATH);
String timeOut = getParameterValue(EMBEDDING_STORE_TIMEOUT); String timeOut = getParameterValue(EMBEDDING_STORE_TIMEOUT);
String databaseName = getParameterValue(EMBEDDING_STORE_DATABASE_NAME);
Integer dimension = null; Integer dimension = null;
if (StringUtils.isNumeric(getParameterValue(EMBEDDING_STORE_DIMENSION))) { if (StringUtils.isNumeric(getParameterValue(EMBEDDING_STORE_DIMENSION))) {
dimension = Integer.valueOf(getParameterValue(EMBEDDING_STORE_DIMENSION)); dimension = Integer.valueOf(getParameterValue(EMBEDDING_STORE_DIMENSION));
} }
return EmbeddingStoreConfig.builder().provider(provider) return EmbeddingStoreConfig.builder().provider(provider).baseUrl(baseUrl)
.baseUrl(baseUrl).apiKey(apiKey).persistPath(persistPath) .apiKey(apiKey).persistPath(persistPath).databaseName(databaseName)
.timeOut(Long.valueOf(timeOut)).dimension(dimension).build(); .timeOut(Long.valueOf(timeOut)).dimension(dimension).build();
} }
@@ -106,4 +112,11 @@ public class EmbeddingStoreParameterConfig extends ParameterConfig {
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "384") ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "384")
); );
} }
private static List<Parameter.Dependency> getDatabaseNameDependency() {
return getDependency(EMBEDDING_STORE_PROVIDER.getName(),
Lists.newArrayList(EmbeddingStoreType.MILVUS.name()),
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "")
);
}
} }

View File

@@ -21,4 +21,5 @@ public class EmbeddingStoreConfig implements Serializable {
private String apiKey; private String apiKey;
private Long timeOut = 60L; private Long timeOut = 60L;
private Integer dimension; private Integer dimension;
private String databaseName;
} }

View File

@@ -22,6 +22,7 @@ public class MilvusEmbeddingStoreFactory extends BaseEmbeddingStoreFactory {
EmbeddingStoreProperties embeddingStore = new EmbeddingStoreProperties(); EmbeddingStoreProperties embeddingStore = new EmbeddingStoreProperties();
BeanUtils.copyProperties(storeConfig, embeddingStore); BeanUtils.copyProperties(storeConfig, embeddingStore);
embeddingStore.setUri(storeConfig.getBaseUrl()); embeddingStore.setUri(storeConfig.getBaseUrl());
embeddingStore.setToken(storeConfig.getApiKey());
return embeddingStore; return embeddingStore;
} }