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

feat:add opensearch (#2049)
Some checks are pending
supersonic CentOS CI / build (21) (push) Waiting to run
Details
supersonic mac CI / build (21) (push) Waiting to run
Details
supersonic ubuntu CI / build (21) (push) Waiting to run
Details
supersonic windows CI / build (21) (push) Waiting to run
Details

Browse Source
This commit is contained in:
zyclove
2025-02-12 20:30:58 +08:00
committed by GitHub
parent 198c7c69e6
commit d942d35c93
9 changed files with 173 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
common/pom.xml
View File

@@ -174,6 +174,10 @@
<groupId>dev.langchain4j</groupId>
<artifactId>langchain4j-milvus</artifactId>
</dependency>
<dependency>
<groupId>dev.langchain4j</groupId>
<artifactId>langchain4j-opensearch</artifactId>
</dependency>
<dependency>
<groupId>dev.langchain4j</groupId>
<artifactId>langchain4j-pgvector</artifactId>
@@ -242,6 +246,10 @@
<groupId>com.google.code.gson</groupId>
<artifactId>gson</artifactId>
</dependency>
<dependency>
<groupId>com.amazonaws</groupId>
<artifactId>aws-java-sdk</artifactId>
</dependency>
<dependency>
<groupId>org.codehaus.woodstox</groupId>

35
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、PGVECTOR", "list", MODULE_NAME, getCandidateValues());
"目前支持四种类型：IN_MEMORY、MILVUS、CHROMA、PGVECTOR、OPENSEARCH", "list", MODULE_NAME, getCandidateValues());
public static final Parameter EMBEDDING_STORE_BASE_URL =
new Parameter("s2.embedding.store.base.url", "", "BaseUrl", "", "string", MODULE_NAME,
@@ -87,16 +87,19 @@ public class EmbeddingStoreParameterConfig extends ParameterConfig {
private static ArrayList<String> getCandidateValues() {
return Lists.newArrayList(EmbeddingStoreType.IN_MEMORY.name(),
EmbeddingStoreType.MILVUS.name(), EmbeddingStoreType.CHROMA.name(),
EmbeddingStoreType.PGVECTOR.name());
EmbeddingStoreType.PGVECTOR.name(), EmbeddingStoreType.OPENSEARCH.name());
}
private static List<Parameter.Dependency> getBaseUrlDependency() {
return getDependency(EMBEDDING_STORE_PROVIDER.getName(),
Lists.newArrayList(EmbeddingStoreType.MILVUS.name(),
EmbeddingStoreType.CHROMA.name(), EmbeddingStoreType.PGVECTOR.name()),
EmbeddingStoreType.CHROMA.name(),
EmbeddingStoreType.PGVECTOR.name(),
EmbeddingStoreType.OPENSEARCH.name()),
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "http://localhost:19530",
EmbeddingStoreType.CHROMA.name(), "http://localhost:8000",
EmbeddingStoreType.PGVECTOR.name(), "127.0.0.1"));
EmbeddingStoreType.PGVECTOR.name(), "127.0.0.1",
EmbeddingStoreType.OPENSEARCH.name(), "http://localhost:9200"));
}
private static List<Parameter.Dependency> getApiKeyDependency() {
@@ -114,17 +117,21 @@ public class EmbeddingStoreParameterConfig extends ParameterConfig {
private static List<Parameter.Dependency> getDimensionDependency() {
return getDependency(EMBEDDING_STORE_PROVIDER.getName(),
Lists.newArrayList(EmbeddingStoreType.MILVUS.name(),
EmbeddingStoreType.PGVECTOR.name()),
EmbeddingStoreType.PGVECTOR.name(),
EmbeddingStoreType.OPENSEARCH.name()),
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "384",
EmbeddingStoreType.PGVECTOR.name(), "512"));
EmbeddingStoreType.PGVECTOR.name(), "512",
EmbeddingStoreType.OPENSEARCH.name(), "512"));
}
private static List<Parameter.Dependency> getDatabaseNameDependency() {
return getDependency(EMBEDDING_STORE_PROVIDER.getName(),
Lists.newArrayList(EmbeddingStoreType.MILVUS.name(),
EmbeddingStoreType.PGVECTOR.name()),
EmbeddingStoreType.PGVECTOR.name(),
EmbeddingStoreType.OPENSEARCH.name()),
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "",
EmbeddingStoreType.PGVECTOR.name(), "postgres"));
EmbeddingStoreType.PGVECTOR.name(), "postgres",
EmbeddingStoreType.OPENSEARCH.name(), "ai_sql"));
}
private static List<Parameter.Dependency> getPortDependency() {
@@ -136,16 +143,20 @@ public class EmbeddingStoreParameterConfig extends ParameterConfig {
private static List<Parameter.Dependency> getUserDependency() {
return getDependency(EMBEDDING_STORE_PROVIDER.getName(),
Lists.newArrayList(EmbeddingStoreType.MILVUS.name(),
EmbeddingStoreType.PGVECTOR.name()),
EmbeddingStoreType.PGVECTOR.name(),
EmbeddingStoreType.OPENSEARCH.name()),
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "milvus",
EmbeddingStoreType.PGVECTOR.name(), "postgres"));
EmbeddingStoreType.PGVECTOR.name(), "postgres",
EmbeddingStoreType.OPENSEARCH.name(), "opensearch"));
}
private static List<Parameter.Dependency> getPasswordDependency() {
return getDependency(EMBEDDING_STORE_PROVIDER.getName(),
Lists.newArrayList(EmbeddingStoreType.MILVUS.name(),
EmbeddingStoreType.PGVECTOR.name()),
EmbeddingStoreType.PGVECTOR.name(),
EmbeddingStoreType.OPENSEARCH.name()),
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "milvus",
EmbeddingStoreType.PGVECTOR.name(), "postgres"));
EmbeddingStoreType.PGVECTOR.name(), "postgres",
EmbeddingStoreType.OPENSEARCH.name(), "opensearch"));
}
}

