Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add AWS OpenSearch AutoConfiguration

Browse Source 
Currently, in order to use an OpenSearch instance provided by AWS,
 additional steps are needed. This commit introduces the required
 configuration.

 Add new starter and update docs
This commit is contained in:
Eddú Meléndez
2024-06-19 20:18:25 -05:00
committed by Christian Tzolov
parent 84737ecc01
commit d276d17c1f
9 changed files with 379 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
pom.xml
View File

@@ -77,6 +77,7 @@
<module>models/spring-ai-zhipuai</module>
<module>models/spring-ai-moonshot</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-anthropic</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-aws-opensearch-store</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-azure-openai</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-bedrock-ai</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-huggingface</module>
@@ -177,6 +178,7 @@
<qdrant.version>1.9.1</qdrant.version>
<typesense.version>0.5.0</typesense.version>
<opensearch-client.version>2.10.1</opensearch-client.version>
<awssdk.version>2.20.161</awssdk.version>
<!-- testing dependencies -->
<httpclient5.version>5.3.1</httpclient5.version>

6
spring-ai-bom/pom.xml
View File

@@ -307,6 +307,12 @@
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-aws-opensearch-store-spring-boot-starter</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-opensearch-store-spring-boot-starter</artifactId>

31
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/opensearch.adoc
View File

@@ -114,7 +114,7 @@ List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Sprin
=== Configuration properties
You can use the following properties in your Spring Boot configuration to customize the PGVector vector store.
You can use the following properties in your Spring Boot configuration to customize the OpenSearch vector store.
[cols="2,5,1"]
|===
@@ -134,6 +134,11 @@ fields are stored and indexed. |
}
}
}
|`spring.opensearch.aws.host`| Hostname of the OpenSearch instance. | -
|`spring.opensearch.aws.service-name`| AWS service name for the OpenSearch instance. | -
|`spring.opensearch.aws.access-key`| AWS access key for the OpenSearch instance. | -
|`spring.opensearch.aws.secret-key`| AWS secret key for the OpenSearch instance. | -
|`spring.opensearch.aws.region`| AWS region for the OpenSearch instance. | -
|===
=== Customizing OpenSearch Client Configuration
@@ -148,9 +153,8 @@ To enable it, add the following dependency to your project's Maven `pom.xml` fil
[source,xml]
----
<dependency>
<groupId>software.amazon.awssdk</groupId>
<artifactId>apache-client</artifactId>
<version>2.25.40</version>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-aws-opensearch-store-spring-boot-starter</artifactId>
</dependency>
----
@@ -159,24 +163,7 @@ or to your Gradle `build.gradle` build file.
[source,groovy]
----
dependencies {
implementation 'software.amazon.awssdk:apache-client:2.25.40'
}
----
Here is an example of the needed bean:
[source,java]
----
@Bean
public OpenSearchClient openSearchClient() {
return new OpenSearchClient(
new AwsSdk2Transport(
ApacheHttpClient.builder().build(),
"search-...us-west-2.es.amazonaws.com", // OpenSearch endpoint, without https://
"es",
Region.US_WEST_2, // signing service region
AwsSdk2TransportOptions.builder().build())
);
implementation 'org.springframework.ai:spring-ai-aws-opensearch-store-spring-boot-starter'
}
----

13
spring-ai-spring-boot-autoconfigure/pom.xml
View File

@@ -342,6 +342,13 @@
<optional>true</optional>
</dependency>
<dependency>
<groupId>software.amazon.awssdk</groupId>
<artifactId>apache-client</artifactId>
<version>${awssdk.version}</version>
<optional>true</optional>
</dependency>
<!-- test dependencies -->
<dependency>
@@ -435,6 +442,12 @@
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>localstack</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>milvus</artifactId>

107
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/opensearch/OpenSearchVectorStoreAutoConfiguration.java
View File

