Skip to content

Commit 9065008

Browse files
strogiyotecstrogiyotec
committed
JAVA-6099 Add new fields for Auto embedding
1 parent 6c4ab1b commit 9065008

1 file changed

Lines changed: 218 additions & 1 deletion

File tree

‎driver-sync/src/test/functional/com/mongodb/client/vector/AbstractAutomatedEmbeddingVectorSearchFunctionalTest.java‎

Lines changed: 218 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@
3434
import org.junit.jupiter.api.BeforeEach;
3535
import org.junit.jupiter.api.DisplayName;
3636
import org.junit.jupiter.api.Test;
37+
import org.junit.jupiter.params.ParameterizedTest;
38+
import org.junit.jupiter.params.provider.ValueSource;
3739

3840
import java.util.ArrayList;
3941
import java.util.Collections;
@@ -44,6 +46,7 @@
4446
import static com.mongodb.client.model.Aggregates.vectorSearch;
4547
import static com.mongodb.client.model.search.SearchPath.fieldPath;
4648
import static com.mongodb.client.model.search.VectorSearchOptions.approximateVectorSearchOptions;
49+
import static com.mongodb.client.model.search.VectorSearchOptions.exactVectorSearchOptions;
4750
import static com.mongodb.client.model.search.VectorSearchQuery.textQuery;
4851
import static java.util.Arrays.asList;
4952
import static org.bson.codecs.configuration.CodecRegistries.fromProviders;
@@ -74,7 +77,7 @@ public void setUp() {
7477
//TODO-JAVA-6059 remove this line when Atlas Vector Search with automated embedding is generally available
7578
// right now atlas search with automated embedding is in private preview and
7679
// only available via a custom docker image
77-
Assumptions.assumeTrue(false);
80+
Assumptions.assumeTrue(true);
7881

7982
super.beforeEach();
8083
mongoClient = getMongoClient(getMongoClientSettingsBuilder()
@@ -210,6 +213,220 @@ private void insertDocumentsForEmbedding() {
210213
));
211214
}
212215

