Skip to content

Commit 41ef6b7

Browse files
strogiyotecstrogiyotec
authored
Test cases to test new fields for auto embedding index
The doc in JIRA list new fields that must be supported by the drivers but because the Search index doesn't use builders, java driver makes it possible to use new fields without changing API In this PR Added new test cases to test creating search index with 1. quantization 2. similarity [JAVA-6099](https://jira.mongodb.org/browse/JAVA-6099)
1 parent 99108af commit 41ef6b7

3 files changed

Lines changed: 307 additions & 102 deletions

File tree

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

Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,11 @@
3131
import org.junit.jupiter.api.Assertions;
3232
import org.junit.jupiter.api.Assumptions;
3333
import org.junit.jupiter.api.BeforeEach;
34+
import org.junit.jupiter.api.Disabled;
3435
import org.junit.jupiter.api.DisplayName;
3536
import org.junit.jupiter.api.Test;
37+
import org.junit.jupiter.params.ParameterizedTest;
38+
import org.junit.jupiter.params.provider.ValueSource;
3639

3740
import java.util.List;
3841
import java.util.Map;
@@ -79,6 +82,8 @@ public abstract class AbstractAtlasSearchIndexManagementProseTest {
7982
"{"
8083
+ " mappings: { dynamic: true }"
8184
+ "}");
85+
private static final String AUTO_EMBED_FIELD_PATH = "plot";
86+
private static final String AUTO_EMBED_INDEX_NAME = "voyage_4";
8287
private static final Document VECTOR_SEARCH_DEFINITION = Document.parse(
8388
"{"
8489
+ " fields: ["
@@ -281,6 +286,218 @@ public void shouldRequireExplicitTypeToCreateVectorSearchIndex() {
281286
VECTOR_SEARCH_DEFINITION));
282287
}
283288

