From 1ca729ae0352ca4b474d5a6aeff1dd2fa1450a5b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Edd=C3=BA=20Mel=C3=A9ndez?= Date: Sat, 16 Mar 2024 22:22:48 -0600 Subject: [PATCH 1/2] Add MongoDB Atlas vector store auto-configuration It also provides a starter. --- pom.xml | 1 + spring-ai-spring-boot-autoconfigure/pom.xml | 8 ++ ...goDBAtlasVectorStoreAutoConfiguration.java | 49 +++++++ .../MongoDBAtlasVectorStoreProperties.java | 56 ++++++++ ...ot.autoconfigure.AutoConfiguration.imports | 1 + ...DBAtlasVectorStoreAutoConfigurationIT.java | 125 ++++++++++++++++++ .../pom.xml | 42 ++++++ .../spring-ai-mongodb-atlas-store/pom.xml | 2 +- .../MongoDBAtlasVectorStoreIT.java | 14 +- 9 files changed, 288 insertions(+), 10 deletions(-) create mode 100644 spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreAutoConfiguration.java create mode 100644 spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreProperties.java create mode 100644 spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreAutoConfigurationIT.java create mode 100644 spring-ai-spring-boot-starters/spring-ai-starter-mongodb-atlas-store/pom.xml diff --git a/pom.xml b/pom.xml index 5b73838a291..a959362981e 100644 --- a/pom.xml +++ b/pom.xml @@ -61,6 +61,7 @@ spring-ai-spring-boot-starters/spring-ai-starter-mistral-ai spring-ai-retry vector-stores/spring-ai-mongodb-atlas-store + spring-ai-spring-boot-starters/spring-ai-starter-mongodb-atlas-store diff --git a/spring-ai-spring-boot-autoconfigure/pom.xml b/spring-ai-spring-boot-autoconfigure/pom.xml index 80c6aacf12d..303f178caea 100644 --- a/spring-ai-spring-boot-autoconfigure/pom.xml +++ b/spring-ai-spring-boot-autoconfigure/pom.xml @@ -223,6 +223,14 @@ true + + + org.springframework.ai + spring-ai-mongodb-atlas-store + ${project.parent.version} + true + + diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreAutoConfiguration.java new file mode 100644 index 00000000000..62bc91cd250 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreAutoConfiguration.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.mongo; + +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.ai.vectorstore.MongoDBAtlasVectorStore; +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.data.mongo.MongoDataAutoConfiguration; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.context.annotation.Bean; +import org.springframework.data.mongodb.core.MongoTemplate; + +/** + * @author Eddú Meléndez + */ +@AutoConfiguration(after = MongoDataAutoConfiguration.class) +@ConditionalOnClass({ MongoDBAtlasVectorStore.class, EmbeddingClient.class, MongoTemplate.class }) +@EnableConfigurationProperties(MongoDBAtlasVectorStoreProperties.class) +public class MongoDBAtlasVectorStoreAutoConfiguration { + + @Bean + @ConditionalOnMissingBean + MongoDBAtlasVectorStore vectorStore(MongoTemplate mongoTemplate, EmbeddingClient embeddingClient, + MongoDBAtlasVectorStoreProperties properties) { + MongoDBAtlasVectorStore.MongoDBVectorStoreConfig config = MongoDBAtlasVectorStore.MongoDBVectorStoreConfig + .builder() + .withCollectionName(properties.getCollectionName()) + .withPathName(properties.getPathName()) + .withVectorIndexName(properties.getIndexName()) + .build(); + return new MongoDBAtlasVectorStore(mongoTemplate, embeddingClient, config); + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreProperties.java new file mode 100644 index 00000000000..c31c26a6d65 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreProperties.java @@ -0,0 +1,56 @@ +/* + * 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.mongo; + +import org.springframework.boot.context.properties.ConfigurationProperties; + +/** + * @author Eddú Meléndez + */ +@ConfigurationProperties("spring.ai.vectorstore.mongodb") +public class MongoDBAtlasVectorStoreProperties { + + private String collectionName; + + private String pathName; + + private String indexName; + + public String getCollectionName() { + return this.collectionName; + } + + public void setCollectionName(String collectionName) { + this.collectionName = collectionName; + } + + public String getPathName() { + return this.pathName; + } + + public void setPathName(String pathName) { + this.pathName = pathName; + } + + 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 4c136de2508..ddc79f0ef9f 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 @@ -25,3 +25,4 @@ org.springframework.ai.autoconfigure.vectorstore.neo4j.Neo4jVectorStoreAutoConfi org.springframework.ai.autoconfigure.vectorstore.qdrant.QdrantVectorStoreAutoConfiguration org.springframework.ai.autoconfigure.retry.SpringAiRetryAutoConfiguration org.springframework.ai.autoconfigure.postgresml.PostgresMlAutoConfiguration +org.springframework.ai.autoconfigure.vectorstore.mongo.MongoDBAtlasVectorStoreAutoConfiguration diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreAutoConfigurationIT.java new file mode 100644 index 00000000000..95477bbd466 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreAutoConfigurationIT.java @@ -0,0 +1,125 @@ +/* + * 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.mongo; + +import org.junit.jupiter.api.Test; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.ai.transformers.TransformersEmbeddingClient; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; +import org.springframework.boot.autoconfigure.AutoConfigurations; +import org.springframework.boot.autoconfigure.data.mongo.MongoDataAutoConfiguration; +import org.springframework.boot.autoconfigure.mongo.MongoAutoConfiguration; +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.GenericContainer; +import org.testcontainers.containers.wait.strategy.Wait; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Eddú Meléndez + */ +@Testcontainers +class MongoDBAtlasVectorStoreAutoConfigurationIT { + + @Container + static GenericContainer mongo = new GenericContainer<>("mongodb/atlas:v1.15.1").withPrivilegedMode(true) + .withCommand("/bin/bash", "-c", + "atlas deployments setup local-test --type local --port 27778 --bindIpAll --username root --password root --force && tail -f /dev/null") + .withExposedPorts(27778) + .waitingFor(Wait.forLogMessage(".*Deployment created!.*\\n", 1)) + .withStartupTimeout(Duration.ofMinutes(5)); + + List documents = List.of( + new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", + Collections.singletonMap("meta1", "meta1")), + new Document("Hello World Hello World Hello World Hello World Hello World Hello World Hello World"), + new Document( + "Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression", + Collections.singletonMap("meta2", "meta2"))); + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(MongoAutoConfiguration.class, MongoDataAutoConfiguration.class, + MongoDBAtlasVectorStoreAutoConfiguration.class)) + .withUserConfiguration(Config.class) + .withPropertyValues("spring.data.mongodb.database=springaisample", + "spring.ai.vectorstore.mongodb.collection-name=test_collection", + "spring.ai.vectorstore.mongodb.path-name=test_path", + "spring.ai.vectorstore.mongodb.index-name=text_index", + String.format( + "spring.data.mongodb.uri=" + String.format("mongodb://root:root@%s:%s/?directConnection=true", + mongo.getHost(), mongo.getMappedPort(27778)))); + + @Test + public void addAndSearch() { + contextRunner.run(context -> { + + VectorStore vectorStore = context.getBean(VectorStore.class); + + vectorStore.add(documents); + Thread.sleep(5000); // Await a second for the document to be indexed + + List results = vectorStore.similaritySearch(SearchRequest.query("Great").withTopK(1)); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId()); + assertThat(resultDoc.getContent()).isEqualTo( + "Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression"); + assertThat(resultDoc.getMetadata()).containsEntry("meta2", "meta2"); + + // Remove all documents from the store + vectorStore.delete(documents.stream().map(Document::getId).collect(Collectors.toList())); + + List results2 = vectorStore.similaritySearch(SearchRequest.query("Great").withTopK(1)); + assertThat(results2).isEmpty(); + }); + } + + public static 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 EmbeddingClient embeddingClient() { + return new TransformersEmbeddingClient(); + } + + } + +} \ No newline at end of file diff --git a/spring-ai-spring-boot-starters/spring-ai-starter-mongodb-atlas-store/pom.xml b/spring-ai-spring-boot-starters/spring-ai-starter-mongodb-atlas-store/pom.xml new file mode 100644 index 00000000000..52cb6d05805 --- /dev/null +++ b/spring-ai-spring-boot-starters/spring-ai-starter-mongodb-atlas-store/pom.xml @@ -0,0 +1,42 @@ + + + 4.0.0 + + org.springframework.ai + spring-ai + 1.0.0-SNAPSHOT + ../../pom.xml + + spring-ai-mongodb-atlas-store-spring-boot-starter + jar + Spring AI Starter - MongoDB Atlas Store + Spring AI MongoDB Atlas 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-mongodb-atlas-store + ${project.parent.version} + + + + diff --git a/vector-stores/spring-ai-mongodb-atlas-store/pom.xml b/vector-stores/spring-ai-mongodb-atlas-store/pom.xml index 691bebece9f..fdb24d930c4 100644 --- a/vector-stores/spring-ai-mongodb-atlas-store/pom.xml +++ b/vector-stores/spring-ai-mongodb-atlas-store/pom.xml @@ -39,7 +39,7 @@ org.springframework.ai - spring-ai-openai-spring-boot-starter + spring-ai-openai ${parent.version} test diff --git a/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStoreIT.java b/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStoreIT.java index afc2dcbd743..a8151dacf58 100644 --- a/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStoreIT.java +++ b/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStoreIT.java @@ -19,13 +19,11 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; -import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingClient; import org.springframework.ai.openai.OpenAiEmbeddingClient; import org.springframework.ai.openai.api.OpenAiApi; import org.springframework.boot.SpringBootConfiguration; -import org.springframework.boot.autoconfigure.AutoConfigurations; import org.springframework.boot.autoconfigure.EnableAutoConfiguration; import org.springframework.boot.test.context.runner.ApplicationContextRunner; import org.springframework.context.annotation.Bean; @@ -44,7 +42,6 @@ /** * @author Chris Smith */ - @Testcontainers class MongoDBAtlasVectorStoreIT { @@ -53,13 +50,12 @@ class MongoDBAtlasVectorStoreIT { private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() .withUserConfiguration(TestApplication.class) - .withPropertyValues("spring.ai.openai.apiKey=" + System.getenv("OPENAI_API_KEY"), - String.format("spring.data.mongodb.database=" + "springaisample"), + .withPropertyValues("spring.data.mongodb.database=springaisample", String.format("spring.data.mongodb.uri=" + container.getConnectionString())); @BeforeEach public void beforeEach() { - contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> { + contextRunner.run(context -> { MongoTemplate mongoTemplate = context.getBean(MongoTemplate.class); mongoTemplate.getCollection("vector_store").deleteMany(new org.bson.Document()); }); @@ -67,7 +63,7 @@ public void beforeEach() { @Test void vectorStoreTest() { - contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> { + contextRunner.run(context -> { VectorStore vectorStore = context.getBean(VectorStore.class); List documents = List.of( @@ -102,7 +98,7 @@ void vectorStoreTest() { @Test void documentUpdateTest() { - contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> { + contextRunner.run(context -> { VectorStore vectorStore = context.getBean(VectorStore.class); Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!", @@ -137,7 +133,7 @@ void documentUpdateTest() { @Test void searchWithFilters() { - contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> { + contextRunner.run(context -> { VectorStore vectorStore = context.getBean(VectorStore.class); var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", From 7fe26e129c64ea9f7f95afa2069b8ba306aa2efb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Edd=C3=BA=20Mel=C3=A9ndez?= Date: Sun, 17 Mar 2024 00:10:14 -0600 Subject: [PATCH 2/2] Fix IT --- ...DBAtlasVectorStoreAutoConfigurationIT.java | 37 ++++--------------- 1 file changed, 7 insertions(+), 30 deletions(-) diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreAutoConfigurationIT.java index 95477bbd466..ec92591a6bc 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreAutoConfigurationIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreAutoConfigurationIT.java @@ -16,25 +16,21 @@ package org.springframework.ai.autoconfigure.vectorstore.mongo; import org.junit.jupiter.api.Test; +import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration; +import org.springframework.ai.autoconfigure.retry.SpringAiRetryAutoConfiguration; import org.springframework.ai.document.Document; -import org.springframework.ai.embedding.EmbeddingClient; -import org.springframework.ai.transformers.TransformersEmbeddingClient; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.VectorStore; import org.springframework.boot.autoconfigure.AutoConfigurations; import org.springframework.boot.autoconfigure.data.mongo.MongoDataAutoConfiguration; import org.springframework.boot.autoconfigure.mongo.MongoAutoConfiguration; +import org.springframework.boot.autoconfigure.web.client.RestClientAutoConfiguration; 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.GenericContainer; import org.testcontainers.containers.wait.strategy.Wait; import org.testcontainers.junit.jupiter.Container; import org.testcontainers.junit.jupiter.Testcontainers; -import java.io.IOException; -import java.nio.charset.StandardCharsets; import java.time.Duration; import java.util.Collections; import java.util.List; @@ -66,12 +62,13 @@ class MongoDBAtlasVectorStoreAutoConfigurationIT { private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() .withConfiguration(AutoConfigurations.of(MongoAutoConfiguration.class, MongoDataAutoConfiguration.class, - MongoDBAtlasVectorStoreAutoConfiguration.class)) - .withUserConfiguration(Config.class) + MongoDBAtlasVectorStoreAutoConfiguration.class, RestClientAutoConfiguration.class, + SpringAiRetryAutoConfiguration.class, OpenAiAutoConfiguration.class)) .withPropertyValues("spring.data.mongodb.database=springaisample", "spring.ai.vectorstore.mongodb.collection-name=test_collection", - "spring.ai.vectorstore.mongodb.path-name=test_path", + "spring.ai.vectorstore.mongodb.path-name=testembedding", "spring.ai.vectorstore.mongodb.index-name=text_index", + "spring.ai.openai.api-key=" + System.getenv("OPENAI_API_KEY"), String.format( "spring.data.mongodb.uri=" + String.format("mongodb://root:root@%s:%s/?directConnection=true", mongo.getHost(), mongo.getMappedPort(27778)))); @@ -102,24 +99,4 @@ public void addAndSearch() { }); } - public static 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 EmbeddingClient embeddingClient() { - return new TransformersEmbeddingClient(); - } - - } - } \ No newline at end of file