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

Add Elasticsearch vector store auto-configuration

Browse Source 
* Add spring-ai-elasticsearch-store-spring-boot-starter
* Register spring-ai-elasticsearch-store-spring-boot-starter in bom

Fixes #475
This commit is contained in:
Eddú Meléndez
2024-03-20 15:26:33 -06:00
committed by Christian Tzolov
parent 2fbbaa3a6b
commit b3f2516b7d
8 changed files with 270 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
pom.xml
View File

@@ -69,6 +69,7 @@
<module>spring-ai-spring-boot-starters/spring-ai-starter-anthropic</module>
<module>vector-stores/spring-ai-elasticsearch-store</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-watsonx-ai</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-elasticsearch-store</module>
</modules>
<organization>

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

@@ -354,6 +354,12 @@
<artifactId>spring-ai-spring-boot-testcontainers</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-elasticsearch-store-spring-boot-starter</artifactId>
<version>${project.version}</version>
</dependency>
</dependencies>
</dependencyManagement>

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

@@ -244,6 +244,13 @@
<optional>true</optional>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-elasticsearch-store</artifactId>
<version>${project.parent.version}</version>
<optional>true</optional>
</dependency>
<!-- test dependencies -->
<dependency>
@@ -324,6 +331,13 @@
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>elasticsearch</artifactId>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.skyscreamer</groupId>
<artifactId>jsonassert</artifactId>

49
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfiguration.java Normal file
View File

@@ -0,0 +1,49 @@
/*
* 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.elasticsearch;
import org.elasticsearch.client.RestClient;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.vectorstore.ElasticsearchVectorStore;
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.elasticsearch.ElasticsearchClientAutoConfiguration;
import org.springframework.boot.autoconfigure.elasticsearch.ElasticsearchRestClientAutoConfiguration;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.util.StringUtils;
/**
* @author Eddú Meléndez
* @since 1.0.0
*/
@AutoConfiguration(after = ElasticsearchRestClientAutoConfiguration.class)
@ConditionalOnClass({ ElasticsearchVectorStore.class, EmbeddingClient.class, RestClient.class })
@EnableConfigurationProperties(ElasticsearchVectorStoreProperties.class)
class ElasticsearchVectorStoreAutoConfiguration {
@Bean
@ConditionalOnMissingBean
ElasticsearchVectorStore vectorStore(ElasticsearchVectorStoreProperties properties, RestClient restClient,
EmbeddingClient embeddingClient) {
if (StringUtils.hasText(properties.getIndexName())) {
return new ElasticsearchVectorStore(properties.getIndexName(), restClient, embeddingClient);
}
return new ElasticsearchVectorStore(restClient, embeddingClient);
}
}

37
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreProperties.java Normal file
View File

@@ -0,0 +1,37 @@
/*
* 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.elasticsearch;
import org.springframework.boot.context.properties.ConfigurationProperties;
/**
* @author Eddú Meléndez
* @since 1.0.0
*/
@ConfigurationProperties(prefix = "spring.ai.vectorstore.elasticsearch")
public class ElasticsearchVectorStoreProperties {
private String indexName;
public String getIndexName() {
return this.indexName;
}
public void setIndexName(String indexName) {
this.indexName = indexName;
}
}

1
spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports
View File

@@ -29,3 +29,4 @@ org.springframework.ai.autoconfigure.postgresml.PostgresMlAutoConfiguration
org.springframework.ai.autoconfigure.vectorstore.mongo.MongoDBAtlasVectorStoreAutoConfiguration
org.springframework.ai.autoconfigure.anthropic.AnthropicAutoConfiguration
org.springframework.ai.autoconfigure.watsonxai.WatsonxAiAutoConfiguration
org.springframework.ai.autoconfigure.vectorstore.elasticsearch.ElasticsearchVectorStoreAutoConfiguration

120
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfigurationIT.java Normal file
View File

@@ -0,0 +1,120 @@
/*
* 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.elasticsearch;
import org.awaitility.Awaitility;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration;
import org.springframework.ai.autoconfigure.retry.SpringAiRetryAutoConfiguration;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.ElasticsearchVectorStore;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.autoconfigure.elasticsearch.ElasticsearchRestClientAutoConfiguration;
import org.springframework.boot.autoconfigure.web.client.RestClientAutoConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.core.io.DefaultResourceLoader;
import org.testcontainers.elasticsearch.ElasticsearchContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import static org.assertj.core.api.Assertions.assertThat;
import static org.hamcrest.Matchers.hasSize;
@Testcontainers
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
class ElasticsearchVectorStoreAutoConfigurationIT {
@Container
private static final ElasticsearchContainer elasticsearchContainer = new ElasticsearchContainer(
"docker.elastic.co/elasticsearch/elasticsearch:8.12.2")
.withEnv("xpack.security.enabled", "false");
private static final String DEFAULT = "default cosine similarity";
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(ElasticsearchRestClientAutoConfiguration.class,
ElasticsearchVectorStoreAutoConfiguration.class, RestClientAutoConfiguration.class,
SpringAiRetryAutoConfiguration.class, OpenAiAutoConfiguration.class))
.withPropertyValues("spring.elasticsearch.uris=" + elasticsearchContainer.getHttpHostAddress(),
"spring.ai.openai.api-key=" + System.getenv("OPENAI_API_KEY"));
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { DEFAULT, """
double value = dotProduct(params.query_vector, 'embedding');
return sigmoid(1, Math.E, -value);
""", "1 / (1 + l1norm(params.query_vector, 'embedding'))",
"1 / (1 + l2norm(params.query_vector, 'embedding'))" })
public void addAndSearchTest(String similarityFunction) {
this.contextRunner.run(context -> {
ElasticsearchVectorStore vectorStore = context.getBean(ElasticsearchVectorStore.class);
if (!DEFAULT.equals(similarityFunction)) {
vectorStore.withSimilarityFunction(similarityFunction);
}
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);
}
}
}

42
spring-ai-spring-boot-starters/spring-ai-starter-elasticsearch-store/pom.xml Normal file
View File

@@ -0,0 +1,42 @@
<?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-elasticsearch-store-spring-boot-starter</artifactId>
<packaging>jar</packaging>
<name>Spring AI Starter - Elasticsearch Store</name>
<description>Spring AI Elasticsearch 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-elasticsearch-store</artifactId>
<version>${project.parent.version}</version>
</dependency>
</dependencies>
</project>