Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat:add opensearch #2049

Merged
merged 1 commit into from
Feb 12, 2025
Merged

feat:add opensearch #2049

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions common/pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down Expand Up @@ -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>
Expand Down
35 changes: 23 additions & 12 deletions common/src/main/java/com/tencent/supersonic/common/config/EmbeddingStoreParameterConfig.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -87,16 +87,19 @@ public EmbeddingStoreConfig convert() {
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() {
Expand All @@ -114,17 +117,21 @@ private static List<Parameter.Dependency> getPathDependency() {
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() {
Expand All @@ -136,16 +143,20 @@ private static List<Parameter.Dependency> getPortDependency() {
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 changes: 29 additions & 0 deletions common/src/main/java/dev/langchain4j/opensearch/spring/EmbeddingStoreProperties.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 changes: 19 additions & 0 deletions common/src/main/java/dev/langchain4j/opensearch/spring/OpenSearchAutoConfig.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 changes: 59 additions & 0 deletions common/src/main/java/dev/langchain4j/opensearch/spring/OpenSearchEmbeddingStoreFactory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 changes: 17 additions & 0 deletions common/src/main/java/dev/langchain4j/opensearch/spring/Properties.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 changes: 6 additions & 0 deletions common/src/main/java/dev/langchain4j/store/embedding/EmbeddingStoreFactoryProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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());
}
Expand Down
2 changes: 1 addition & 1 deletion common/src/main/java/dev/langchain4j/store/embedding/EmbeddingStoreType.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
package dev.langchain4j.store.embedding;

public enum EmbeddingStoreType {
IN_MEMORY, MILVUS, CHROMA, PGVECTOR
IN_MEMORY, MILVUS, CHROMA, PGVECTOR, OPENSEARCH
}
11 changes: 11 additions & 0 deletions pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down Expand Up @@ -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>
Expand Down Expand Up @@ -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>

Expand Down
Loading