From 3bd1fae1d160926cf569e44904f45b00bfb4eb0c Mon Sep 17 00:00:00 2001 From: Soby Chacko Date: Tue, 20 Aug 2024 17:37:10 -0400 Subject: [PATCH] Fixing MongoDB vector store support test failures Remove the requirement of MongoDB vector store auto-configuration only after `MongoDataAutoConfiguration`. This enables the `MongoDBAtlasVectorStoreAutoConfiguration` to properly provide custom Mongo conversions. In `MongoDBAtlasVectorStoreIT` and `MongoDbVectorStoreObservationIT` tests, properly provision the `MongoTemplate` with custom conversions as these tests are not relying on auto-configuration. --- ...goDBAtlasVectorStoreAutoConfiguration.java | 21 ++------ ...DBAtlasVectorStoreAutoConfigurationIT.java | 3 +- .../MongoDBAtlasVectorStoreIT.java | 48 +++++++++++++++++-- .../MongoDbVectorStoreObservationIT.java | 48 +++++++++++++++++-- 4 files changed, 92 insertions(+), 28 deletions(-) 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 index 70e8f96070a..bdae34efb40 100644 --- 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 @@ -1,5 +1,5 @@ /* - * Copyright 2023 - 2024 the original author or authors. + * 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. @@ -18,13 +18,10 @@ import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.vectorstore.MongoDBAtlasVectorStore; import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention; -import org.springframework.beans.BeansException; import org.springframework.beans.factory.ObjectProvider; -import org.springframework.beans.factory.config.BeanPostProcessor; 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.core.convert.converter.Converter; @@ -40,9 +37,10 @@ /** * @author Eddú Meléndez * @author Christian Tzolov + * @author Soby Chacko * @since 1.0.0 */ -@AutoConfiguration(after = MongoDataAutoConfiguration.class) +@AutoConfiguration @ConditionalOnClass({ MongoDBAtlasVectorStore.class, EmbeddingModel.class, MongoTemplate.class }) @EnableConfigurationProperties(MongoDBAtlasVectorStoreProperties.class) public class MongoDBAtlasVectorStoreAutoConfiguration { @@ -92,17 +90,8 @@ public MimeType convert(String source) { } @Bean - public BeanPostProcessor mongoCustomConversionsPostProcessor() { - return new BeanPostProcessor() { - @Override - public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException { - if (bean instanceof MongoCustomConversions) { - return new MongoCustomConversions( - Arrays.asList(mimeTypeToStringConverter(), stringToMimeTypeConverter())); - } - return bean; - } - }; + public MongoCustomConversions mongoCustomConversions() { + return new MongoCustomConversions(Arrays.asList(mimeTypeToStringConverter(), stringToMimeTypeConverter())); } } 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 4d1cc4c4205..47d0a70d890 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 @@ -1,5 +1,5 @@ /* - * Copyright 2023 - 2024 the original author or authors. + * 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. @@ -53,7 +53,6 @@ */ @Testcontainers @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") -@Disabled("Disabled due to https://github.com/spring-projects/spring-ai/issues/698") class MongoDBAtlasVectorStoreAutoConfigurationIT { @Container 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 e804ff4fdb9..9215e259675 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 @@ -1,5 +1,5 @@ /* - * Copyright 2023 - 2024 the original author or authors. + * 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. @@ -17,7 +17,6 @@ import com.mongodb.client.MongoClient; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import org.springframework.ai.document.Document; @@ -28,10 +27,17 @@ import org.springframework.boot.autoconfigure.EnableAutoConfiguration; import org.springframework.boot.test.context.runner.ApplicationContextRunner; import org.springframework.context.annotation.Bean; +import org.springframework.core.convert.converter.Converter; import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.convert.MappingMongoConverter; +import org.springframework.data.mongodb.core.convert.MongoCustomConversions; +import org.springframework.data.mongodb.core.mapping.MongoMappingContext; +import org.springframework.util.MimeType; + import org.testcontainers.junit.jupiter.Container; import org.testcontainers.junit.jupiter.Testcontainers; +import java.util.Arrays; import java.util.Collections; import java.util.List; import java.util.Map; @@ -42,9 +48,9 @@ /** * @author Chris Smith + * @author Soby Chacko */ @Testcontainers -@Disabled("Disabled due to https://github.com/spring-projects/spring-ai/issues/698") class MongoDBAtlasVectorStoreIT { @Container @@ -201,8 +207,14 @@ public VectorStore vectorStore(MongoTemplate mongoTemplate, EmbeddingModel embed } @Bean - public MongoTemplate mongoTemplate(MongoClient mongoClient) { - return new MongoTemplate(mongoClient, "springaisample"); + public MongoTemplate mongoTemplate(MongoClient mongoClient, MongoCustomConversions mongoCustomConversions) { + MongoTemplate mongoTemplate = new MongoTemplate(mongoClient, "springaisample"); + MappingMongoConverter converter = (MappingMongoConverter) mongoTemplate.getConverter(); + converter.setCustomConversions(mongoCustomConversions); + ((MongoMappingContext) converter.getMappingContext()) + .setSimpleTypeHolder(mongoCustomConversions.getSimpleTypeHolder()); + converter.afterPropertiesSet(); + return mongoTemplate; } @Bean @@ -210,6 +222,32 @@ public EmbeddingModel embeddingModel() { return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY"))); } + @Bean + public Converter mimeTypeToStringConverter() { + return new Converter() { + @Override + public String convert(MimeType source) { + return source.toString(); + } + }; + } + + @Bean + public Converter stringToMimeTypeConverter() { + return new Converter() { + @Override + public MimeType convert(String source) { + return MimeType.valueOf(source); + } + }; + } + + @Bean + public MongoCustomConversions mongoCustomConversions(Converter mimeTypeToStringConverter, + Converter stringToMimeTypeConverter) { + return new MongoCustomConversions(Arrays.asList(mimeTypeToStringConverter, stringToMimeTypeConverter)); + } + } } \ No newline at end of file diff --git a/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDbVectorStoreObservationIT.java b/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDbVectorStoreObservationIT.java index 08577f6532f..37423ba0a2d 100644 --- a/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDbVectorStoreObservationIT.java +++ b/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDbVectorStoreObservationIT.java @@ -19,12 +19,11 @@ import java.io.IOException; import java.nio.charset.StandardCharsets; +import java.util.Arrays; import java.util.List; import java.util.Map; -import java.util.Vector; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; import org.springframework.ai.document.Document; @@ -39,8 +38,14 @@ import org.springframework.boot.autoconfigure.EnableAutoConfiguration; import org.springframework.boot.test.context.runner.ApplicationContextRunner; import org.springframework.context.annotation.Bean; +import org.springframework.core.convert.converter.Converter; import org.springframework.core.io.DefaultResourceLoader; import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.convert.MappingMongoConverter; +import org.springframework.data.mongodb.core.convert.MongoCustomConversions; +import org.springframework.data.mongodb.core.mapping.MongoMappingContext; +import org.springframework.util.MimeType; + import org.testcontainers.junit.jupiter.Container; import org.testcontainers.junit.jupiter.Testcontainers; @@ -52,10 +57,10 @@ /** * @author Christian Tzolov + * @author Soby Chacko */ @Testcontainers @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") -@Disabled("Disabled due to https://github.com/spring-projects/spring-ai/issues/698") public class MongoDbVectorStoreObservationIT { @Container @@ -100,6 +105,8 @@ void observationVectorStoreAddAndQueryOperations() { vectorStore.add(documents); + Thread.sleep(5000); + TestObservationRegistryAssert.assertThat(observationRegistry) .doesNotHaveAnyRemainingCurrentObservation() .hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME) @@ -175,8 +182,14 @@ public VectorStore vectorStore(MongoTemplate mongoTemplate, EmbeddingModel embed } @Bean - public MongoTemplate mongoTemplate(MongoClient mongoClient) { - return new MongoTemplate(mongoClient, "springaisample"); + public MongoTemplate mongoTemplate(MongoClient mongoClient, MongoCustomConversions mongoCustomConversions) { + MongoTemplate mongoTemplate = new MongoTemplate(mongoClient, "springaisample"); + MappingMongoConverter converter = (MappingMongoConverter) mongoTemplate.getConverter(); + converter.setCustomConversions(mongoCustomConversions); + ((MongoMappingContext) converter.getMappingContext()) + .setSimpleTypeHolder(mongoCustomConversions.getSimpleTypeHolder()); + converter.afterPropertiesSet(); + return mongoTemplate; } @Bean @@ -184,6 +197,31 @@ public EmbeddingModel embeddingModel() { return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY"))); } + @Bean + public Converter mimeTypeToStringConverter() { + return new Converter() { + @Override + public String convert(MimeType source) { + return source.toString(); + } + }; + } + + @Bean + public Converter stringToMimeTypeConverter() { + return new Converter() { + @Override + public MimeType convert(String source) { + return MimeType.valueOf(source); + } + }; + } + + @Bean + public MongoCustomConversions mongoCustomConversions() { + return new MongoCustomConversions(Arrays.asList(mimeTypeToStringConverter(), stringToMimeTypeConverter())); + } + } }