Skip to content

Commit 9e19b4d

Browse files
author
Milder Hernandez Cagua
committed
Add search for Redis Hash Set
1 parent 6746c86 commit 9e19b4d

12 files changed

Lines changed: 315 additions & 181 deletions

File tree

api-test/integration-tests/src/test/java/com/microsoft/semantickernel/tests/connectors/memory/redis/RedisHashSetVectorStoreRecordCollectionTest.java

Lines changed: 68 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -189,7 +189,7 @@ public void getBatchAsync(RecordCollectionOptions options) {
189189
List<String> ids = new ArrayList<>();
190190
hotels.forEach(hotel -> ids.add(hotel.getId()));
191191

192-
List<Hotel> retrievedHotels = recordCollection.getBatchAsync(ids, null).block();
192+
List<Hotel> retrievedHotels = recordCollection.getBatchAsync(ids, new GetRecordOptions(true)).block();
193193

194194
assertNotNull(retrievedHotels);
195195
assertEquals(hotels.size(), retrievedHotels.size());
@@ -341,38 +341,71 @@ private static Stream<Arguments> provideSearchParameters() {
341341
);
342342
}
343343

344-
// @ParameterizedTest
345-
// @MethodSource("provideSearchParameters")
346-
// public void search(RecordCollectionOptions options, String embeddingName) {
347-
// String collectionName = "search" + embeddingName;
348-
// RedisHashSetVectorStoreRecordCollection<Hotel> recordCollection = buildRecordCollection(optionsMap.get(options), collectionName);
349-
//
350-
// List<Hotel> hotels = getHotels();
351-
// recordCollection.upsertBatchAsync(hotels, null).block();
352-
//
353-
// VectorSearchOptions searchOptions = VectorSearchOptions.builder()
354-
// .withVectorFieldName(embeddingName)
355-
// .withLimit(3)
356-
// .build();
357-
//
358-
// // Embeddings similar to the third hotel
359-
// List<VectorSearchResult<Hotel>> results = recordCollection.searchAsync(SEARCH_EMBEDDINGS, searchOptions).block();
360-
// assertNotNull(results);
361-
// assertEquals(3, results.size());
362-
// // The third hotel should be the most similar
363-
// assertEquals(hotels.get(2).getId(), results.get(0).getRecord().getId());
364-
//
365-
// searchOptions = VectorSearchOptions.builder()
366-
// .withVectorFieldName(embeddingName)
367-
// .withOffset(1)
368-
// .withLimit(-100)
369-
// .build();
370-
//
371-
// // Skip the first result
372-
// results = recordCollection.searchAsync(SEARCH_EMBEDDINGS, searchOptions).block();
373-
// assertNotNull(results);
374-
// assertEquals(1, results.size());
375-
// // The first hotel should be the most similar
376-
// assertEquals(hotels.get(0).getId(), results.get(0).getRecord().getId());
377-
// }
344+
@ParameterizedTest
345+
@MethodSource("provideSearchParameters")
346+
public void search(RecordCollectionOptions options, String embeddingName) {
347+
String collectionName = "search" + embeddingName;
348+
RedisHashSetVectorStoreRecordCollection<Hotel> recordCollection = buildRecordCollection(optionsMap.get(options), collectionName);
349+
350+
List<Hotel> hotels = getHotels();
351+
recordCollection.upsertBatchAsync(hotels, null).block();
352+
353+
VectorSearchOptions searchOptions = VectorSearchOptions.builder()
354+
.withVectorFieldName(embeddingName)
355+
.build();
356+
357+
// Embeddings similar to the third hotel
358+
List<VectorSearchResult<Hotel>> results = recordCollection.searchAsync(SEARCH_EMBEDDINGS, searchOptions).block();
359+
assertNotNull(results);
360+
assertEquals(VectorSearchOptions.DEFAULT_RESULT_LIMIT, results.size());
361+
// The third hotel should be the most similar
362+
assertEquals(hotels.get(2).getId(), results.get(0).getRecord().getId());
363+
assertNull(results.get(0).getRecord().getEuclidean());
364+
}
365+
366+
@ParameterizedTest
367+
@MethodSource("provideSearchParameters")
368+
public void searchWithVectors(RecordCollectionOptions options, String embeddingName) {
369+
String collectionName = "search" + embeddingName;
370+
RedisHashSetVectorStoreRecordCollection<Hotel> recordCollection = buildRecordCollection(optionsMap.get(options), collectionName);
371+
372+
List<Hotel> hotels = getHotels();
373+
recordCollection.upsertBatchAsync(hotels, null).block();
374+
375+
VectorSearchOptions searchOptions = VectorSearchOptions.builder()
376+
.withVectorFieldName(embeddingName)
377+
.withIncludeVectors(true)
378+
.build();
379+
380+
// Embeddings similar to the third hotel
381+
List<VectorSearchResult<Hotel>> results = recordCollection.searchAsync(SEARCH_EMBEDDINGS, searchOptions).block();
382+
assertNotNull(results);
383+
assertEquals(VectorSearchOptions.DEFAULT_RESULT_LIMIT, results.size());
384+
// The third hotel should be the most similar
385+
assertEquals(hotels.get(2).getId(), results.get(0).getRecord().getId());
386+
assertNotNull(results.get(0).getRecord().getEuclidean());
387+
}
388+
389+
@ParameterizedTest
390+
@MethodSource("provideSearchParameters")
391+
public void searchWithOffSet(RecordCollectionOptions options, String embeddingName) {
392+
String collectionName = "search" + embeddingName;
393+
RedisHashSetVectorStoreRecordCollection<Hotel> recordCollection = buildRecordCollection(optionsMap.get(options), collectionName);
394+
395+
List<Hotel> hotels = getHotels();
396+
recordCollection.upsertBatchAsync(hotels, null).block();
397+
398+
VectorSearchOptions searchOptions = VectorSearchOptions.builder()
399+
.withVectorFieldName(embeddingName)
400+
.withOffset(1)
401+
.withLimit(4)
402+
.build();
403+
404+
// Embeddings similar to the third hotel
405+
List<VectorSearchResult<Hotel>> results = recordCollection.searchAsync(SEARCH_EMBEDDINGS, searchOptions).block();
406+
assertNotNull(results);
407+
assertEquals(4, results.size());
408+
// The first hotel should be the most similar
409+
assertEquals(hotels.get(0).getId(), results.get(0).getRecord().getId());
410+
}
378411
}