29
common/src/main/java/dev/langchain4j/opensearch/spring/EmbeddingStoreProperties.java Normal file
View File

@@ -0,0 +1,29 @@
package dev.langchain4j.opensearch.spring;
import io.milvus.common.clientenum.ConsistencyLevelEnum;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import lombok.Getter;
import lombok.Setter;
@Getter
@Setter
class EmbeddingStoreProperties {
private String uri;
private String host;
private Integer port;
private String serviceName;
private String region;
private String collectionName;
private Integer dimension;
private IndexType indexType;
private MetricType metricType;
private String token;
private String user;
private String password;
private ConsistencyLevelEnum consistencyLevel;
private Boolean retrieveEmbeddingsOnSearch;
private String databaseName;
private Boolean autoFlushOnInsert;
}

19
common/src/main/java/dev/langchain4j/opensearch/spring/OpenSearchAutoConfig.java Normal file
View File

@@ -0,0 +1,19 @@
package dev.langchain4j.opensearch.spring;
import dev.langchain4j.store.embedding.EmbeddingStoreFactory;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import static dev.langchain4j.opensearch.spring.Properties.PREFIX;
@Configuration
@EnableConfigurationProperties(dev.langchain4j.opensearch.spring.Properties.class)
public class OpenSearchAutoConfig {
@Bean
@ConditionalOnProperty(PREFIX + ".embedding-store.uri")
EmbeddingStoreFactory milvusChatModel(Properties properties) {
return new OpenSearchEmbeddingStoreFactory(properties.getEmbeddingStore());
}
}

59
common/src/main/java/dev/langchain4j/opensearch/spring/OpenSearchEmbeddingStoreFactory.java Normal file
View File

