From a1afb0db4644f216b3fead6142357fa468b7ca19 Mon Sep 17 00:00:00 2001
From: Mark Pollack <mpollack@vmware.com>
Date: Fri, 5 Jul 2024 16:26:02 -0400
Subject: [PATCH] Fix Mongo VectorStore persistance issue

Fixes #698
---
 ...goDBAtlasVectorStoreAutoConfiguration.java | 41 +++++++++++++++++++
 ...DBAtlasVectorStoreAutoConfigurationIT.java |  4 +-
 2 files changed, 42 insertions(+), 3 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 dfbebca77..9d5bd857a 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
@@ -17,15 +17,22 @@ package org.springframework.ai.autoconfigure.vectorstore.mongo;
 
 import org.springframework.ai.embedding.EmbeddingModel;
 import org.springframework.ai.vectorstore.MongoDBAtlasVectorStore;
+import org.springframework.beans.BeansException;
+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;
 import org.springframework.data.mongodb.core.MongoTemplate;
+import org.springframework.data.mongodb.core.convert.MongoCustomConversions;
+import org.springframework.util.MimeType;
 import org.springframework.util.StringUtils;
 
+import java.util.Arrays;
+
 /**
  * @author Eddú Meléndez
  * @author Christian Tzolov
@@ -57,4 +64,38 @@ public class MongoDBAtlasVectorStoreAutoConfiguration {
 		return new MongoDBAtlasVectorStore(mongoTemplate, embeddingModel, config, properties.isInitializeSchema());
 	}
 
+	@Bean
+	public Converter<MimeType, String> mimeTypeToStringConverter() {
+		return new Converter<MimeType, String>() {
+			@Override
+			public String convert(MimeType source) {
+				return source.toString();
+			}
+		};
+	}
+
+	@Bean
+	public Converter<String, MimeType> stringToMimeTypeConverter() {
+		return new Converter<String, MimeType>() {
+			@Override
+			public MimeType convert(String source) {
+				return MimeType.valueOf(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;
+			}
+		};
+	}
+
 }
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 dfcae56b6..efef61d8c 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
@@ -15,7 +15,6 @@
  */
 package org.springframework.ai.autoconfigure.vectorstore.mongo;
 
-import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration;
 import org.springframework.ai.autoconfigure.retry.SpringAiRetryAutoConfiguration;
@@ -43,11 +42,10 @@ import static org.assertj.core.api.Assertions.assertThat;
  * @author Eddú Meléndez
  */
 @Testcontainers
-@Disabled("Disabled due to https://github.com/spring-projects/spring-ai/issues/698")
 class MongoDBAtlasVectorStoreAutoConfigurationIT {
 
 	@Container
-	static GenericContainer<?> mongo = new GenericContainer<>("mongodb/atlas:v1.15.1").withPrivilegedMode(true)
+	static GenericContainer<?> mongo = new GenericContainer<>("mongodb/atlas:v1.24.0").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)