Add MongoDB as a vector store

pom.xml

@@ -7,7 +7,6 @@
 
 <packaging>pom</packaging>
 <url>https://github.com/spring-projects-experimental/spring-ai</url>
-
 <name>Spring AI</name>
 <description>Building AI applications with Spring Boot</description>
 
@@ -20,6 +19,7 @@
 <module>spring-ai-spring-boot-starters/spring-ai-starter-azure-openai</module>
 <module>spring-ai-docs</module>
 <module>vector-stores/spring-ai-pgvector-store</module>
+<module>vector-stores/spring-ai-mongodb-store</module>
 </modules>
 
  <organization>

spring-ai-core/src/main/java/org/springframework/ai/vectorstore/impl/InMemoryVectorStore.java

@@ -95,6 +95,14 @@ public Similarity(String key, double similarity) {
 this.similarity = similarity;
 }
 
+public double getSimilarity() {
+return similarity;
+}
+
+public String getKey() {
+return key;
+}
+
 }
 
 public class EmbeddingMath {

vector-stores/spring-ai-mongodb-store/README.md

@@ -0,0 +1,4 @@
+
+# MongoDB VectorStore
+
+This use mongo db as a vector store using the same math as the in-memory vector store.

vector-stores/spring-ai-mongodb-store/pom.xml

@@ -0,0 +1,74 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
+<modelVersion>4.0.0</modelVersion>
+<parent>
+<groupId>org.springframework.experimental.ai</groupId>
+<artifactId>spring-ai</artifactId>
+<version>0.2.0-SNAPSHOT</version>
+<relativePath>../../pom.xml</relativePath>
+</parent>
+<artifactId>spring-ai-mongodb-store</artifactId>
+<packaging>jar</packaging>
+<name>Spring AI Vector Store - mongodb</name>
+<description>Spring AI MongoDB Vector Store</description>
+<url>https://github.com/spring-projects-experimental/spring-ai</url>
+
+<scm>
+<url>https://github.com/spring-projects-experimental/spring-ai</url>
+<connection>git://github.com/spring-projects-experimental/spring-ai.git</connection>
+<developerConnection>git@github.com:spring-projects-experimental/spring-ai.git</developerConnection>
+</scm>
+
+<properties>
+<spring-ai.version>0.2.0-SNAPSHOT</spring-ai.version>
+<!-- testing -->
+<testcontainers.version>1.19.0</testcontainers.version>
+</properties>
+
+<dependencies>
+<dependency>
+<groupId>org.springframework.experimental.ai</groupId>
+<artifactId>spring-ai-core</artifactId>
+<version>${spring-ai.version}</version>
+</dependency>
+
+<dependency>
+<groupId>org.springframework.data</groupId>
+<artifactId>spring-data-mongodb</artifactId>
+</dependency>
+
+<dependency>
+<groupId>org.mongodb</groupId>
+<artifactId>mongodb-driver-sync</artifactId>
+</dependency>
+<!-- TESTING -->
+<dependency>
+<groupId>org.springframework.experimental.ai</groupId>
+<artifactId>spring-ai-openai-spring-boot-starter</artifactId>
+<version>${spring-ai.version}</version>
+<scope>test</scope>
+</dependency>
+
+<dependency>
+<groupId>org.springframework.boot</groupId>
+<artifactId>spring-boot-starter-test</artifactId>
+<scope>test</scope>
+</dependency>
+
+<dependency>
+<groupId>org.testcontainers</groupId>
+<artifactId>testcontainers</artifactId>
+<version>${testcontainers.version}</version>
+<scope>test</scope>
+</dependency>
+
+<dependency>
+<groupId>org.testcontainers</groupId>
+<artifactId>junit-jupiter</artifactId>
+<version>${testcontainers.version}</version>
+<scope>test</scope>
+</dependency>
+</dependencies>
+
+</project>

vector-stores/spring-ai-mongodb-store/src/main/java/org/springframework/ai/vectorstore/MongoDBVectorStore.java

@@ -0,0 +1,128 @@
+/*
+ * Copyright 2023-2023 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.vectorstore;
+
+import com.mongodb.BasicDBObject;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.embedding.EmbeddingClient;
+import org.springframework.ai.vectorstore.impl.InMemoryVectorStore;
+import org.springframework.data.domain.Sort;
+import org.springframework.data.mongodb.core.MongoTemplate;
+import org.springframework.data.mongodb.core.aggregation.Aggregation;
+import org.springframework.data.mongodb.core.query.Query;
+
+import java.util.Comparator;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.springframework.data.mongodb.core.query.Criteria.where;
+
+/**
+ * @author Chris Smith
+ */
+public class MongoDBVectorStore implements VectorStore {
+
+private MongoTemplate mongoTemplate;
+
+private EmbeddingClient embeddingClient;
+
+private static final String VECTOR_COLLECTION_NAME = "SampleCollection";
+
+public MongoDBVectorStore(MongoTemplate mongoTemplate, EmbeddingClient embeddingClient) {
+this.mongoTemplate = mongoTemplate;
+this.embeddingClient = embeddingClient;
+if (!mongoTemplate.collectionExists(VECTOR_COLLECTION_NAME)) {
+mongoTemplate.createCollection(VECTOR_COLLECTION_NAME);
+}
+}
+
+/**
+ * Maps a basicDBObject to a Spring AI Document
+ * @param basicDBObject
+ * @return
+ */
+private Document mapBasicDbObject(BasicDBObject basicDBObject) {
+String id = basicDBObject.getString("_id");
+String content = basicDBObject.getString("text");
+Map<String, Object> metadata = (Map<String, Object>) basicDBObject.get("metadata");
+List<Double> embedding = (List<Double>) basicDBObject.get("embedding");
+
+Document document = new Document(id, content, metadata);
+document.setEmbedding(embedding);
+
+return document;
+}
+
+@Override
+public void add(List<Document> documents) {
+for (Document document : documents) {
+List<Double> embedding = this.embeddingClient.embed(document);
+document.setEmbedding(embedding);
+
+this.mongoTemplate.save(document, VECTOR_COLLECTION_NAME);
+}
+}
+
+@Override
+public Optional<Boolean> delete(List<String> idList) {
+Query query = new Query(where("_id").in(idList));
+
+var deleteRes = this.mongoTemplate.remove(query, VECTOR_COLLECTION_NAME);
+long deleteCount = deleteRes.getDeletedCount();
+
+return Optional.of(deleteCount == idList.size());
+}
+
+@Override
+public List<Document> similaritySearch(String query) {
+return this.similaritySearch(query, 4);
+}
+
+@Override
+public List<Document> similaritySearch(String query, int k) {
+List<Double> queryEmbedding = this.embeddingClient.embed(query);
+
+//Build aggregation to leverage searching through mongodb
+Aggregation aggregation = Aggregation.newAggregation(
+new VectorSearchAggregation(k, queryEmbedding), Aggregation.sort(Sort.Direction.DESC, "score"));
+return this.mongoTemplate.aggregate(aggregation, VECTOR_COLLECTION_NAME,BasicDBObject.class)
+.getMappedResults().stream()
+.map(this::mapBasicDbObject)
+.toList();
+}
+
+@Override
+public List<Document> similaritySearch(String query, int k, double threshold) {
+List<Double> queryEmbedding = this.embeddingClient.embed(query);
+
+return this.mongoTemplate.findAll(BasicDBObject.class, VECTOR_COLLECTION_NAME)
+.stream()
+.map(this::mapBasicDbObject)
+.map(entry -> new InMemoryVectorStore.Similarity(entry.getId(),
+InMemoryVectorStore.EmbeddingMath.cosineSimilarity(queryEmbedding, entry.getEmbedding())))
+.filter(s -> s.getSimilarity() >= threshold)
+.sorted(Comparator.<InMemoryVectorStore.Similarity>comparingDouble(s -> s.getSimilarity()).reversed())
+.limit(k)
+.map(s -> this.mongoTemplate.findById(s.getKey(), BasicDBObject.class, VECTOR_COLLECTION_NAME))
+.map(this::mapBasicDbObject)
+.toList();
+}
+
+
+
+}

vector-stores/spring-ai-mongodb-store/src/main/java/org/springframework/ai/vectorstore/VectorSearchAggregation.java

@@ -0,0 +1,30 @@
+package org.springframework.ai.vectorstore;
+
+import org.bson.Document;
+import org.springframework.data.mongodb.core.aggregation.AggregationOperation;
+import org.springframework.data.mongodb.core.aggregation.AggregationOperationContext;
+
+import java.util.List;
+
+class VectorSearchAggregation implements AggregationOperation {
+
+ private final int count;
+ private final List<Double> embeddings;
+
+
+
+ public VectorSearchAggregation(int count, List<Double> embeddings){
+ this.count = count;
+ this.embeddings = embeddings;
+ }
+ @Override
+ public org.bson.Document toDocument(AggregationOperationContext context) {
+ var doc = new Document("$search",
+ new Document("index", "default")
+ .append("path", "embedding")
+ .append("numCandidates", 100)
+ .append("limit",count)
+ .append("queryVector", embeddings));
+ return context.getMappedObject(doc);
+ }
+}