From e47ca9afa457e4fd8cc6a84fdf6235427402efe5 Mon Sep 17 00:00:00 2001 From: zhaoyingchao Date: Tue, 11 Feb 2025 20:59:38 +0800 Subject: [PATCH] feat:add opensearch --- common/pom.xml | 8 +++ .../config/EmbeddingStoreParameterConfig.java | 35 +++++++---- .../spring/EmbeddingStoreProperties.java | 29 +++++++++ .../spring/OpenSearchAutoConfig.java | 19 ++++++ .../OpenSearchEmbeddingStoreFactory.java | 59 +++++++++++++++++++ .../opensearch/spring/Properties.java | 17 ++++++ .../EmbeddingStoreFactoryProvider.java | 6 ++ .../store/embedding/EmbeddingStoreType.java | 2 +- pom.xml | 11 ++++ 9 files changed, 173 insertions(+), 13 deletions(-) create mode 100644 common/src/main/java/dev/langchain4j/opensearch/spring/EmbeddingStoreProperties.java create mode 100644 common/src/main/java/dev/langchain4j/opensearch/spring/OpenSearchAutoConfig.java create mode 100644 common/src/main/java/dev/langchain4j/opensearch/spring/OpenSearchEmbeddingStoreFactory.java create mode 100644 common/src/main/java/dev/langchain4j/opensearch/spring/Properties.java diff --git a/common/pom.xml b/common/pom.xml index c26e05dead..1651892020 100644 --- a/common/pom.xml +++ b/common/pom.xml @@ -174,6 +174,10 @@ dev.langchain4j langchain4j-milvus + + dev.langchain4j + langchain4j-opensearch + dev.langchain4j langchain4j-pgvector @@ -242,6 +246,10 @@ com.google.code.gson gson + + com.amazonaws + aws-java-sdk + org.codehaus.woodstox 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 5591e63cb5..9c22d0437d 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 @@ -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 EmbeddingStoreConfig convert() { private static ArrayList 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 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 getApiKeyDependency() { @@ -114,17 +117,21 @@ private static List getPathDependency() { private static List 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 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 getPortDependency() { @@ -136,16 +143,20 @@ private static List getPortDependency() { private static List 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 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")); } } diff --git a/common/src/main/java/dev/langchain4j/opensearch/spring/EmbeddingStoreProperties.java b/common/src/main/java/dev/langchain4j/opensearch/spring/EmbeddingStoreProperties.java new file mode 100644 index 0000000000..6ceeea0be0 --- /dev/null +++ b/common/src/main/java/dev/langchain4j/opensearch/spring/EmbeddingStoreProperties.java @@ -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; +} diff --git a/common/src/main/java/dev/langchain4j/opensearch/spring/OpenSearchAutoConfig.java b/common/src/main/java/dev/langchain4j/opensearch/spring/OpenSearchAutoConfig.java new file mode 100644 index 0000000000..a7252517b2 --- /dev/null +++ b/common/src/main/java/dev/langchain4j/opensearch/spring/OpenSearchAutoConfig.java @@ -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()); + } +} diff --git a/common/src/main/java/dev/langchain4j/opensearch/spring/OpenSearchEmbeddingStoreFactory.java b/common/src/main/java/dev/langchain4j/opensearch/spring/OpenSearchEmbeddingStoreFactory.java new file mode 100644 index 0000000000..b6c4b73f8c --- /dev/null +++ b/common/src/main/java/dev/langchain4j/opensearch/spring/OpenSearchEmbeddingStoreFactory.java @@ -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 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(); + } +} diff --git a/common/src/main/java/dev/langchain4j/opensearch/spring/Properties.java b/common/src/main/java/dev/langchain4j/opensearch/spring/Properties.java new file mode 100644 index 0000000000..9fff1980d7 --- /dev/null +++ b/common/src/main/java/dev/langchain4j/opensearch/spring/Properties.java @@ -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; +} diff --git a/common/src/main/java/dev/langchain4j/store/embedding/EmbeddingStoreFactoryProvider.java b/common/src/main/java/dev/langchain4j/store/embedding/EmbeddingStoreFactoryProvider.java index 828a0e8fcb..5e4f44e53d 100644 --- a/common/src/main/java/dev/langchain4j/store/embedding/EmbeddingStoreFactoryProvider.java +++ b/common/src/main/java/dev/langchain4j/store/embedding/EmbeddingStoreFactoryProvider.java @@ -6,6 +6,7 @@ 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 static EmbeddingStoreFactory getFactory(EmbeddingStoreConfig embeddingSto 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()); } diff --git a/common/src/main/java/dev/langchain4j/store/embedding/EmbeddingStoreType.java b/common/src/main/java/dev/langchain4j/store/embedding/EmbeddingStoreType.java index bb533b0f65..5cbda10003 100644 --- a/common/src/main/java/dev/langchain4j/store/embedding/EmbeddingStoreType.java +++ b/common/src/main/java/dev/langchain4j/store/embedding/EmbeddingStoreType.java @@ -1,5 +1,5 @@ package dev.langchain4j.store.embedding; public enum EmbeddingStoreType { - IN_MEMORY, MILVUS, CHROMA, PGVECTOR + IN_MEMORY, MILVUS, CHROMA, PGVECTOR, OPENSEARCH } diff --git a/pom.xml b/pom.xml index 0f707286bd..f05f09a7d4 100644 --- a/pom.xml +++ b/pom.xml @@ -79,6 +79,7 @@ false 4.2.1 3.0.0 + 1.12.780 @@ -173,6 +174,11 @@ langchain4j-milvus ${langchain4j.version} + + dev.langchain4j + langchain4j-opensearch + ${langchain4j.version} + dev.langchain4j langchain4j-pgvector @@ -213,6 +219,11 @@ springdoc-openapi-starter-webmvc-ui 2.1.0 + + com.amazonaws + aws-java-sdk + ${aws-java-sdk.version} +