Skip to content

Commit d942d35

Browse files
zyclovezyclove
authored
feat:add opensearch (#2049)
1 parent 198c7c6 commit d942d35

File tree

9 files changed

+173
-13
lines changed

9 files changed

+173
-13
lines changed

‎common/pom.xml

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,10 @@
174174
<groupId>dev.langchain4j</groupId>
175175
<artifactId>langchain4j-milvus</artifactId>
176176
</dependency>
177+
<dependency>
178+
<groupId>dev.langchain4j</groupId>
179+
<artifactId>langchain4j-opensearch</artifactId>
180+
</dependency>
177181
<dependency>
178182
<groupId>dev.langchain4j</groupId>
179183
<artifactId>langchain4j-pgvector</artifactId>
@@ -242,6 +246,10 @@
242246
<groupId>com.google.code.gson</groupId>
243247
<artifactId>gson</artifactId>
244248
</dependency>
249+
<dependency>
250+
<groupId>com.amazonaws</groupId>
251+
<artifactId>aws-java-sdk</artifactId>
252+
</dependency>
245253

246254
<dependency>
247255
<groupId>org.codehaus.woodstox</groupId>

‎common/src/main/java/com/tencent/supersonic/common/config/EmbeddingStoreParameterConfig.java

Lines changed: 23 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ public class EmbeddingStoreParameterConfig extends ParameterConfig {
1919

2020
public static final Parameter EMBEDDING_STORE_PROVIDER = new Parameter(
2121
"s2.embedding.store.provider", EmbeddingStoreType.IN_MEMORY.name(), "向量库类型",
22-
"目前支持四种类型：IN_MEMORY、MILVUS、CHROMA、PGVECTOR", "list", MODULE_NAME, getCandidateValues());
22+
"目前支持四种类型：IN_MEMORY、MILVUS、CHROMA、PGVECTOR、OPENSEARCH", "list", MODULE_NAME, getCandidateValues());
2323

2424
public static final Parameter EMBEDDING_STORE_BASE_URL =
2525
new Parameter("s2.embedding.store.base.url", "", "BaseUrl", "", "string", MODULE_NAME,
@@ -87,16 +87,19 @@ public EmbeddingStoreConfig convert() {
8787
private static ArrayList<String> getCandidateValues() {
8888
return Lists.newArrayList(EmbeddingStoreType.IN_MEMORY.name(),
8989
EmbeddingStoreType.MILVUS.name(), EmbeddingStoreType.CHROMA.name(),
90-
EmbeddingStoreType.PGVECTOR.name());
90+
EmbeddingStoreType.PGVECTOR.name(), EmbeddingStoreType.OPENSEARCH.name());
9191
}
9292

9393
private static List<Parameter.Dependency> getBaseUrlDependency() {
9494
return getDependency(EMBEDDING_STORE_PROVIDER.getName(),
9595
Lists.newArrayList(EmbeddingStoreType.MILVUS.name(),
96-
EmbeddingStoreType.CHROMA.name(), EmbeddingStoreType.PGVECTOR.name()),
96+
EmbeddingStoreType.CHROMA.name(),
97+
EmbeddingStoreType.PGVECTOR.name(),
98+
EmbeddingStoreType.OPENSEARCH.name()),
9799
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "http://localhost:19530",
98100
EmbeddingStoreType.CHROMA.name(), "http://localhost:8000",
99-
EmbeddingStoreType.PGVECTOR.name(), "127.0.0.1"));
101+
EmbeddingStoreType.PGVECTOR.name(), "127.0.0.1",
102+
EmbeddingStoreType.OPENSEARCH.name(), "http://localhost:9200"));
100103
}
101104

