diff --git a/pom.xml b/pom.xml index 83bd0e1ff..5b73838a2 100644 --- a/pom.xml +++ b/pom.xml @@ -60,6 +60,7 @@ spring-ai-spring-boot-starters/spring-ai-starter-bedrock-ai spring-ai-spring-boot-starters/spring-ai-starter-mistral-ai spring-ai-retry + vector-stores/spring-ai-mongodb-atlas-store diff --git a/spring-ai-bom/pom.xml b/spring-ai-bom/pom.xml index 03655b644..26ab0d414 100644 --- a/spring-ai-bom/pom.xml +++ b/spring-ai-bom/pom.xml @@ -174,6 +174,12 @@ ${project.version} + + org.springframework.ai + spring-ai-mongodb-atlas-store + ${project.version} + + org.springframework.ai diff --git a/vector-stores/spring-ai-mongodb-atlas-store/pom.xml b/vector-stores/spring-ai-mongodb-atlas-store/pom.xml new file mode 100644 index 000000000..691bebece --- /dev/null +++ b/vector-stores/spring-ai-mongodb-atlas-store/pom.xml @@ -0,0 +1,59 @@ + + + 4.0.0 + + org.springframework.ai + spring-ai + 1.0.0-SNAPSHOT + ../../pom.xml + + spring-ai-mongodb-atlas-store + jar + Spring AI MongoDB Atlas Vector Store + Spring AI MongoDB Atlas Vector Store + https://github.com/spring-projects-experimental/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.ai + spring-ai-core + ${parent.version} + + + + org.springframework.data + spring-data-mongodb + + + org.mongodb + mongodb-driver-sync + + + + + org.springframework.ai + spring-ai-openai-spring-boot-starter + ${parent.version} + test + + + + org.springframework.boot + spring-boot-starter-test + test + + + org.testcontainers + junit-jupiter + ${testcontainers.version} + test + + + diff --git a/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/MongoDBAtlasFilterExpressionConverter.java b/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/MongoDBAtlasFilterExpressionConverter.java new file mode 100644 index 000000000..f541b0002 --- /dev/null +++ b/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/MongoDBAtlasFilterExpressionConverter.java @@ -0,0 +1,97 @@ +/* + * 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.vectorstore; + +import org.springframework.ai.vectorstore.filter.Filter; +import org.springframework.ai.vectorstore.filter.converter.AbstractFilterExpressionConverter; + +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.AND; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.OR; + +/** + * Converts {@link Filter.Expression} into MongDB Atlas metadata filter expression format. + * (https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-stage/#std-label-vectorSearch-agg-pipeline-filter) + * + * @author Chris Smith + * @since 1.0.0 + */ +public class MongoDBAtlasFilterExpressionConverter extends AbstractFilterExpressionConverter { + + @Override + protected void doExpression(Filter.Expression expression, StringBuilder context) { + // Handling AND/OR + if (AND.equals(expression.type()) || OR.equals(expression.type())) { + doCompoundExpressionType(expression, context); + } + else { + doSingleExpressionType(expression, context); + } + } + + private void doCompoundExpressionType(Filter.Expression expression, StringBuilder context) { + context.append("{"); + context.append(getOperationSymbol(expression)); + context.append(":["); + this.convertOperand(expression.left(), context); + context.append(","); + this.convertOperand(expression.right(), context); + context.append("]}"); + } + + private void doSingleExpressionType(Filter.Expression expression, StringBuilder context) { + context.append("{"); + this.convertOperand(expression.left(), context); + context.append(":{"); + context.append(getOperationSymbol(expression)); + context.append(":"); + this.convertOperand(expression.right(), context); + context.append("}}"); + } + + private String getOperationSymbol(Filter.Expression exp) { + switch (exp.type()) { + case AND: + return "$and"; + case OR: + return "$or"; + case EQ: + return "$eq"; + case NE: + return "$ne"; + case LT: + return "$lt"; + case LTE: + return "$lte"; + case GT: + return "$gt"; + case GTE: + return "$gte"; + case IN: + return "$in"; + case NIN: + return "$nin"; + default: + throw new RuntimeException("Not supported expression type:" + exp.type()); + } + } + + @Override + protected void doKey(Filter.Key filterKey, StringBuilder context) { + var identifier = (hasOuterQuotes(filterKey.key())) ? removeOuterQuotes(filterKey.key()) : filterKey.key(); + context.append("\"metadata." + identifier + "\""); + } + +} diff --git a/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStore.java b/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStore.java new file mode 100644 index 000000000..928cbb140 --- /dev/null +++ b/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStore.java @@ -0,0 +1,271 @@ +/* + * 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.vectorstore; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import com.mongodb.BasicDBObject; + +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingClient; +import org.springframework.beans.factory.InitializingBean; +import org.springframework.data.mongodb.core.MongoTemplate; +import org.springframework.data.mongodb.core.aggregation.Aggregation; +import org.springframework.data.mongodb.core.query.Criteria; +import org.springframework.data.mongodb.core.query.Query; +import org.springframework.util.Assert; + +import static org.springframework.data.mongodb.core.query.Criteria.where; + +/** + * @author Chris Smith + * @since 1.0.0 + */ +public class MongoDBAtlasVectorStore implements VectorStore, InitializingBean { + + public static final String ID_FIELD_NAME = "_id"; + + public static final String METADATA_FIELD_NAME = "metadata"; + + public static final String CONTENT_FIELD_NAME = "content"; + + public static final String SCORE_FIELD_NAME = "score"; + + private static final String DEFAULT_VECTOR_COLLECTION_NAME = "vector_store"; + + private static final String DEFAULT_VECTOR_INDEX_NAME = "vector_index"; + + private static final String DEFAULT_PATH_NAME = "embedding"; + + private static final int DEFAULT_NUM_CANDIDATES = 200; + + private final MongoTemplate mongoTemplate; + + private final EmbeddingClient embeddingClient; + + private final MongoDBVectorStoreConfig config; + + private final MongoDBAtlasFilterExpressionConverter filterExpressionConverter = new MongoDBAtlasFilterExpressionConverter(); + + public MongoDBAtlasVectorStore(MongoTemplate mongoTemplate, EmbeddingClient embeddingClient) { + this(mongoTemplate, embeddingClient, MongoDBVectorStoreConfig.defaultConfig()); + } + + public MongoDBAtlasVectorStore(MongoTemplate mongoTemplate, EmbeddingClient embeddingClient, + MongoDBVectorStoreConfig config) { + this.mongoTemplate = mongoTemplate; + this.embeddingClient = embeddingClient; + this.config = config; + + } + + @Override + public void afterPropertiesSet() throws Exception { + // Create the collection if it does not exist + if (!mongoTemplate.collectionExists(this.config.collectionName)) { + mongoTemplate.createCollection(this.config.collectionName); + } + // Create search index, command doesn't do anything if already existing + mongoTemplate.executeCommand(createSearchIndex()); + } + + /** + * Provides the Definition for the search index + */ + private org.bson.Document createSearchIndex() { + List vectorFields = new ArrayList<>(); + vectorFields.add(new org.bson.Document().append("type", "vector") + .append("path", this.config.pathName) + .append("numDimensions", 1536) + .append("similarity", "cosine")); + vectorFields.addAll(this.config.metadataFieldsToFilter.stream() + .map(fieldName -> new org.bson.Document().append("type", "filter").append("path", "metadata." + fieldName)) + .toList()); + + return new org.bson.Document().append("createSearchIndexes", this.config.collectionName) + .append("indexes", + List.of(new org.bson.Document().append("name", this.config.vectorIndexName) + .append("type", "vectorSearch") + .append("definition", new org.bson.Document("fields", vectorFields)))); + } + + /** + * Maps a BasicDBObject to a Spring AI Document + * @param basicDBObject the basicDBObject to map to a spring ai document + * @return the spring ai document + */ + @SuppressWarnings("unchecked") + private Document mapBasicDbObject(BasicDBObject basicDBObject) { + String id = basicDBObject.getString(ID_FIELD_NAME); + String content = basicDBObject.getString(CONTENT_FIELD_NAME); + Map metadata = (Map) basicDBObject.get(METADATA_FIELD_NAME); + List embedding = (List) basicDBObject.get(this.config.pathName); + + Document document = new Document(id, content, metadata); + document.setEmbedding(embedding); + + return document; + } + + @Override + public void add(List documents) { + for (Document document : documents) { + List embedding = this.embeddingClient.embed(document); + document.setEmbedding(embedding); + this.mongoTemplate.save(document, this.config.collectionName); + } + } + + @Override + public Optional delete(List idList) { + Query query = new Query(where(ID_FIELD_NAME).in(idList)); + + var deleteRes = this.mongoTemplate.remove(query, this.config.collectionName); + long deleteCount = deleteRes.getDeletedCount(); + + return Optional.of(deleteCount == idList.size()); + } + + @Override + public List similaritySearch(String query) { + return similaritySearch(SearchRequest.query(query)); + } + + @Override + public List similaritySearch(SearchRequest request) { + + String nativeFilterExpressions = (request.getFilterExpression() != null) + ? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : ""; + + List queryEmbedding = this.embeddingClient.embed(request.getQuery()); + var vectorSearch = new VectorSearchAggregation(queryEmbedding, this.config.pathName, this.config.numCandidates, + this.config.vectorIndexName, request.getTopK(), nativeFilterExpressions); + + Aggregation aggregation = Aggregation.newAggregation(vectorSearch, + Aggregation.addFields() + .addField(SCORE_FIELD_NAME) + .withValueOfExpression("{\"$meta\":\"vectorSearchScore\"}") + .build(), + Aggregation.match(new Criteria(SCORE_FIELD_NAME).gte(request.getSimilarityThreshold()))); + + return this.mongoTemplate.aggregate(aggregation, this.config.collectionName, BasicDBObject.class) + .getMappedResults() + .stream() + .map(this::mapBasicDbObject) + .toList(); + } + + public static class MongoDBVectorStoreConfig { + + private final String collectionName; + + private final String vectorIndexName; + + private final String pathName; + + private final List metadataFieldsToFilter; + + private final int numCandidates; + + private MongoDBVectorStoreConfig(Builder builder) { + this.collectionName = builder.collectionName; + this.vectorIndexName = builder.vectorIndexName; + this.pathName = builder.pathName; + this.numCandidates = builder.numCandidates; + this.metadataFieldsToFilter = builder.metadataFieldsToFilter; + } + + public static Builder builder() { + return new Builder(); + } + + public static MongoDBVectorStoreConfig defaultConfig() { + return builder().build(); + } + + public static class Builder { + + private String collectionName = DEFAULT_VECTOR_COLLECTION_NAME; + + private String vectorIndexName = DEFAULT_VECTOR_INDEX_NAME; + + private String pathName = DEFAULT_PATH_NAME; + + private int numCandidates = DEFAULT_NUM_CANDIDATES; + + private List metadataFieldsToFilter = Collections.emptyList(); + + private Builder() { + } + + /** + * Configures the collection to use This must match the name of the collection + * for the Vector Search Index in Atlas + * @param collectionName + * @return this builder + */ + public Builder withCollectionName(String collectionName) { + Assert.notNull(collectionName, "Collection Name must not be null"); + Assert.notNull(collectionName, "Collection Name must not be empty"); + this.collectionName = collectionName; + return this; + } + + /** + * Configures the vector index name. This must match the name of the Vector + * Search Index Name in Atlas + * @param vectorIndexName + * @return this builder + */ + public Builder withVectorIndexName(String vectorIndexName) { + Assert.notNull(vectorIndexName, "Vector Index Name must not be null"); + Assert.notNull(vectorIndexName, "Vector Index Name must not be empty"); + this.vectorIndexName = vectorIndexName; + return this; + } + + /** + * Configures the path name. This must match the name of the field indexed for + * the Vector Search Index in Atlas + * @param pathName + * @return this builder + */ + public Builder withPathName(String pathName) { + Assert.notNull(pathName, "Path Name must not be null"); + Assert.notNull(pathName, "Path Name must not be empty"); + this.pathName = pathName; + return this; + } + + public Builder withMetadataFieldsToFilter(List metadataFieldsToFilter) { + Assert.notEmpty(metadataFieldsToFilter, "Fields list must not be empty"); + this.metadataFieldsToFilter = metadataFieldsToFilter; + return this; + } + + public MongoDBVectorStoreConfig build() { + return new MongoDBVectorStoreConfig(this); + } + + } + + } + +} \ No newline at end of file diff --git a/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/VectorSearchAggregation.java b/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/VectorSearchAggregation.java new file mode 100644 index 000000000..572749af4 --- /dev/null +++ b/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/VectorSearchAggregation.java @@ -0,0 +1,42 @@ +/* + * 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.vectorstore; + +import org.bson.Document; +import org.springframework.data.mongodb.core.aggregation.AggregationOperation; +import org.springframework.data.mongodb.core.aggregation.AggregationOperationContext; +import org.springframework.lang.NonNull; + +import java.util.List; + +record VectorSearchAggregation(List embeddings, String path, int numCandidates, String index, int count, + String filter) implements AggregationOperation { + + @SuppressWarnings("null") + @Override + public org.bson.Document toDocument(@NonNull AggregationOperationContext context) { + var vectorSearch = new Document("queryVector", embeddings).append("path", path) + .append("numCandidates", numCandidates) + .append("index", index) + .append("limit", count); + if (!filter.isEmpty()) { + vectorSearch.append("filter", Document.parse(filter)); + } + var doc = new Document("$vectorSearch", vectorSearch); + + return context.getMappedObject(doc); + } +} \ No newline at end of file diff --git a/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDBAtlasContainer.java b/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDBAtlasContainer.java new file mode 100644 index 000000000..2d2dd6aa5 --- /dev/null +++ b/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDBAtlasContainer.java @@ -0,0 +1,39 @@ +/* + * 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.vectorstore; + +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.containers.wait.strategy.Wait; + +import java.time.Duration; + +public class MongoDBAtlasContainer extends GenericContainer { + + public MongoDBAtlasContainer() { + super("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)).withReuse(true); + } + + public String getConnectionString() { + return String.format("mongodb://root:root@%s:%s/?directConnection=true", getHost(), getMappedPort(27778)); + } + +} \ No newline at end of file diff --git a/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDBAtlasFilterConverterTest.java b/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDBAtlasFilterConverterTest.java new file mode 100644 index 000000000..6ab38c55c --- /dev/null +++ b/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDBAtlasFilterConverterTest.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.vectorstore; + +import java.util.List; + +import org.junit.jupiter.api.Test; + +import org.springframework.ai.vectorstore.filter.Filter.Expression; +import org.springframework.ai.vectorstore.filter.Filter.Group; +import org.springframework.ai.vectorstore.filter.Filter.Key; +import org.springframework.ai.vectorstore.filter.Filter.Value; +import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.AND; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.EQ; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.GTE; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.IN; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.LTE; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.NE; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.NIN; +import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.OR; + +/** + * @author Christopher Smith + */ +public class MongoDBAtlasFilterConverterTest { + + FilterExpressionConverter converter = new MongoDBAtlasFilterExpressionConverter(); + + @Test + public void testEQ() { + // country == "BG" + String vectorExpr = converter.convertExpression(new Expression(EQ, new Key("country"), new Value("BG"))); + assertThat(vectorExpr).isEqualTo("{\"metadata.country\":{$eq:\"BG\"}}"); + } + + @Test + public void tesEqAndGte() { + // genre == "drama" AND year >= 2020 + String vectorExpr = converter + .convertExpression(new Expression(AND, new Expression(EQ, new Key("genre"), new Value("drama")), + new Expression(GTE, new Key("year"), new Value(2020)))); + assertThat(vectorExpr) + .isEqualTo("{$and:[{\"metadata.genre\":{$eq:\"drama\"}},{\"metadata.year\":{$gte:2020}}]}"); + } + + @Test + public void tesIn() { + // genre in ["comedy", "documentary", "drama"] + String vectorExpr = converter.convertExpression( + new Expression(IN, new Key("genre"), new Value(List.of("comedy", "documentary", "drama")))); + assertThat(vectorExpr).isEqualTo("{\"metadata.genre\":{$in:[\"comedy\",\"documentary\",\"drama\"]}}"); + } + + @Test + public void testNe() { + // year >= 2020 OR country == "BG" AND city != "Sofia" + String vectorExpr = converter + .convertExpression(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)), + new Expression(AND, new Expression(EQ, new Key("country"), new Value("BG")), + new Expression(NE, new Key("city"), new Value("Sofia"))))); + assertThat(vectorExpr).isEqualTo( + "{$or:[{\"metadata.year\":{$gte:2020}},{$and:[{\"metadata.country\":{$eq:\"BG\"}},{\"metadata.city\":{$ne:\"Sofia\"}}]}]}"); + } + + @Test + public void testGroup() { + // (year >= 2020 OR country == "BG") AND city NIN ["Sofia", "Plovdiv"] + String vectorExpr = converter.convertExpression(new Expression(AND, + new Group(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)), + new Expression(EQ, new Key("country"), new Value("BG")))), + new Expression(NIN, new Key("city"), new Value(List.of("Sofia", "Plovdiv"))))); + assertThat(vectorExpr).isEqualTo( + "{$and:[{$or:[{\"metadata.year\":{$gte:2020}},{\"metadata.country\":{$eq:\"BG\"}}]},{\"metadata.city\":{$nin:[\"Sofia\",\"Plovdiv\"]}}]}"); + } + + @Test + public void testBoolean() { + // isOpen == true AND year >= 2020 AND country IN ["BG", "NL", "US"] + String vectorExpr = converter.convertExpression(new Expression(AND, + new Expression(AND, new Expression(EQ, new Key("isOpen"), new Value(true)), + new Expression(GTE, new Key("year"), new Value(2020))), + new Expression(IN, new Key("country"), new Value(List.of("BG", "NL", "US"))))); + + assertThat(vectorExpr).isEqualTo( + "{$and:[{$and:[{\"metadata.isOpen\":{$eq:true}},{\"metadata.year\":{$gte:2020}}]},{\"metadata.country\":{$in:[\"BG\",\"NL\",\"US\"]}}]}"); + } + + @Test + public void testDecimal() { + // temperature >= -15.6 && temperature <= +20.13 + String vectorExpr = converter + .convertExpression(new Expression(AND, new Expression(GTE, new Key("temperature"), new Value(-15.6)), + new Expression(LTE, new Key("temperature"), new Value(20.13)))); + + assertThat(vectorExpr) + .isEqualTo("{$and:[{\"metadata.temperature\":{$gte:-15.6}},{\"metadata.temperature\":{$lte:20.13}}]}"); + } + + @Test + public void testComplexIdentifiers() { + String vectorExpr = converter + .convertExpression(new Expression(EQ, new Key("\"country 1 2 3\""), new Value("BG"))); + assertThat(vectorExpr).isEqualTo("{\"metadata.country 1 2 3\":{$eq:\"BG\"}}"); + + vectorExpr = converter.convertExpression(new Expression(EQ, new Key("'country 1 2 3'"), new Value("BG"))); + assertThat(vectorExpr).isEqualTo("{\"metadata.country 1 2 3\":{$eq:\"BG\"}}"); + } + +} 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 new file mode 100644 index 000000000..afc2dcbd7 --- /dev/null +++ b/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/MongoDBAtlasVectorStoreIT.java @@ -0,0 +1,216 @@ +/* + * 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.vectorstore; + +import com.mongodb.client.MongoClient; +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; +import org.springframework.data.mongodb.core.MongoTemplate; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; + +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Chris Smith + */ + +@Testcontainers +class MongoDBAtlasVectorStoreIT { + + @Container + private static MongoDBAtlasContainer container = new MongoDBAtlasContainer(); + + 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"), + String.format("spring.data.mongodb.uri=" + container.getConnectionString())); + + @BeforeEach + public void beforeEach() { + contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> { + MongoTemplate mongoTemplate = context.getBean(MongoTemplate.class); + mongoTemplate.getCollection("vector_store").deleteMany(new org.bson.Document()); + }); + } + + @Test + void vectorStoreTest() { + contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + + 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"))); + + 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(); + + }); + } + + @Test + void documentUpdateTest() { + contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + + Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!", + Collections.singletonMap("meta1", "meta1")); + + vectorStore.add(List.of(document)); + Thread.sleep(5000); // Await a second for the document to be indexed + + List results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(5)); + + assertThat(results).hasSize(1); + Document resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(document.getId()); + assertThat(resultDoc.getContent()).isEqualTo("Spring AI rocks!!"); + assertThat(resultDoc.getMetadata()).containsEntry("meta1", "meta1"); + + Document sameIdDocument = new Document(document.getId(), + "The World is Big and Salvation Lurks Around the Corner", + Collections.singletonMap("meta2", "meta2")); + + vectorStore.add(List.of(sameIdDocument)); + + results = vectorStore.similaritySearch(SearchRequest.query("FooBar").withTopK(5)); + + assertThat(results).hasSize(1); + resultDoc = results.get(0); + assertThat(resultDoc.getId()).isEqualTo(document.getId()); + assertThat(resultDoc.getContent()).isEqualTo("The World is Big and Salvation Lurks Around the Corner"); + assertThat(resultDoc.getMetadata()).containsEntry("meta2", "meta2"); + }); + } + + @Test + void searchWithFilters() { + contextRunner.withConfiguration(AutoConfigurations.of(OpenAiAutoConfiguration.class)).run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + + var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", + Map.of("country", "BG", "year", 2020)); + var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", + Map.of("country", "NL")); + var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", + Map.of("country", "BG", "year", 2023)); + + vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); + Thread.sleep(5000); // Await a second for the document to be indexed + + List results = vectorStore.similaritySearch(SearchRequest.query("The World").withTopK(5)); + assertThat(results).hasSize(3); + + results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("country == 'NL'")); + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId()); + + results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("country == 'BG'")); + + assertThat(results).hasSize(2); + assertThat(results.get(0).getId()).isIn(bgDocument.getId(), bgDocument2.getId()); + assertThat(results.get(1).getId()).isIn(bgDocument.getId(), bgDocument2.getId()); + + results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("country == 'BG' && year == 2020")); + + assertThat(results).hasSize(1); + assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId()); + + results = vectorStore.similaritySearch(SearchRequest.query("The World") + .withTopK(5) + .withSimilarityThresholdAll() + .withFilterExpression("NOT(country == 'BG' && year == 2020)")); + + assertThat(results).hasSize(2); + assertThat(results.get(0).getId()).isIn(nlDocument.getId(), bgDocument2.getId()); + assertThat(results.get(1).getId()).isIn(nlDocument.getId(), bgDocument2.getId()); + + }); + } + + @SpringBootConfiguration + @EnableAutoConfiguration + public static class TestApplication { + + @Bean + public VectorStore vectorStore(MongoTemplate mongoTemplate, EmbeddingClient embeddingClient) { + return new MongoDBAtlasVectorStore(mongoTemplate, embeddingClient, + MongoDBAtlasVectorStore.MongoDBVectorStoreConfig.builder() + .withMetadataFieldsToFilter(List.of("country", "year")) + .build()); + } + + @Bean + public MongoTemplate mongoTemplate(MongoClient mongoClient) { + return new MongoTemplate(mongoClient, "springaisample"); + } + + @Bean + public EmbeddingClient embeddingClient() { + return new OpenAiEmbeddingClient(new OpenAiApi(System.getenv("OPENAI_API_KEY"))); + } + + } + +} \ No newline at end of file diff --git a/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/VectorSearchAggregationTest.java b/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/VectorSearchAggregationTest.java new file mode 100644 index 000000000..f489df3da --- /dev/null +++ b/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/VectorSearchAggregationTest.java @@ -0,0 +1,63 @@ +/* + * 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.vectorstore; + +import org.bson.Document; +import org.junit.jupiter.api.Test; +import org.springframework.data.mongodb.core.aggregation.Aggregation; + +import java.util.List; + +import static org.junit.jupiter.api.Assertions.*; + +class VectorSearchAggregationTest { + + @Test + void toDocumentNoFilter() { + var vectorSearchAggregation = new VectorSearchAggregation(List.of(1.0, 2.0, 3.0), "embedding", 10, + "vector_store", 10, ""); + var aggregation = Aggregation.newAggregation(vectorSearchAggregation); + var document = aggregation.toDocument("vector_store", Aggregation.DEFAULT_CONTEXT); + + var vectorSearchDocument = new Document("$vectorSearch", + new Document("queryVector", List.of(1.0, 2.0, 3.0)).append("path", "embedding") + .append("numCandidates", 10) + .append("index", "vector_store") + .append("limit", 10)); + var expected = new Document().append("aggregate", "vector_store") + .append("pipeline", List.of(vectorSearchDocument)); + assertEquals(expected, document); + } + + @Test + void toDocumentWithFilter() { + var vectorSearchAggregation = new VectorSearchAggregation(List.of(1.0, 2.0, 3.0), "embedding", 10, + "vector_store", 10, "{\"metadata.country\":{$eq:\"BG\"}}"); + var aggregation = Aggregation.newAggregation(vectorSearchAggregation); + var document = aggregation.toDocument("vector_store", Aggregation.DEFAULT_CONTEXT); + + var vectorSearchDocument = new Document("$vectorSearch", + new Document("queryVector", List.of(1.0, 2.0, 3.0)).append("path", "embedding") + .append("numCandidates", 10) + .append("index", "vector_store") + .append("filter", new Document("metadata.country", new Document().append("$eq", "BG"))) + .append("limit", 10)); + var expected = new Document().append("aggregate", "vector_store") + .append("pipeline", List.of(vectorSearchDocument)); + assertEquals(expected, document); + } + +} \ No newline at end of file