@@ -0,0 +1,59 @@
package dev.langchain4j.opensearch.spring;
import com.tencent.supersonic.common.pojo.EmbeddingStoreConfig;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.BaseEmbeddingStoreFactory;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.opensearch.OpenSearchEmbeddingStore;
import org.apache.hc.client5.http.auth.AuthScope;
import org.apache.hc.client5.http.auth.UsernamePasswordCredentials;
import org.apache.hc.client5.http.impl.auth.BasicCredentialsProvider;
import org.apache.hc.core5.http.HttpHost;
import org.opensearch.client.transport.aws.AwsSdk2TransportOptions;
import org.springframework.beans.BeanUtils;
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import java.net.URI;
/**
* @author zyc
*/
public class OpenSearchEmbeddingStoreFactory extends BaseEmbeddingStoreFactory {
private final EmbeddingStoreProperties storeProperties;
public OpenSearchEmbeddingStoreFactory(EmbeddingStoreConfig storeConfig) {
this(createPropertiesFromConfig(storeConfig));
}
public OpenSearchEmbeddingStoreFactory(EmbeddingStoreProperties storeProperties) {
this.storeProperties = storeProperties;
}
private static EmbeddingStoreProperties createPropertiesFromConfig(
EmbeddingStoreConfig storeConfig) {
EmbeddingStoreProperties embeddingStore = new EmbeddingStoreProperties();
BeanUtils.copyProperties(storeConfig, embeddingStore);
embeddingStore.setUri(storeConfig.getBaseUrl());
embeddingStore.setToken(storeConfig.getApiKey());
embeddingStore.setDatabaseName(storeConfig.getDatabaseName());
return embeddingStore;
}
@Override
public EmbeddingStore<TextSegment> createEmbeddingStore(String collectionName) {
final AwsSdk2TransportOptions options = AwsSdk2TransportOptions.builder()
.setCredentials(StaticCredentialsProvider.create(AwsBasicCredentials.create(storeProperties.getUser(), storeProperties.getPassword())))
.build();
final String indexName = storeProperties.getDatabaseName() + "_" + collectionName;
return OpenSearchEmbeddingStore.builder().serviceName(storeProperties.getServiceName())
.serverUrl(storeProperties.getUri())
.region(storeProperties.getRegion())
.indexName(indexName)
.userName(storeProperties.getUser())
.password(storeProperties.getPassword())
.apiKey(storeProperties.getToken())
.options(options)
.build();
}
}

17
common/src/main/java/dev/langchain4j/opensearch/spring/Properties.java Normal file
View File

@@ -0,0 +1,17 @@
package dev.langchain4j.opensearch.spring;
import lombok.Getter;
import lombok.Setter;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.boot.context.properties.NestedConfigurationProperty;
@Getter
@Setter
@ConfigurationProperties(prefix = Properties.PREFIX)
public class Properties {
static final String PREFIX = "langchain4j.opensearch";
@NestedConfigurationProperty
dev.langchain4j.opensearch.spring.EmbeddingStoreProperties embeddingStore;
}

6
common/src/main/java/dev/langchain4j/store/embedding/EmbeddingStoreFactoryProvider.java
View File

@@ -6,6 +6,7 @@ import com.tencent.supersonic.common.util.ContextUtils;
import dev.langchain4j.chroma.spring.ChromaEmbeddingStoreFactory;
import dev.langchain4j.inmemory.spring.InMemoryEmbeddingStoreFactory;
import dev.langchain4j.milvus.spring.MilvusEmbeddingStoreFactory;
import dev.langchain4j.opensearch.spring.OpenSearchEmbeddingStoreFactory;
import dev.langchain4j.pgvector.spring.PgvectorEmbeddingStoreFactory;
import org.apache.commons.lang3.StringUtils;
@@ -45,6 +46,11 @@ public class EmbeddingStoreFactoryProvider {
return factoryMap.computeIfAbsent(embeddingStoreConfig,
storeConfig -> new InMemoryEmbeddingStoreFactory(storeConfig));
}
if (EmbeddingStoreType.OPENSEARCH.name()
.equalsIgnoreCase(embeddingStoreConfig.getProvider())) {
return factoryMap.computeIfAbsent(embeddingStoreConfig,
storeConfig -> new OpenSearchEmbeddingStoreFactory(storeConfig));
}
throw new RuntimeException("Unsupported EmbeddingStoreFactory provider: "
+ embeddingStoreConfig.getProvider());
}

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

@@ -1,5 +1,5 @@
package dev.langchain4j.store.embedding;
public enum EmbeddingStoreType {
IN_MEMORY, MILVUS, CHROMA, PGVECTOR
IN_MEMORY, MILVUS, CHROMA, PGVECTOR, OPENSEARCH
}