diff --git a/pom.xml b/pom.xml index 162c6ff8c70..33e48a02c0e 100644 --- a/pom.xml +++ b/pom.xml @@ -69,6 +69,7 @@ spring-ai-spring-boot-starters/spring-ai-starter-anthropic vector-stores/spring-ai-elasticsearch-store spring-ai-spring-boot-starters/spring-ai-starter-watsonx-ai + spring-ai-spring-boot-starters/spring-ai-starter-elasticsearch-store diff --git a/spring-ai-bom/pom.xml b/spring-ai-bom/pom.xml index 19c7235de1f..bc1ebc3aa2b 100644 --- a/spring-ai-bom/pom.xml +++ b/spring-ai-bom/pom.xml @@ -354,6 +354,12 @@ spring-ai-spring-boot-testcontainers ${project.version} + + + org.springframework.ai + spring-ai-elasticsearch-store-spring-boot-starter + ${project.version} + diff --git a/spring-ai-spring-boot-autoconfigure/pom.xml b/spring-ai-spring-boot-autoconfigure/pom.xml index 5b7bbc2a9f7..84fe689ef29 100644 --- a/spring-ai-spring-boot-autoconfigure/pom.xml +++ b/spring-ai-spring-boot-autoconfigure/pom.xml @@ -244,6 +244,13 @@ true + + org.springframework.ai + spring-ai-elasticsearch-store + ${project.parent.version} + true + + @@ -324,6 +331,13 @@ test + + org.testcontainers + elasticsearch + ${testcontainers.version} + test + + org.skyscreamer jsonassert diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfiguration.java new file mode 100644 index 00000000000..a2706e27aee --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfiguration.java @@ -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); + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreProperties.java new file mode 100644 index 00000000000..5f4f5ccc796 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreProperties.java @@ -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; + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports b/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports index 418ba7a1af8..f614b7711ec 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports +++ b/spring-ai-spring-boot-autoconfigure/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports @@ -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 diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfigurationIT.java new file mode 100644 index 00000000000..31963727a73 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfigurationIT.java @@ -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 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 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); + } + } + +} diff --git a/spring-ai-spring-boot-starters/spring-ai-starter-elasticsearch-store/pom.xml b/spring-ai-spring-boot-starters/spring-ai-starter-elasticsearch-store/pom.xml new file mode 100644 index 00000000000..363ca6d1151 --- /dev/null +++ b/spring-ai-spring-boot-starters/spring-ai-starter-elasticsearch-store/pom.xml @@ -0,0 +1,42 @@ + + + 4.0.0 + + org.springframework.ai + spring-ai + 1.0.0-SNAPSHOT + ../../pom.xml + + spring-ai-elasticsearch-store-spring-boot-starter + jar + Spring AI Starter - Elasticsearch Store + Spring AI Elasticsearch Store Auto Configuration + https://github.com/spring-projects/spring-ai + + + https://github.com/spring-projects/spring-ai + git://github.com/spring-projects/spring-ai.git + git@github.com:spring-projects/spring-ai.git + + + + + + org.springframework.boot + spring-boot-starter + + + + org.springframework.ai + spring-ai-spring-boot-autoconfigure + ${project.parent.version} + + + + org.springframework.ai + spring-ai-elasticsearch-store + ${project.parent.version} + + + +