@@ -20,14 +20,24 @@ 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.opensearch.OpenSearchClient;
import org.opensearch.client.transport.OpenSearchTransport;
import org.opensearch.client.transport.aws.AwsSdk2Transport;
import org.opensearch.client.transport.aws.AwsSdk2TransportOptions;
import org.opensearch.client.transport.httpclient5.ApacheHttpClient5TransportBuilder;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.OpenSearchVectorStore;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingClass;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import software.amazon.awssdk.http.SdkHttpClient;
import software.amazon.awssdk.http.apache.ApacheHttpClient;
import software.amazon.awssdk.regions.Region;
import java.net.URISyntaxException;
import java.util.List;
@@ -48,45 +58,76 @@ public class OpenSearchVectorStoreAutoConfiguration {
@ConditionalOnMissingBean
OpenSearchVectorStore vectorStore(OpenSearchVectorStoreProperties properties, OpenSearchClient openSearchClient,
EmbeddingModel embeddingModel) {
return new OpenSearchVectorStore(
Optional.ofNullable(properties.getIndexName()).orElse(OpenSearchVectorStore.DEFAULT_INDEX_NAME),
openSearchClient, embeddingModel, Optional.ofNullable(properties.getMappingJson())
.orElse(OpenSearchVectorStore.DEFAULT_MAPPING_EMBEDDING_TYPE_KNN_VECTOR_DIMENSION_1536));
var indexName = Optional.ofNullable(properties.getIndexName()).orElse(OpenSearchVectorStore.DEFAULT_INDEX_NAME);
var mappingJson = Optional.ofNullable(properties.getMappingJson())
.orElse(OpenSearchVectorStore.DEFAULT_MAPPING_EMBEDDING_TYPE_KNN_VECTOR_DIMENSION_1536);
return new OpenSearchVectorStore(indexName, openSearchClient, embeddingModel, mappingJson);
}
@Bean
@ConditionalOnMissingBean
OpenSearchClient openSearchClient(OpenSearchConnectionDetails connectionDetails) {
HttpHost[] httpHosts = connectionDetails.getUris()
.stream()
.map(s -> createHttpHost(s))
.toArray(HttpHost[]::new);
ApacheHttpClient5TransportBuilder transportBuilder = ApacheHttpClient5TransportBuilder.builder(httpHosts);
@Configuration(proxyBeanMethods = false)
@ConditionalOnMissingClass({ "software.amazon.awssdk.regions.Region",
"software.amazon.awssdk.http.apache.ApacheHttpClient" })
static class OpenSearchConfiguration {
Optional.ofNullable(connectionDetails.getUsername())
.map(username -> createBasicCredentialsProvider(httpHosts[0], username, connectionDetails.getPassword()))
.ifPresent(basicCredentialsProvider -> transportBuilder
.setHttpClientConfigCallback(httpAsyncClientBuilder -> httpAsyncClientBuilder
.setDefaultCredentialsProvider(basicCredentialsProvider)));
@Bean
@ConditionalOnMissingBean
OpenSearchClient openSearchClient(OpenSearchVectorStoreProperties properties) {
HttpHost[] httpHosts = properties.getUris().stream().map(s -> createHttpHost(s)).toArray(HttpHost[]::new);
ApacheHttpClient5TransportBuilder transportBuilder = ApacheHttpClient5TransportBuilder.builder(httpHosts);
return new OpenSearchClient(transportBuilder.build());
}
private BasicCredentialsProvider createBasicCredentialsProvider(HttpHost httpHost, String username,
String password) {
BasicCredentialsProvider basicCredentialsProvider = new BasicCredentialsProvider();
basicCredentialsProvider.setCredentials(new AuthScope(httpHost),
new UsernamePasswordCredentials(username, password.toCharArray()));
return basicCredentialsProvider;
}
private HttpHost createHttpHost(String s) {
try {
return HttpHost.create(s);
Optional.ofNullable(properties.getUsername())
.map(username -> createBasicCredentialsProvider(httpHosts[0], username, properties.getPassword()))
.ifPresent(basicCredentialsProvider -> transportBuilder
.setHttpClientConfigCallback(httpAsyncClientBuilder -> httpAsyncClientBuilder
.setDefaultCredentialsProvider(basicCredentialsProvider)));
return new OpenSearchClient(transportBuilder.build());
}
catch (URISyntaxException e) {
throw new RuntimeException(e);
private BasicCredentialsProvider createBasicCredentialsProvider(HttpHost httpHost, String username,
String password) {
BasicCredentialsProvider basicCredentialsProvider = new BasicCredentialsProvider();
basicCredentialsProvider.setCredentials(new AuthScope(httpHost),
new UsernamePasswordCredentials(username, password.toCharArray()));
return basicCredentialsProvider;
}
private HttpHost createHttpHost(String s) {
try {
return HttpHost.create(s);
}
catch (URISyntaxException e) {
throw new RuntimeException(e);
}
}
}
@Configuration(proxyBeanMethods = false)
@ConditionalOnClass({ Region.class, ApacheHttpClient.class })
static class AwsOpenSearchConfiguration {
@Bean
@ConditionalOnMissingBean
OpenSearchClient openSearchClient(OpenSearchVectorStoreProperties properties, AwsSdk2TransportOptions options) {
OpenSearchVectorStoreProperties.Aws aws = properties.getAws();
Region region = Region.of(aws.getRegion());
SdkHttpClient httpClient = ApacheHttpClient.builder().build();
OpenSearchTransport transport = new AwsSdk2Transport(httpClient, aws.getHost(), aws.getServiceName(),
region, options);
return new OpenSearchClient(transport);
}
@Bean
@ConditionalOnMissingBean
AwsSdk2TransportOptions options(OpenSearchVectorStoreProperties properties) {
OpenSearchVectorStoreProperties.Aws aws = properties.getAws();
return AwsSdk2TransportOptions.builder()
.setCredentials(StaticCredentialsProvider
.create(AwsBasicCredentials.create(aws.getAccessKey(), aws.getSecretKey())))
.build();
}
}
static class PropertiesOpenSearchConnectionDetails implements OpenSearchConnectionDetails {

64
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/opensearch/OpenSearchVectorStoreProperties.java
View File

@@ -37,6 +37,8 @@ public class OpenSearchVectorStoreProperties {
private String mappingJson;
private Aws aws = new Aws();
public List<String> getUris() {
return uris;
}
@@ -77,4 +79,66 @@ public class OpenSearchVectorStoreProperties {
this.mappingJson = mappingJson;
}
public Aws getAws() {
return this.aws;
}
public void setAws(Aws aws) {
this.aws = aws;
}
static class Aws {
private String host;
private String serviceName;
private String accessKey;
private String secretKey;
private String region;
public String getHost() {
return this.host;
}
public void setHost(String host) {
this.host = host;
}
public String getServiceName() {
return this.serviceName;
}
public void setServiceName(String serviceName) {
this.serviceName = serviceName;
}
public String getAccessKey() {
return this.accessKey;
}
public void setAccessKey(String accessKey) {
this.accessKey = accessKey;
}
public String getSecretKey() {
return this.secretKey;
}
public void setSecretKey(String secretKey) {
this.secretKey = secretKey;
}
public String getRegion() {
return this.region;
}
public void setRegion(String region) {
this.region = region;
}
}
}

159
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/opensearch/AwsOpenSearchVectorStoreAutoConfigurationIT.java Normal file
View File

@@ -0,0 +1,159 @@
/*
* Copyright 2023 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.autoconfigure.vectorstore.opensearch;
import com.jayway.jsonpath.JsonPath;
import net.minidev.json.JSONArray;
import org.awaitility.Awaitility;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.springframework.ai.autoconfigure.retry.SpringAiRetryAutoConfiguration;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.transformers.TransformersEmbeddingModel;
import org.springframework.ai.vectorstore.OpenSearchVectorStore;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.DefaultResourceLoader;
import org.testcontainers.containers.localstack.LocalStackContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.utility.DockerImageName;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import static org.assertj.core.api.Assertions.assertThat;
import static org.awaitility.Awaitility.await;
import static org.hamcrest.Matchers.hasSize;
@Testcontainers
class AwsOpenSearchVectorStoreAutoConfigurationIT {
@Container
private static final LocalStackContainer localstack = new LocalStackContainer(
DockerImageName.parse("localstack/localstack:3.5.0"))
.withEnv("LOCALSTACK_HOST", "localhost.localstack.cloud");
private static final String DOCUMENT_INDEX = "auto-spring-ai-document-index";
private List<Document> documents = List.of(
new Document("1", getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")),
new Document("2", getText("classpath:/test/data/time.shelter.txt"), Map.of()),
new Document("3", getText("classpath:/test/data/great.depression.txt"), Map.of("meta2", "meta2")));
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(OpenSearchVectorStoreAutoConfiguration.class,
SpringAiRetryAutoConfiguration.class))
.withUserConfiguration(Config.class)
.withPropertyValues(
OpenSearchVectorStoreProperties.CONFIG_PREFIX + ".aws.host="
+ String.format("testcontainers-domain.%s.opensearch.localhost.localstack.cloud:%s",
localstack.getRegion(), localstack.getMappedPort(4566)),
OpenSearchVectorStoreProperties.CONFIG_PREFIX + ".aws.service-name=opensearch",
OpenSearchVectorStoreProperties.CONFIG_PREFIX + ".aws.region=" + localstack.getRegion(),
OpenSearchVectorStoreProperties.CONFIG_PREFIX + ".aws.access-key=" + localstack.getAccessKey(),
OpenSearchVectorStoreProperties.CONFIG_PREFIX + ".aws.secret-key=" + localstack.getSecretKey(),
OpenSearchVectorStoreProperties.CONFIG_PREFIX + ".indexName=" + DOCUMENT_INDEX,
OpenSearchVectorStoreProperties.CONFIG_PREFIX + ".mappingJson=" + """
{
"properties":{
"embedding":{
"type":"knn_vector",
"dimension":384
}
}
}
""");
@BeforeAll
static void beforeAll() throws IOException, InterruptedException {
String[] createDomainCmd = { "awslocal", "opensearch", "create-domain", "--domain-name",
"testcontainers-domain", "--region", localstack.getRegion() };
localstack.execInContainer(createDomainCmd);
String[] describeDomainCmd = { "awslocal", "opensearch", "describe-domain", "--domain-name",
"testcontainers-domain", "--region", localstack.getRegion() };
await().pollInterval(Duration.ofSeconds(30)).atMost(Duration.ofSeconds(300)).untilAsserted(() -> {
org.testcontainers.containers.Container.ExecResult execResult = localstack
.execInContainer(describeDomainCmd);
String response = execResult.getStdout();
JSONArray processed = JsonPath.read(response, "$.DomainStatus[?(@.Processing == false)]");
assertThat(processed).isNotEmpty();
});
}
@Test
public void addAndSearchTest() {
this.contextRunner.run(context -> {
OpenSearchVectorStore vectorStore = context.getBean(OpenSearchVectorStore.class);
vectorStore.add(documents);
Awaitility.await()
.until(() -> vectorStore
.similaritySearch(SearchRequest.query("Great Depression").withTopK(1).withSimilarityThreshold(0)),
hasSize(1));
List<Document> results = vectorStore
.similaritySearch(SearchRequest.query("Great Depression").withTopK(1).withSimilarityThreshold(0));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getContent()).contains("The Great Depression (1929–1939) was an economic shock");
assertThat(resultDoc.getMetadata()).hasSize(2);
assertThat(resultDoc.getMetadata()).containsKey("meta2");
assertThat(resultDoc.getMetadata()).containsKey("distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(Document::getId).toList());
Awaitility.await()
.until(() -> vectorStore
.similaritySearch(SearchRequest.query("Great Depression").withTopK(1).withSimilarityThreshold(0)),
hasSize(0));
});
}
private String getText(String uri) {
var resource = new DefaultResourceLoader().getResource(uri);
try {
return resource.getContentAsString(StandardCharsets.UTF_8);
}
catch (IOException e) {
throw new RuntimeException(e);
}
}
@Configuration(proxyBeanMethods = false)
static class Config {
@Bean
public EmbeddingModel embeddingModel() {
return new TransformersEmbeddingModel();
}
}
}

4
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/opensearch/OpenSearchVectorStoreAutoConfigurationIT.java
View File

@@ -25,6 +25,7 @@ import org.springframework.ai.transformers.TransformersEmbeddingModel;
import org.springframework.ai.vectorstore.OpenSearchVectorStore;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.test.context.FilteredClassLoader;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
@@ -32,6 +33,8 @@ import org.springframework.core.io.DefaultResourceLoader;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.utility.DockerImageName;
import software.amazon.awssdk.http.apache.ApacheHttpClient;
import software.amazon.awssdk.regions.Region;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
@@ -58,6 +61,7 @@ class OpenSearchVectorStoreAutoConfigurationIT {
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(OpenSearchVectorStoreAutoConfiguration.class,
SpringAiRetryAutoConfiguration.class))
.withClassLoader(new FilteredClassLoader(Region.class, ApacheHttpClient.class))
.withUserConfiguration(Config.class)
.withPropertyValues(
OpenSearchVectorStoreProperties.CONFIG_PREFIX + ".uris=" + opensearchContainer.getHttpHostAddress(),

48
spring-ai-spring-boot-starters/spring-ai-starter-aws-opensearch-store/pom.xml Normal file
View File

@@ -0,0 +1,48 @@
<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
<modelVersion>4.0.0</modelVersion>
<parent>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai</artifactId>
<version>1.0.0-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>spring-ai-aws-opensearch-store-spring-boot-starter</artifactId>
<packaging>jar</packaging>
<name>Spring AI Starter - AWS OpenSearch Store</name>
<description>Spring AI AWS OpenSearch Store Auto Configuration</description>
<url>https://github.com/spring-projects/spring-ai</url>
<scm>
<url>https://github.com/spring-projects/spring-ai</url>
<connection>git://github.com/spring-projects/spring-ai.git</connection>
<developerConnection>git@github.com:spring-projects/spring-ai.git</developerConnection>
</scm>
<dependencies>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter</artifactId>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-spring-boot-autoconfigure</artifactId>
<version>${project.parent.version}</version>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-opensearch-store</artifactId>
<version>${project.parent.version}</version>
</dependency>
<dependency>
<groupId>software.amazon.awssdk</groupId>
<artifactId>apache-client</artifactId>
<version>${awssdk.version}</version>
</dependency>
</dependencies>
</project>