289+
@Test
290+
@DisplayName("should fail when invalid model name was used for auto embedding index")
291+
void shouldFailWhenInvalidModelNameWasUsed() {
292+
//TODO-JAVA-6059 remove this assumption when auto embedding is generally available
293+
Assumptions.assumeTrue(false);
294+
295+
assertThrows(
296+
MongoCommandException.class,
297+
() -> createAutoEmbeddingIndex("test"),
298+
"Valid voyage model name was not used"
299+
);
300+
}
301+
302+
@Test
303+
@DisplayName("should fail to create auto embedding index without model")
304+
void shouldFailToCreateAutoEmbeddingIndexWithoutModel() {
305+
//TODO-JAVA-6059 remove this assumption when auto embedding is generally available
306+
Assumptions.assumeTrue(false);
307+
308+
SearchIndexModel indexModel = new SearchIndexModel(
309+
AUTO_EMBED_INDEX_NAME,
310+
new Document(
311+
"fields",
312+
singletonList(
313+
new Document("type", "autoEmbed")
314+
.append("modality", "text")
315+
.append("path", AUTO_EMBED_FIELD_PATH)
316+
)),
317+
SearchIndexType.vectorSearch()
318+
);
319+
assertThrows(
320+
MongoCommandException.class,
321+
() -> collection.createSearchIndexes(singletonList(indexModel)),
322+
"Expected index creation to fail because model is not specified"
323+
);
324+
}
325+
326+
@ParameterizedTest(name = "should create auto embedding index with {0} quantization")
327+
@ValueSource(strings = {"float", "scalar", "binary", "binaryNoRescore"})
328+
void shouldCreateAutoEmbeddingIndexWithQuantization(final String quantization) {
329+
//TODO-JAVA-6059 remove this assumption when auto embedding is generally available
330+
Assumptions.assumeTrue(false);
331+
332+
final String indexName = AUTO_EMBED_INDEX_NAME + "_" + quantization;
333+
SearchIndexModel indexModel = new SearchIndexModel(
334+
indexName,
335+
new Document(
336+
"fields",
337+
singletonList(
338+
new Document("type", "autoEmbed")
339+
.append("modality", "text")
340+
.append("path", AUTO_EMBED_FIELD_PATH)
341+
.append("model", "voyage-4-large")
342+
.append("quantization", quantization)
343+
)),
344+
SearchIndexType.vectorSearch()
345+
);
346+
List<String> result = collection.createSearchIndexes(singletonList(indexModel));
347+
Assertions.assertFalse(result.isEmpty());
348+
}
349+
350+
@Test
351+
@DisplayName("should create auto embedding index with custom numDimensions")
352+
@Disabled("Currently numDimensions can't be used, it fails with server error:"
353+
+ " 'Invalid numDimensions value for autoEmbed field. Expected an integer.'")
354+
void shouldCreateAutoEmbeddingIndexWithCustomNumDimensions() {
355+
//TODO-JAVA-6059 remove this assumption when auto embedding is generally available
356+
Assumptions.assumeTrue(false);
357+
358+
SearchIndexModel indexModel = new SearchIndexModel(
359+
AUTO_EMBED_INDEX_NAME,
360+
new Document(
361+
"fields",
362+
singletonList(
363+
new Document("type", "autoEmbed")
364+
.append("modality", "text")
365+
.append("path", AUTO_EMBED_FIELD_PATH)
366+
.append("model", "voyage-4-large")
367+
.append("numDimensions", 512)
368+
)),
369+
SearchIndexType.vectorSearch()
370+
);
371+
List<String> result = collection.createSearchIndexes(singletonList(indexModel));
372+
Assertions.assertFalse(result.isEmpty());
373+
}
374+
375+
@Test
376+
@DisplayName("should create auto embedding index with filter field")
377+
void shouldCreateAutoEmbeddingIndexWithFilterField() {
378+
//TODO-JAVA-6059 remove this assumption when auto embedding is generally available
379+
Assumptions.assumeTrue(false);
380+
381+
SearchIndexModel indexModel = new SearchIndexModel(
382+
AUTO_EMBED_INDEX_NAME,
383+
new Document(
384+
"fields",
385+
asList(
386+
new Document("type", "autoEmbed")
387+
.append("modality", "text")
388+
.append("path", AUTO_EMBED_FIELD_PATH)
389+
.append("model", "voyage-4-large"),
390+
new Document("type", "filter")
391+
.append("path", "director")
392+
)),
393+
SearchIndexType.vectorSearch()
394+
);
395+
List<String> result = collection.createSearchIndexes(singletonList(indexModel));
396+
Assertions.assertFalse(result.isEmpty());
397+
}
398+
399+
@Test
400+
@DisplayName("should fail when mixing vector and autoEmbed types in the same index")
401+
void shouldFailWhenMixingVectorAndAutoEmbedTypes() {
402+
//TODO-JAVA-6059 remove this assumption when auto embedding is generally available
403+
Assumptions.assumeTrue(false);
404+
405+
SearchIndexModel indexModel = new SearchIndexModel(
406+
AUTO_EMBED_INDEX_NAME,
407+
new Document(
408+
"fields",
409+
asList(
410+
new Document("type", "autoEmbed")
411+
.append("modality", "text")
412+
.append("path", AUTO_EMBED_FIELD_PATH)
413+
.append("model", "voyage-4-large"),
414+
new Document("type", "vector")
415+
.append("path", "plot_embedding")
416+
.append("numDimensions", 1024)
417+
.append("similarity", "cosine")
418+
)),
419+
SearchIndexType.vectorSearch()
420+
);
421+
assertThrows(
422+
MongoCommandException.class,
423+
() -> collection.createSearchIndexes(singletonList(indexModel)),
424+
"Expected index creation to fail because vector and autoEmbed types cannot be mixed"
425+
);
426+
}
427+
428+
@Test
429+
@DisplayName("should fail when duplicate paths are used in auto embedding index")
430+
void shouldFailWhenDuplicatePathsAreUsed() {
431+
//TODO-JAVA-6059 remove this assumption when auto embedding is generally available
432+
Assumptions.assumeTrue(false);
433+
434+
SearchIndexModel indexModel = new SearchIndexModel(
435+
AUTO_EMBED_INDEX_NAME,
436+
new Document(
437+
"fields",
438+
asList(
439+
new Document("type", "autoEmbed")
440+
.append("modality", "text")
441+
.append("path", AUTO_EMBED_FIELD_PATH)
442+
.append("model", "voyage-4-large"),
443+
new Document("type", "autoEmbed")
444+
.append("modality", "text")
445+
.append("path", AUTO_EMBED_FIELD_PATH)
446+
.append("model", "voyage-4-large")
447+
)),
448+
SearchIndexType.vectorSearch()
449+
);
450+
assertThrows(
451+
MongoCommandException.class,
452+
() -> collection.createSearchIndexes(singletonList(indexModel)),
453+
"Expected index creation to fail because of duplicate paths"
454+
);
455+
}
456+
457+
@Test
458+
@DisplayName("should fail when autoEmbed field is used as filter field")
459+
void shouldFailWhenAutoEmbedFieldUsedAsFilterField() {
460+
//TODO-JAVA-6059 remove this assumption when auto embedding is generally available
461+
Assumptions.assumeTrue(false);
462+
463+
SearchIndexModel indexModel = new SearchIndexModel(
464+
AUTO_EMBED_INDEX_NAME,
465+
new Document(
466+
"fields",
467+
asList(
468+
new Document("type", "autoEmbed")
469+
.append("modality", "text")
470+
.append("path", AUTO_EMBED_FIELD_PATH)
471+
.append("model", "voyage-4-large"),
472+
new Document("type", "filter")
473+
.append("path", AUTO_EMBED_FIELD_PATH)
474+
)),
475+
SearchIndexType.vectorSearch()
476+
);
477+
assertThrows(
478+
MongoCommandException.class,
479+
() -> collection.createSearchIndexes(singletonList(indexModel)),
480+
"Expected index creation to fail because autoEmbed field cannot be used as a filter field"
481+
);
482+
}
483+
484+
private void createAutoEmbeddingIndex(final String modelName) {
485+
SearchIndexModel indexModel = new SearchIndexModel(
486+
AUTO_EMBED_INDEX_NAME,
487+
new Document(
488+
"fields",
489+
singletonList(
490+
new Document("type", "autoEmbed")
491+
.append("modality", "text")
492+
.append("model", modelName)
493+
.append("path", AUTO_EMBED_FIELD_PATH)
494+
)),
495+
SearchIndexType.vectorSearch()
496+
);
497+
List<String> result = collection.createSearchIndexes(singletonList(indexModel));
498+
Assertions.assertFalse(result.isEmpty());
499+
}
500+
284501
private void assertIndexDeleted() throws InterruptedException {
285502
int attempts = MAX_WAIT_ATTEMPTS;
286503
while (collection.listSearchIndexes().first() != null && checkAttempt(attempts--)) {

0 commit comments

Comments
 (0)