102105
private static List<Parameter.Dependency> getApiKeyDependency() {
@@ -114,17 +117,21 @@ private static List<Parameter.Dependency> getPathDependency() {
114117
private static List<Parameter.Dependency> getDimensionDependency() {
115118
return getDependency(EMBEDDING_STORE_PROVIDER.getName(),
116119
Lists.newArrayList(EmbeddingStoreType.MILVUS.name(),
117-
EmbeddingStoreType.PGVECTOR.name()),
120+
EmbeddingStoreType.PGVECTOR.name(),
121+
EmbeddingStoreType.OPENSEARCH.name()),
118122
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "384",
119-
EmbeddingStoreType.PGVECTOR.name(), "512"));
123+
EmbeddingStoreType.PGVECTOR.name(), "512",
124+
EmbeddingStoreType.OPENSEARCH.name(), "512"));
120125
}
121126

122127
private static List<Parameter.Dependency> getDatabaseNameDependency() {
123128
return getDependency(EMBEDDING_STORE_PROVIDER.getName(),
124129
Lists.newArrayList(EmbeddingStoreType.MILVUS.name(),
125-
EmbeddingStoreType.PGVECTOR.name()),
130+
EmbeddingStoreType.PGVECTOR.name(),
131+
EmbeddingStoreType.OPENSEARCH.name()),
126132
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "",
127-
EmbeddingStoreType.PGVECTOR.name(), "postgres"));
133+
EmbeddingStoreType.PGVECTOR.name(), "postgres",
134+
EmbeddingStoreType.OPENSEARCH.name(), "ai_sql"));
128135
}
129136

130137
private static List<Parameter.Dependency> getPortDependency() {
@@ -136,16 +143,20 @@ private static List<Parameter.Dependency> getPortDependency() {
136143
private static List<Parameter.Dependency> getUserDependency() {
137144
return getDependency(EMBEDDING_STORE_PROVIDER.getName(),
138145
Lists.newArrayList(EmbeddingStoreType.MILVUS.name(),
139-
EmbeddingStoreType.PGVECTOR.name()),
146+
EmbeddingStoreType.PGVECTOR.name(),
147+
EmbeddingStoreType.OPENSEARCH.name()),
140148
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "milvus",
141-
EmbeddingStoreType.PGVECTOR.name(), "postgres"));
149+
EmbeddingStoreType.PGVECTOR.name(), "postgres",
150+
EmbeddingStoreType.OPENSEARCH.name(), "opensearch"));
142151
}
143152

144153
private static List<Parameter.Dependency> getPasswordDependency() {
145154
return getDependency(EMBEDDING_STORE_PROVIDER.getName(),
146155
Lists.newArrayList(EmbeddingStoreType.MILVUS.name(),
147-
EmbeddingStoreType.PGVECTOR.name()),
156+
EmbeddingStoreType.PGVECTOR.name(),
157+
EmbeddingStoreType.OPENSEARCH.name()),
148158
ImmutableMap.of(EmbeddingStoreType.MILVUS.name(), "milvus",
149-
EmbeddingStoreType.PGVECTOR.name(), "postgres"));
159+
EmbeddingStoreType.PGVECTOR.name(), "postgres",
160+
EmbeddingStoreType.OPENSEARCH.name(), "opensearch"));
150161
}
151162
}

‎common/src/main/java/dev/langchain4j/opensearch/spring/EmbeddingStoreProperties.java

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
package dev.langchain4j.opensearch.spring;
2+
3+
import io.milvus.common.clientenum.ConsistencyLevelEnum;
4+
import io.milvus.param.IndexType;
5+
import io.milvus.param.MetricType;
6+
import lombok.Getter;
7+
import lombok.Setter;
8+
9+
@Getter
10+
@Setter
11+
class EmbeddingStoreProperties {
12+
13+
private String uri;
14+
private String host;
15+
private Integer port;
16+
private String serviceName;
17+
private String region;
18+
private String collectionName;
19+
private Integer dimension;
20+
private IndexType indexType;
21+
private MetricType metricType;
22+
private String token;
23+
private String user;
24+
private String password;
25+
private ConsistencyLevelEnum consistencyLevel;
26+
private Boolean retrieveEmbeddingsOnSearch;
27+
private String databaseName;
28+
private Boolean autoFlushOnInsert;
29+
}