216+
@Test
217+
@DisplayName("should create auto embedding index with all optional fields")
218+
void shouldCreateAutoEmbeddingIndexWithAllOptionalFields() {
219+
mongoClient.getDatabase(getDatabaseName()).createCollection(getCollectionName());
220+
SearchIndexModel indexModel = new SearchIndexModel(
221+
INDEX_NAME,
222+
new Document(
223+
"fields",
224+
Collections.singletonList(
225+
new Document("type", "autoEmbed")
226+
.append("modality", "text")
227+
.append("path", FIELD_SEARCH_PATH)
228+
.append("model", "voyage-4-large")
229+
.append("quantization", "binary")
230+
.append("similarity", "euclidean")
231+
)),
232+
SearchIndexType.vectorSearch()
233+
);
234+
List<String> result = documentCollection.createSearchIndexes(Collections.singletonList(indexModel));
235+
Assertions.assertFalse(result.isEmpty());
236+
}
237+
238+
@ParameterizedTest(name = "should create auto embedding index with {0} quantization")
239+
@ValueSource(strings = {"float", "scalar", "binary", "binaryNoRescore"})
240+
void shouldCreateAutoEmbeddingIndexWithQuantization(final String quantization) {
241+
final String indexName = INDEX_NAME + "_" + quantization;
242+
mongoClient.getDatabase(getDatabaseName()).createCollection(getCollectionName());
243+
SearchIndexModel indexModel = new SearchIndexModel(
244+
indexName,
245+
new Document(
246+
"fields",
247+
Collections.singletonList(
248+
new Document("type", "autoEmbed")
249+
.append("modality", "text")
250+
.append("path", FIELD_SEARCH_PATH)
251+
.append("model", "voyage-4-large")
252+
.append("quantization", quantization)
253+
)),
254+
SearchIndexType.vectorSearch()
255+
);
256+
List<String> result = documentCollection.createSearchIndexes(Collections.singletonList(indexModel));
257+
Assertions.assertFalse(result.isEmpty());
258+
}
259+
260+
@Test
261+
@DisplayName("should create auto embedding index with custom numDimensions")
262+
void shouldCreateAutoEmbeddingIndexWithCustomNumDimensions() {
263+
mongoClient.getDatabase(getDatabaseName()).createCollection(getCollectionName());
264+
SearchIndexModel indexModel = new SearchIndexModel(
265+
INDEX_NAME,
266+
new Document(
267+
"fields",
268+
Collections.singletonList(
269+
new Document("type", "autoEmbed")
270+
.append("modality", "text")
271+
.append("path", FIELD_SEARCH_PATH)
272+
.append("model", "voyage-4-large")
273+
.append("numDimensions", 512)
274+
)),
275+
SearchIndexType.vectorSearch()
276+
);
277+
List<String> result = documentCollection.createSearchIndexes(Collections.singletonList(indexModel));
278+
Assertions.assertFalse(result.isEmpty());
279+
}
280+
281+
@Test
282+
@DisplayName("should create auto embedding index with filter field")
283+
void shouldCreateAutoEmbeddingIndexWithFilterField() {
284+
mongoClient.getDatabase(getDatabaseName()).createCollection(getCollectionName());
285+
SearchIndexModel indexModel = new SearchIndexModel(
286+
INDEX_NAME,
287+
new Document(
288+
"fields",
289+
asList(
290+
new Document("type", "autoEmbed")
291+
.append("modality", "text")
292+
.append("path", FIELD_SEARCH_PATH)
293+
.append("model", "voyage-4-large"),
294+
new Document("type", "filter")
295+
.append("path", "director")
296+
)),
297+
SearchIndexType.vectorSearch()
298+
);
299+
List<String> result = documentCollection.createSearchIndexes(Collections.singletonList(indexModel));
300+
Assertions.assertFalse(result.isEmpty());
301+
}
302+
303+
@Test
304+
@DisplayName("should fail when mixing vector and autoEmbed types in the same index")
305+
void shouldFailWhenMixingVectorAndAutoEmbedTypes() {
306+
mongoClient.getDatabase(getDatabaseName()).createCollection(getCollectionName());
307+
SearchIndexModel indexModel = new SearchIndexModel(
308+
INDEX_NAME,
309+
new Document(
310+
"fields",
311+
asList(
312+
new Document("type", "autoEmbed")
313+
.append("modality", "text")
314+
.append("path", FIELD_SEARCH_PATH)
315+
.append("model", "voyage-4-large"),
316+
new Document("type", "vector")
317+
.append("path", "plot_embedding")
318+
.append("numDimensions", 1024)
319+
.append("similarity", "cosine")
320+
)),
321+
SearchIndexType.vectorSearch()
322+
);
323+
Assertions.assertThrows(
324+
MongoCommandException.class,
325+
() -> documentCollection.createSearchIndexes(Collections.singletonList(indexModel)),
326+
"Expected index creation to fail because vector and autoEmbed types cannot be mixed"
327+
);
328+
}
329+
330+
@Test
331+
@DisplayName("should fail when duplicate paths are used")
332+
void shouldFailWhenDuplicatePathsAreUsed() {
333+
mongoClient.getDatabase(getDatabaseName()).createCollection(getCollectionName());
334+
SearchIndexModel indexModel = new SearchIndexModel(
335+
INDEX_NAME,
336+
new Document(
337+
"fields",
338+
asList(
339+
new Document("type", "autoEmbed")
340+
.append("modality", "text")
341+
.append("path", FIELD_SEARCH_PATH)
342+
.append("model", "voyage-4-large"),
343+
new Document("type", "autoEmbed")
344+
.append("modality", "text")
345+
.append("path", FIELD_SEARCH_PATH)
346+
.append("model", "voyage-4-large")
347+
)),
348+
SearchIndexType.vectorSearch()
349+
);
350+
Assertions.assertThrows(
351+
MongoCommandException.class,
352+
() -> documentCollection.createSearchIndexes(Collections.singletonList(indexModel)),
353+
"Expected index creation to fail because of duplicate paths"
354+
);
355+
}
356+
357+
@Test
358+
@DisplayName("should fail when autoEmbed field is used as filter field")
359+
void shouldFailWhenAutoEmbedFieldUsedAsFilterField() {
360+
mongoClient.getDatabase(getDatabaseName()).createCollection(getCollectionName());
361+
SearchIndexModel indexModel = new SearchIndexModel(
362+
INDEX_NAME,
363+
new Document(
364+
"fields",
365+
asList(
366+
new Document("type", "autoEmbed")
367+
.append("modality", "text")
368+
.append("path", FIELD_SEARCH_PATH)
369+
.append("model", "voyage-4-large"),
370+
new Document("type", "filter")
371+
.append("path", FIELD_SEARCH_PATH)
372+
)),
373+
SearchIndexType.vectorSearch()
374+
);
375+
Assertions.assertThrows(
376+
MongoCommandException.class,
377+
() -> documentCollection.createSearchIndexes(Collections.singletonList(indexModel)),
378+
"Expected index creation to fail because autoEmbed field cannot be used as a filter field"
379+
);
380+
}
381+
382+
@Test
383+
@DisplayName("should create auto embedding index and run query with model override")
384+
void shouldCreateAutoEmbeddingIndexAndRunQueryWithModelOverride() throws InterruptedException {
385+
mongoClient.getDatabase(getDatabaseName()).createCollection(getCollectionName());
386+
createAutoEmbeddingIndex("voyage-4-large");
387+
TimeUnit.SECONDS.sleep(2L);
388+
insertDocumentsForEmbedding();
389+
TimeUnit.SECONDS.sleep(2L);
390+
391+
List<Bson> pipeline = asList(
392+
vectorSearch(
393+
fieldPath(FIELD_SEARCH_PATH),
394+
textQuery("movies about love").model("voyage-4-large"),
395+
INDEX_NAME,
396+
5L,
397+
approximateVectorSearchOptions(5L)
398+
)
399+
);
400+
List<Document> documents = documentCollection.aggregate(pipeline).into(new ArrayList<>());
401+
402+
Assertions.assertFalse(documents.isEmpty(), "Expected to get some results from vector search query");
403+
Assertions.assertEquals(MOVIE_NAME, documents.get(0).getString("title"));
404+
}
405+
406+
@Test
407+
@DisplayName("should create auto embedding index and run exact vector search query")
408+
void shouldCreateAutoEmbeddingIndexAndRunExactVectorSearchQuery() throws InterruptedException {
409+
mongoClient.getDatabase(getDatabaseName()).createCollection(getCollectionName());
410+
createAutoEmbeddingIndex("voyage-4-large");
411+
TimeUnit.SECONDS.sleep(2L);
412+
insertDocumentsForEmbedding();
413+
TimeUnit.SECONDS.sleep(2L);
414+
415+
List<Bson> pipeline = asList(
416+
vectorSearch(
417+
fieldPath(FIELD_SEARCH_PATH),
418+
textQuery("movies about love"),
419+
INDEX_NAME,
420+
5L,
421+
exactVectorSearchOptions()
422+
)
423+
);
424+
List<Document> documents = documentCollection.aggregate(pipeline).into(new ArrayList<>());
425+
426+
Assertions.assertFalse(documents.isEmpty(), "Expected to get some results from exact vector search query");
427+
Assertions.assertEquals(MOVIE_NAME, documents.get(0).getString("title"));
428+
}
429+
213430
private void createAutoEmbeddingIndex(final String modelName) {
214431
SearchIndexModel indexModel = new SearchIndexModel(
215432
INDEX_NAME,

0 commit comments

Comments
 (0)