api-test/integration-tests/src/test/java/com/microsoft/semantickernel/tests/connectors/memory/redis/RedisJsonVectorStoreRecordCollectionTest.java

Lines changed: 41 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -351,27 +351,61 @@ public void search(RecordCollectionOptions options, String embeddingName) {
351351

352352
VectorSearchOptions searchOptions = VectorSearchOptions.builder()
353353
.withVectorFieldName(embeddingName)
354-
.withLimit(3)
355354
.build();
356355

357356
// Embeddings similar to the third hotel
358357
List<VectorSearchResult<Hotel>> results = recordCollection.searchAsync(SEARCH_EMBEDDINGS, searchOptions).block();
359358
assertNotNull(results);
360-
assertEquals(3, results.size());
359+
assertEquals(VectorSearchOptions.DEFAULT_RESULT_LIMIT, results.size());
361360
// The third hotel should be the most similar
362361
assertEquals(hotels.get(2).getId(), results.get(0).getRecord().getId());
362+
assertNull(results.get(0).getRecord().getEuclidean());
363+
}
364+
365+
@ParameterizedTest
366+
@MethodSource("provideSearchParameters")
367+
public void searchWithVectors(RecordCollectionOptions options, String embeddingName) {
368+
String collectionName = "search" + embeddingName;
369+
RedisJsonVectorStoreRecordCollection<Hotel> recordCollection = buildrecordCollection(optionsMap.get(options), collectionName);
370+
371+
List<Hotel> hotels = getHotels();
372+
recordCollection.upsertBatchAsync(hotels, null).block();
373+
374+
VectorSearchOptions searchOptions = VectorSearchOptions.builder()
375+
.withVectorFieldName(embeddingName)
376+
.withIncludeVectors(true)
377+
.build();
378+
379+
// Embeddings similar to the third hotel
380+
List<VectorSearchResult<Hotel>> results = recordCollection.searchAsync(SEARCH_EMBEDDINGS, searchOptions).block();
381+
assertNotNull(results);
382+
assertEquals(VectorSearchOptions.DEFAULT_RESULT_LIMIT, results.size());
383+
// The third hotel should be the most similar
384+
assertEquals(hotels.get(2).getId(), results.get(0).getRecord().getId());
385+
assertNotNull(results.get(0).getRecord().getEuclidean());
386+
}
387+
388+
@ParameterizedTest
389+
@MethodSource("provideSearchParameters")
390+
public void searchWithOffSet(RecordCollectionOptions options, String embeddingName) {
391+
String collectionName = "search" + embeddingName;
392+
RedisJsonVectorStoreRecordCollection<Hotel> recordCollection = buildrecordCollection(optionsMap.get(options), collectionName);
393+
394+
List<Hotel> hotels = getHotels();
395+
recordCollection.upsertBatchAsync(hotels, null).block();
363396

364-
searchOptions = VectorSearchOptions.builder()
397+
VectorSearchOptions searchOptions = VectorSearchOptions.builder()
365398
.withVectorFieldName(embeddingName)
366399
.withOffset(1)
367-
.withLimit(-100)
400+
.withLimit(4)
368401
.build();
369402

370-
// Skip the first result
371-
results = recordCollection.searchAsync(SEARCH_EMBEDDINGS, searchOptions).block();
403+
// Embeddings similar to the third hotel
404+
List<VectorSearchResult<Hotel>> results = recordCollection.searchAsync(SEARCH_EMBEDDINGS, searchOptions).block();
372405
assertNotNull(results);
373-
assertEquals(1, results.size());
406+
assertEquals(4, results.size());
374407
// The first hotel should be the most similar
375408
assertEquals(hotels.get(0).getId(), results.get(0).getRecord().getId());
376409
}
410+
377411
}

connectors/semantickernel-connectors-memory-redis/pom.xml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@
4141
<dependency>
4242
<groupId>redis.clients</groupId>
4343
<artifactId>jedis</artifactId>
44-
<version>5.1.0</version>
44+
<version>5.2.0-beta5</version>
4545
<scope>provided</scope>
4646
</dependency>
4747
<dependency>

semantickernel-bom/pom.xml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -120,7 +120,7 @@
120120
<dependency>
121121
<groupId>redis.clients</groupId>
122122
<artifactId>jedis</artifactId>
123-
<version>5.1.0</version>
123+
<version>5.2.0-beta5</version>
124124
</dependency>
125125

126126
<dependency>

0 commit comments

Comments
 (0)