‎common/src/main/java/dev/langchain4j/opensearch/spring/OpenSearchAutoConfig.java

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
package dev.langchain4j.opensearch.spring;
2+
3+
import dev.langchain4j.store.embedding.EmbeddingStoreFactory;
4+
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
5+
import org.springframework.boot.context.properties.EnableConfigurationProperties;
6+
import org.springframework.context.annotation.Bean;
7+
import org.springframework.context.annotation.Configuration;
8+
9+
import static dev.langchain4j.opensearch.spring.Properties.PREFIX;
10+
@Configuration
11+
@EnableConfigurationProperties(dev.langchain4j.opensearch.spring.Properties.class)
12+
public class OpenSearchAutoConfig {
13+
14+
@Bean
15+
@ConditionalOnProperty(PREFIX + ".embedding-store.uri")
16+
EmbeddingStoreFactory milvusChatModel(Properties properties) {
17+
return new OpenSearchEmbeddingStoreFactory(properties.getEmbeddingStore());
18+
}
19+
}

‎common/src/main/java/dev/langchain4j/opensearch/spring/OpenSearchEmbeddingStoreFactory.java

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
package dev.langchain4j.opensearch.spring;
2+
3+
import com.tencent.supersonic.common.pojo.EmbeddingStoreConfig;
4+
import dev.langchain4j.data.segment.TextSegment;
5+
import dev.langchain4j.store.embedding.BaseEmbeddingStoreFactory;
6+
import dev.langchain4j.store.embedding.EmbeddingStore;
7+
import dev.langchain4j.store.embedding.opensearch.OpenSearchEmbeddingStore;
8+
import org.apache.hc.client5.http.auth.AuthScope;
9+
import org.apache.hc.client5.http.auth.UsernamePasswordCredentials;
10+
import org.apache.hc.client5.http.impl.auth.BasicCredentialsProvider;
11+
import org.apache.hc.core5.http.HttpHost;
12+
import org.opensearch.client.transport.aws.AwsSdk2TransportOptions;
13+
import org.springframework.beans.BeanUtils;
14+
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
15+
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
16+
17+
import java.net.URI;
18+
19+
/**
20+
* @author zyc
21+
*/
22+
public class OpenSearchEmbeddingStoreFactory extends BaseEmbeddingStoreFactory {
23+
private final EmbeddingStoreProperties storeProperties;
24+
25+
public OpenSearchEmbeddingStoreFactory(EmbeddingStoreConfig storeConfig) {
26+
this(createPropertiesFromConfig(storeConfig));
27+
}
28+
29+
public OpenSearchEmbeddingStoreFactory(EmbeddingStoreProperties storeProperties) {
30+
this.storeProperties = storeProperties;
31+
}
32+
33+
private static EmbeddingStoreProperties createPropertiesFromConfig(
34+
EmbeddingStoreConfig storeConfig) {
35+
EmbeddingStoreProperties embeddingStore = new EmbeddingStoreProperties();
36+
BeanUtils.copyProperties(storeConfig, embeddingStore);
37+
embeddingStore.setUri(storeConfig.getBaseUrl());
38+
embeddingStore.setToken(storeConfig.getApiKey());
39+
embeddingStore.setDatabaseName(storeConfig.getDatabaseName());
40+
return embeddingStore;
41+
}
42+
43+
@Override
44+
public EmbeddingStore<TextSegment> createEmbeddingStore(String collectionName) {
45+
final AwsSdk2TransportOptions options = AwsSdk2TransportOptions.builder()
46+
.setCredentials(StaticCredentialsProvider.create(AwsBasicCredentials.create(storeProperties.getUser(), storeProperties.getPassword())))
47+
.build();
48+
final String indexName = storeProperties.getDatabaseName() + "_" + collectionName;
49+
return OpenSearchEmbeddingStore.builder().serviceName(storeProperties.getServiceName())
50+
.serverUrl(storeProperties.getUri())
51+
.region(storeProperties.getRegion())
52+
.indexName(indexName)
53+
.userName(storeProperties.getUser())
54+
.password(storeProperties.getPassword())
55+
.apiKey(storeProperties.getToken())
56+
.options(options)
57+
.build();
58+
}
59+
}

