Skip to content

Commit f0771f0

Browse files
sobychackotzolov
sobychacko
authored and
tzolov
committed
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.
1 parent dad30f0 commit f0771f0

File tree

4 files changed

+92
-28
lines changed

4 files changed

+92
-28
lines changed

‎spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreAutoConfiguration.java

Lines changed: 5 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023 - 2024 the original author or authors.
2+
* Copyright 2023-2024 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -18,13 +18,10 @@
1818
import org.springframework.ai.embedding.EmbeddingModel;
1919
import org.springframework.ai.vectorstore.MongoDBAtlasVectorStore;
2020
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
21-
import org.springframework.beans.BeansException;
2221
import org.springframework.beans.factory.ObjectProvider;
23-
import org.springframework.beans.factory.config.BeanPostProcessor;
2422
import org.springframework.boot.autoconfigure.AutoConfiguration;
2523
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
2624
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
27-
import org.springframework.boot.autoconfigure.data.mongo.MongoDataAutoConfiguration;
2825
import org.springframework.boot.context.properties.EnableConfigurationProperties;
2926
import org.springframework.context.annotation.Bean;
3027
import org.springframework.core.convert.converter.Converter;
@@ -40,9 +37,10 @@
4037
/**
4138
* @author Eddú Meléndez
4239
* @author Christian Tzolov
40+
* @author Soby Chacko
4341
* @since 1.0.0
4442
*/
45-
@AutoConfiguration(after = MongoDataAutoConfiguration.class)
43+
@AutoConfiguration
4644
@ConditionalOnClass({ MongoDBAtlasVectorStore.class, EmbeddingModel.class, MongoTemplate.class })
4745
@EnableConfigurationProperties(MongoDBAtlasVectorStoreProperties.class)
4846
public class MongoDBAtlasVectorStoreAutoConfiguration {
@@ -92,17 +90,8 @@ public MimeType convert(String source) {
9290
}
9391

9492
@Bean
95-
public BeanPostProcessor mongoCustomConversionsPostProcessor() {
96-
return new BeanPostProcessor() {
97-
@Override
98-
public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
99-
if (bean instanceof MongoCustomConversions) {
100-
return new MongoCustomConversions(
101-
Arrays.asList(mimeTypeToStringConverter(), stringToMimeTypeConverter()));
102-
}
103-
return bean;
104-
}
105-
};
93+
public MongoCustomConversions mongoCustomConversions() {
94+
return new MongoCustomConversions(Arrays.asList(mimeTypeToStringConverter(), stringToMimeTypeConverter()));
10695
}
10796

10897
}

