diff --git a/common/src/main/java/com/tencent/supersonic/common/config/EmbeddingStoreParameterConfig.java b/common/src/main/java/com/tencent/supersonic/common/config/EmbeddingStoreParameterConfig.java index 85f663265..2eaf942d7 100644 --- a/common/src/main/java/com/tencent/supersonic/common/config/EmbeddingStoreParameterConfig.java +++ b/common/src/main/java/com/tencent/supersonic/common/config/EmbeddingStoreParameterConfig.java @@ -6,6 +6,7 @@ import com.tencent.supersonic.common.pojo.EmbeddingStoreConfig; import com.tencent.supersonic.common.pojo.Parameter; import dev.langchain4j.store.embedding.EmbeddingStoreType; import lombok.extern.slf4j.Slf4j; +import org.apache.commons.lang3.StringUtils; import org.springframework.stereotype.Service; import java.util.List; @@ -26,50 +27,63 @@ public class EmbeddingStoreParameterConfig extends ParameterConfig { "BaseUrl", "", "string", "向量库配置", null, getDependency(EMBEDDING_STORE_PROVIDER.getName(), - Lists.newArrayList( - EmbeddingStoreType.MILVUS.name(), - EmbeddingStoreType.CHROMA.name() - ), - ImmutableMap.of( - EmbeddingStoreType.MILVUS.name(), "http://localhost:19530", - EmbeddingStoreType.CHROMA.name(), "http://localhost:8000" - ) - )); + Lists.newArrayList( + EmbeddingStoreType.MILVUS.name(), + EmbeddingStoreType.CHROMA.name() + ), + ImmutableMap.of( + EmbeddingStoreType.MILVUS.name(), "http://localhost:19530", + EmbeddingStoreType.CHROMA.name(), "http://localhost:8000" + ) + )); public static final Parameter EMBEDDING_STORE_API_KEY = new Parameter("s2.embedding.store.api.key", "", "ApiKey", "", "string", "向量库配置", null, getDependency(EMBEDDING_STORE_PROVIDER.getName(), - Lists.newArrayList( - EmbeddingStoreType.MILVUS.name() - ), - ImmutableMap.of( - EmbeddingStoreType.MILVUS.name(), "demo" - ) - )); + Lists.newArrayList( + EmbeddingStoreType.MILVUS.name() + ), + ImmutableMap.of( + EmbeddingStoreType.MILVUS.name(), "demo" + ) + )); public static final Parameter EMBEDDING_STORE_PERSIST_PATH = new Parameter("s2.embedding.store.persist.path", "/tmp", "持久化路径", "", "string", "向量库配置", null, getDependency(EMBEDDING_STORE_PROVIDER.getName(), - Lists.newArrayList( - EmbeddingStoreType.IN_MEMORY.name() - ), - ImmutableMap.of( - EmbeddingStoreType.IN_MEMORY.name(), "/tmp" - ))); + Lists.newArrayList( + EmbeddingStoreType.IN_MEMORY.name() + ), + ImmutableMap.of( + EmbeddingStoreType.IN_MEMORY.name(), "/tmp" + ))); public static final Parameter EMBEDDING_STORE_TIMEOUT = new Parameter("s2.embedding.store.timeout", "60", "超时时间(秒)", "", "number", "向量库配置"); + public static final Parameter EMBEDDING_STORE_DIMENSION = + new Parameter("s2.embedding.store.dimension", "", + "纬度", "", + "number", "向量库配置", null, + getDependency(EMBEDDING_STORE_PROVIDER.getName(), + Lists.newArrayList( + EmbeddingStoreType.MILVUS.name() + ), + ImmutableMap.of( + EmbeddingStoreType.MILVUS.name(), "384" + ) + )); + @Override public List getSysParameters() { return Lists.newArrayList( EMBEDDING_STORE_PROVIDER, EMBEDDING_STORE_BASE_URL, EMBEDDING_STORE_API_KEY, - EMBEDDING_STORE_PERSIST_PATH, EMBEDDING_STORE_TIMEOUT + EMBEDDING_STORE_PERSIST_PATH, EMBEDDING_STORE_TIMEOUT, EMBEDDING_STORE_DIMENSION ); } @@ -79,13 +93,17 @@ public class EmbeddingStoreParameterConfig extends ParameterConfig { String apiKey = getParameterValue(EMBEDDING_STORE_API_KEY); String persistPath = getParameterValue(EMBEDDING_STORE_PERSIST_PATH); String timeOut = getParameterValue(EMBEDDING_STORE_TIMEOUT); - + Integer dimension = null; + if (StringUtils.isNumeric(getParameterValue(EMBEDDING_STORE_DIMENSION))) { + dimension = Integer.valueOf(getParameterValue(EMBEDDING_STORE_DIMENSION)); + } return EmbeddingStoreConfig.builder() .provider(provider) .baseUrl(baseUrl) .apiKey(apiKey) .persistPath(persistPath) .timeOut(Long.valueOf(timeOut)) + .dimension(dimension) .build(); } } \ No newline at end of file diff --git a/common/src/main/java/com/tencent/supersonic/common/pojo/EmbeddingStoreConfig.java b/common/src/main/java/com/tencent/supersonic/common/pojo/EmbeddingStoreConfig.java index c92af21e6..cd9ad5275 100644 --- a/common/src/main/java/com/tencent/supersonic/common/pojo/EmbeddingStoreConfig.java +++ b/common/src/main/java/com/tencent/supersonic/common/pojo/EmbeddingStoreConfig.java @@ -20,4 +20,5 @@ public class EmbeddingStoreConfig implements Serializable { private String baseUrl; private String apiKey; private Long timeOut = 60L; + private Integer dimension; } \ No newline at end of file