‎common/src/main/java/dev/langchain4j/opensearch/spring/Properties.java

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
package dev.langchain4j.opensearch.spring;
2+
3+
import lombok.Getter;
4+
import lombok.Setter;
5+
import org.springframework.boot.context.properties.ConfigurationProperties;
6+
import org.springframework.boot.context.properties.NestedConfigurationProperty;
7+
8+
@Getter
9+
@Setter
10+
@ConfigurationProperties(prefix = Properties.PREFIX)
11+
public class Properties {
12+
13+
static final String PREFIX = "langchain4j.opensearch";
14+
15+
@NestedConfigurationProperty
16+
dev.langchain4j.opensearch.spring.EmbeddingStoreProperties embeddingStore;
17+
}

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import dev.langchain4j.chroma.spring.ChromaEmbeddingStoreFactory;
77
import dev.langchain4j.inmemory.spring.InMemoryEmbeddingStoreFactory;
88
import dev.langchain4j.milvus.spring.MilvusEmbeddingStoreFactory;
9+
import dev.langchain4j.opensearch.spring.OpenSearchEmbeddingStoreFactory;
910
import dev.langchain4j.pgvector.spring.PgvectorEmbeddingStoreFactory;
1011
import org.apache.commons.lang3.StringUtils;
1112

@@ -45,6 +46,11 @@ public static EmbeddingStoreFactory getFactory(EmbeddingStoreConfig embeddingSto
4546
return factoryMap.computeIfAbsent(embeddingStoreConfig,
4647
storeConfig -> new InMemoryEmbeddingStoreFactory(storeConfig));
4748
}
49+
if (EmbeddingStoreType.OPENSEARCH.name()
50+
.equalsIgnoreCase(embeddingStoreConfig.getProvider())) {
51+
return factoryMap.computeIfAbsent(embeddingStoreConfig,
52+
storeConfig -> new OpenSearchEmbeddingStoreFactory(storeConfig));
53+
}
4854
throw new RuntimeException("Unsupported EmbeddingStoreFactory provider: "
4955
+ embeddingStoreConfig.getProvider());
5056
}

‎common/src/main/java/dev/langchain4j/store/embedding/EmbeddingStoreType.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
package dev.langchain4j.store.embedding;
22

33
public enum EmbeddingStoreType {
4-
IN_MEMORY, MILVUS, CHROMA, PGVECTOR
4+
IN_MEMORY, MILVUS, CHROMA, PGVECTOR, OPENSEARCH
55
}

‎pom.xml

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@
7979
<spotless.skip>false</spotless.skip>
8080
<stax2.version>4.2.1</stax2.version>
8181
<io.springfox.version>3.0.0</io.springfox.version>
82+
<aws-java-sdk.version>1.12.780</aws-java-sdk.version>
8283
</properties>
8384

8485
<dependencyManagement>
@@ -173,6 +174,11 @@
173174
<artifactId>langchain4j-milvus</artifactId>
174175
<version>${langchain4j.version}</version>
175176
</dependency>
177+
<dependency>
178+
<groupId>dev.langchain4j</groupId>
179+
<artifactId>langchain4j-opensearch</artifactId>
180+
<version>${langchain4j.version}</version>
181+
</dependency>
176182
<dependency>
177183
<groupId>dev.langchain4j</groupId>
178184
<artifactId>langchain4j-pgvector</artifactId>
@@ -213,6 +219,11 @@
213219
<artifactId>springdoc-openapi-starter-webmvc-ui</artifactId>
214220
<version>2.1.0</version>
215221
</dependency>
222+
<dependency>
223+
<groupId>com.amazonaws</groupId>
224+
<artifactId>aws-java-sdk</artifactId>
225+
<version>${aws-java-sdk.version}</version>
226+
</dependency>
216227
</dependencies>
217228
</dependencyManagement>
218229

0 commit comments

Comments
 (0)