feat:add opensearch (#2049)

common/pom.xml

@@ -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>

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 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"));
     }
 }

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;
+}

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());
+    }
+}

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<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();
+    }
+}

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;
+}

common/src/main/java/dev/langchain4j/store/embedding/EmbeddingStoreFactoryProvider.java

@@ -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());
     }

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
 }

pom.xml

@@ -79,6 +79,7 @@
         <spotless.skip>false</spotless.skip>
         <stax2.version>4.2.1</stax2.version>
         <io.springfox.version>3.0.0</io.springfox.version>
+        <aws-java-sdk.version>1.12.780</aws-java-sdk.version>
     </properties>
 
     <dependencyManagement>
@@ -173,6 +174,11 @@
                 <artifactId>langchain4j-milvus</artifactId>
                 <version>${langchain4j.version}</version>
             </dependency>
+            <dependency>
+                <groupId>dev.langchain4j</groupId>
+                <artifactId>langchain4j-opensearch</artifactId>
+                <version>${langchain4j.version}</version>
+            </dependency>
             <dependency>
                 <groupId>dev.langchain4j</groupId>
                 <artifactId>langchain4j-pgvector</artifactId>
@@ -213,6 +219,11 @@
                 <artifactId>springdoc-openapi-starter-webmvc-ui</artifactId>
                 <version>2.1.0</version>
             </dependency>
+            <dependency>
+                <groupId>com.amazonaws</groupId>
+                <artifactId>aws-java-sdk</artifactId>
+                <version>${aws-java-sdk.version}</version>
+            </dependency>
         </dependencies>
     </dependencyManagement>