‎spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/mongo/MongoDBAtlasVectorStoreAutoConfigurationIT.java

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023 - 2024 the original author or authors.
2+
* Copyright 2023-2024 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -53,7 +53,6 @@
5353
*/
5454
@Testcontainers
5555
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
56-
@Disabled("Disabled due to https://github.com/spring-projects/spring-ai/issues/698")
5756
class MongoDBAtlasVectorStoreAutoConfigurationIT {
5857

5958
@Container

‎vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStoreIT.java

Lines changed: 43 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023 - 2024 the original author or authors.
2+
* Copyright 2023-2024 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -17,7 +17,6 @@
1717

1818
import com.mongodb.client.MongoClient;
1919
import org.junit.jupiter.api.BeforeEach;
20-
import org.junit.jupiter.api.Disabled;
2120
import org.junit.jupiter.api.Test;
2221

2322
import org.springframework.ai.document.Document;
@@ -28,10 +27,17 @@
2827
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
2928
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
3029
import org.springframework.context.annotation.Bean;
30+
import org.springframework.core.convert.converter.Converter;
3131
import org.springframework.data.mongodb.core.MongoTemplate;
32+
import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
33+
import org.springframework.data.mongodb.core.convert.MongoCustomConversions;
34+
import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
35+
import org.springframework.util.MimeType;
36+
3237
import org.testcontainers.junit.jupiter.Container;
3338
import org.testcontainers.junit.jupiter.Testcontainers;
3439

40+
import java.util.Arrays;
3541
import java.util.Collections;
3642
import java.util.List;
3743
import java.util.Map;
@@ -42,9 +48,9 @@
4248

4349
/**
4450
* @author Chris Smith
51+
* @author Soby Chacko
4552
*/
4653
@Testcontainers
47-
@Disabled("Disabled due to https://github.com/spring-projects/spring-ai/issues/698")
4854
class MongoDBAtlasVectorStoreIT {
4955

5056
@Container
@@ -201,15 +207,47 @@ public VectorStore vectorStore(MongoTemplate mongoTemplate, EmbeddingModel embed
201207
}
202208

203209
@Bean
204-
public MongoTemplate mongoTemplate(MongoClient mongoClient) {
205-
return new MongoTemplate(mongoClient, "springaisample");
210+
public MongoTemplate mongoTemplate(MongoClient mongoClient, MongoCustomConversions mongoCustomConversions) {
211+
MongoTemplate mongoTemplate = new MongoTemplate(mongoClient, "springaisample");
212+
MappingMongoConverter converter = (MappingMongoConverter) mongoTemplate.getConverter();
213+
converter.setCustomConversions(mongoCustomConversions);
214+
((MongoMappingContext) converter.getMappingContext())
215+
.setSimpleTypeHolder(mongoCustomConversions.getSimpleTypeHolder());
216+
converter.afterPropertiesSet();
217+
return mongoTemplate;
206218
}
207219

208220
@Bean
209221
public EmbeddingModel embeddingModel() {
210222
return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY")));
211223
}
212224

225+
@Bean
226+
public Converter<MimeType, String> mimeTypeToStringConverter() {
227+
return new Converter<MimeType, String>() {
228+
@Override
229+
public String convert(MimeType source) {
230+
return source.toString();
231+
}
232+
};
233+
}
234+
235+
@Bean
236+
public Converter<String, MimeType> stringToMimeTypeConverter() {
237+
return new Converter<String, MimeType>() {
238+
@Override
239+
public MimeType convert(String source) {
240+
return MimeType.valueOf(source);
241+
}
242+
};
243+
}
244+
245+
@Bean
246+
public MongoCustomConversions mongoCustomConversions(Converter<MimeType, String> mimeTypeToStringConverter,
247+
Converter<String, MimeType> stringToMimeTypeConverter) {
248+
return new MongoCustomConversions(Arrays.asList(mimeTypeToStringConverter, stringToMimeTypeConverter));
249+
}
250+
213251
}
214252

215253
}

‎vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDbVectorStoreObservationIT.java

Lines changed: 43 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,11 @@
1919

2020
import java.io.IOException;
2121
import java.nio.charset.StandardCharsets;
22+
import java.util.Arrays;
2223
import java.util.List;
2324
import java.util.Map;
24-
import java.util.Vector;
2525

2626
import org.junit.jupiter.api.BeforeEach;
27-
import org.junit.jupiter.api.Disabled;
2827
import org.junit.jupiter.api.Test;
2928
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
3029
import org.springframework.ai.document.Document;
@@ -39,8 +38,14 @@
3938
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
4039
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
4140
import org.springframework.context.annotation.Bean;
41+
import org.springframework.core.convert.converter.Converter;
4242
import org.springframework.core.io.DefaultResourceLoader;
4343
import org.springframework.data.mongodb.core.MongoTemplate;
44+
import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
45+
import org.springframework.data.mongodb.core.convert.MongoCustomConversions;
46+
import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
47+
import org.springframework.util.MimeType;
48+
4449
import org.testcontainers.junit.jupiter.Container;
4550
import org.testcontainers.junit.jupiter.Testcontainers;
4651

@@ -52,10 +57,10 @@
5257

5358
/**
5459
* @author Christian Tzolov
60+
* @author Soby Chacko
5561
*/
5662
@Testcontainers
5763
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
58-
@Disabled("Disabled due to https://github.com/spring-projects/spring-ai/issues/698")
5964
public class MongoDbVectorStoreObservationIT {
6065

6166
@Container
@@ -100,6 +105,8 @@ void observationVectorStoreAddAndQueryOperations() {
100105

101106
vectorStore.add(documents);
102107

108+
Thread.sleep(5000);
109+
103110
TestObservationRegistryAssert.assertThat(observationRegistry)
104111
.doesNotHaveAnyRemainingCurrentObservation()
105112
.hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME)
@@ -175,15 +182,46 @@ public VectorStore vectorStore(MongoTemplate mongoTemplate, EmbeddingModel embed
175182
}
176183

177184
@Bean
178-
public MongoTemplate mongoTemplate(MongoClient mongoClient) {
179-
return new MongoTemplate(mongoClient, "springaisample");
185+
public MongoTemplate mongoTemplate(MongoClient mongoClient, MongoCustomConversions mongoCustomConversions) {
186+
MongoTemplate mongoTemplate = new MongoTemplate(mongoClient, "springaisample");
187+
MappingMongoConverter converter = (MappingMongoConverter) mongoTemplate.getConverter();
188+
converter.setCustomConversions(mongoCustomConversions);
189+
((MongoMappingContext) converter.getMappingContext())
190+
.setSimpleTypeHolder(mongoCustomConversions.getSimpleTypeHolder());
191+
converter.afterPropertiesSet();
192+
return mongoTemplate;
180193
}
181194

182195
@Bean
183196
public EmbeddingModel embeddingModel() {
184197
return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY")));
185198
}
186199

200+
@Bean
201+
public Converter<MimeType, String> mimeTypeToStringConverter() {
202+
return new Converter<MimeType, String>() {
203+
@Override
204+
public String convert(MimeType source) {
205+
return source.toString();
206+
}
207+
};
208+
}
209+
210+
@Bean
211+
public Converter<String, MimeType> stringToMimeTypeConverter() {
212+
return new Converter<String, MimeType>() {
213+
@Override
214+
public MimeType convert(String source) {
215+
return MimeType.valueOf(source);
216+
}
217+
};
218+
}
219+
220+
@Bean
221+
public MongoCustomConversions mongoCustomConversions() {
222+
return new MongoCustomConversions(Arrays.asList(mimeTypeToStringConverter(), stringToMimeTypeConverter()));
223+
}
224+
187225
}
188226

189227
}

0 commit comments

Comments
 (0)