Skip to content

Commit 2d7368d

Browse files
milderhcmilderhc
authored
Merge pull request #95 from milderhc/redis
Add RedisVectorStore and index/collection management
2 parents 734a40e + e6af0cc commit 2d7368d

File tree

10 files changed

+649
-177
lines changed

10 files changed

+649
-177
lines changed

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

+92-61
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
package com.microsoft.semantickernel.tests.connectors.memory.redis;
22

33
import com.microsoft.semantickernel.connectors.memory.redis.RedisVectorStoreRecordCollection;
4-
import com.microsoft.semantickernel.connectors.memory.redis.RedisVectorStoreOptions;
4+
import com.microsoft.semantickernel.connectors.memory.redis.RedisVectorStoreRecordCollectionOptions;
55
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDataField;
66
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition;
77
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordField;
@@ -11,7 +11,10 @@
1111
import com.microsoft.semantickernel.tests.connectors.memory.Hotel;
1212
import com.redis.testcontainers.RedisContainer;
1313
import org.junit.jupiter.api.BeforeAll;
14+
import org.junit.jupiter.api.MethodOrderer;
15+
import org.junit.jupiter.api.Order;
1416
import org.junit.jupiter.api.Test;
17+
import org.junit.jupiter.api.TestMethodOrder;
1518
import org.junit.jupiter.params.ParameterizedTest;
1619
import org.junit.jupiter.params.provider.EnumSource;
1720
import org.testcontainers.junit.jupiter.Container;
@@ -26,23 +29,25 @@
2629
import java.util.Map;
2730

2831
import static org.junit.jupiter.api.Assertions.assertEquals;
32+
import static org.junit.jupiter.api.Assertions.assertFalse;
2933
import static org.junit.jupiter.api.Assertions.assertNotNull;
3034
import static org.junit.jupiter.api.Assertions.assertNull;
3135

3236
@Testcontainers
37+
@TestMethodOrder(MethodOrderer.OrderAnnotation.class)
3338
public class RedisVectorStoreRecordCollectionTest {
3439

3540
@Container private static final RedisContainer redisContainer = new RedisContainer("redis/redis-stack:latest");
3641

37-
private static final Map<Options, RedisVectorStoreOptions<Hotel>> optionsMap = new HashMap<>();
42+
private static final Map<RecordCollectionOptions, RedisVectorStoreRecordCollectionOptions<Hotel>> optionsMap = new HashMap<>();
3843

39-
public enum Options {
44+
public enum RecordCollectionOptions {
4045
DEFAULT, WITH_CUSTOM_DEFINITION
4146
}
4247

4348
@BeforeAll
4449
static void setup() {
45-
optionsMap.put(Options.DEFAULT, RedisVectorStoreOptions.<Hotel>builder()
50+
optionsMap.put(RecordCollectionOptions.DEFAULT, RedisVectorStoreRecordCollectionOptions.<Hotel>builder()
4651
.withRecordClass(Hotel.class)
4752
.build());
4853

@@ -74,27 +79,21 @@ static void setup() {
7479
.build());
7580
VectorStoreRecordDefinition recordDefinition = VectorStoreRecordDefinition.fromFields(fields);
7681

77-
optionsMap.put(Options.WITH_CUSTOM_DEFINITION, RedisVectorStoreOptions.<Hotel>builder()
82+
optionsMap.put(RecordCollectionOptions.WITH_CUSTOM_DEFINITION, RedisVectorStoreRecordCollectionOptions.<Hotel>builder()
7883
.withRecordClass(Hotel.class)
7984
.withRecordDefinition(recordDefinition)
8085
.build());
8186
}
8287

83-
private RedisVectorStoreRecordCollection<Hotel> buildRecordStore(@Nonnull RedisVectorStoreOptions<Hotel> options, @Nonnull String collectionName) {
84-
return new RedisVectorStoreRecordCollection<>(new JedisPooled(redisContainer.getRedisURI()), collectionName, RedisVectorStoreOptions.<Hotel>builder()
88+
private RedisVectorStoreRecordCollection<Hotel> buildrecordCollection(@Nonnull RedisVectorStoreRecordCollectionOptions<Hotel> options, @Nonnull String collectionName) {
89+
return new RedisVectorStoreRecordCollection<>(new JedisPooled(redisContainer.getRedisURI()), collectionName, RedisVectorStoreRecordCollectionOptions.<Hotel>builder()
8590
.withRecordClass(options.getRecordClass())
8691
.withVectorStoreRecordMapper(options.getVectorStoreRecordMapper())
8792
.withRecordDefinition(options.getRecordDefinition())
88-
.withPrefixCollectionName(options.prefixCollectionName())
93+
.withPrefixCollectionName(options.isPrefixCollectionName())
8994
.build());
9095
}
9196

92-
@ParameterizedTest
93-
@EnumSource(Options.class)
94-
public void buildRecordStore(Options options) {
95-
assertNotNull(buildRecordStore(optionsMap.get(options), "buildTest"));
96-
}
97-
9897
private List<Hotel> getHotels() {
9998
return List.of(
10099
new Hotel("id_1", "Hotel 1", 1, "Hotel 1 description", Arrays.asList(1.0f, 2.0f, 3.0f), 4.0),
@@ -105,37 +104,65 @@ private List<Hotel> getHotels() {
105104
);
106105
}
107106

107+
@Order(1)
108+
@ParameterizedTest
109+
@EnumSource(RecordCollectionOptions.class)
110+
public void buildrecordCollection(RecordCollectionOptions options) {
111+
assertNotNull(buildrecordCollection(optionsMap.get(options), options.name()));
112+
}
113+
114+
@Order(2)
115+
@ParameterizedTest
116+
@EnumSource(RecordCollectionOptions.class)
117+
public void createCollectionAsync(RecordCollectionOptions options) {
118+
RedisVectorStoreRecordCollection<Hotel> recordCollection = buildrecordCollection(optionsMap.get(options), options.name());
119+
120+
assertEquals(false, recordCollection.collectionExistsAsync().block());
121+
recordCollection.createCollectionAsync().block();
122+
assertEquals(true, recordCollection.collectionExistsAsync().block());
123+
}
124+
125+
@Test
126+
public void deleteCollectionAsync() {
127+
RedisVectorStoreRecordCollection<Hotel> recordCollection = buildrecordCollection(optionsMap.get(RecordCollectionOptions.DEFAULT), "deleteCollectionAsync");
128+
129+
assertEquals(false, recordCollection.collectionExistsAsync().block());
130+
recordCollection.createCollectionAsync().block();
131+
recordCollection.deleteCollectionAsync().block();
132+
assertEquals(false, recordCollection.collectionExistsAsync().block());
133+
}
134+
108135
@ParameterizedTest
109-
@EnumSource(Options.class)
110-
public void upsertAndGetRecordAsync(Options options) {
111-
RedisVectorStoreRecordCollection<Hotel> recordStore = buildRecordStore(optionsMap.get(options), "upsertAndGetRecordAsync");
136+
@EnumSource(RecordCollectionOptions.class)
137+
public void upsertAndGetRecordAsync(RecordCollectionOptions options) {
138+
RedisVectorStoreRecordCollection<Hotel> recordCollection = buildrecordCollection(optionsMap.get(options), options.name());
112139

113140
List<Hotel> hotels = getHotels();
114141
for (Hotel hotel : hotels) {
115-
recordStore.upsertAsync(hotel, null).block();
142+
recordCollection.upsertAsync(hotel, null).block();
116143
}
117144

118145
for (Hotel hotel : hotels) {
119-
Hotel retrievedHotel = recordStore.getAsync(hotel.getId(), null).block();
146+
Hotel retrievedHotel = recordCollection.getAsync(hotel.getId(), null).block();
120147
assertNotNull(retrievedHotel);
121148
assertEquals(hotel.getId(), retrievedHotel.getId());
122149
}
123150
}
124151

125152
@ParameterizedTest
126-
@EnumSource(Options.class)
127-
public void getBatchAsync(Options options) {
128-
RedisVectorStoreRecordCollection<Hotel> recordStore = buildRecordStore(optionsMap.get(options), "getBatchAsync");
153+
@EnumSource(RecordCollectionOptions.class)
154+
public void getBatchAsync(RecordCollectionOptions options) {
155+
RedisVectorStoreRecordCollection<Hotel> recordCollection = buildrecordCollection(optionsMap.get(options), options.name());
129156

130157
List<Hotel> hotels = getHotels();
131158
for (Hotel hotel : hotels) {
132-
recordStore.upsertAsync(hotel, null).block();
159+
recordCollection.upsertAsync(hotel, null).block();
133160
}
134161

135162
List<String> ids = new ArrayList<>();
136163
hotels.forEach(hotel -> ids.add(hotel.getId()));
137164

138-
List<Hotel> retrievedHotels = recordStore.getBatchAsync(ids, null).block();
165+
List<Hotel> retrievedHotels = recordCollection.getBatchAsync(ids, null).block();
139166

140167
assertNotNull(retrievedHotels);
141168
assertEquals(hotels.size(), retrievedHotels.size());
@@ -145,15 +172,15 @@ public void getBatchAsync(Options options) {
145172
}
146173

147174
@ParameterizedTest
148-
@EnumSource(Options.class)
149-
public void upsertBatchAsync(Options options) {
150-
RedisVectorStoreRecordCollection<Hotel> recordStore = buildRecordStore(optionsMap.get(options), "upsertBatchAsync");
175+
@EnumSource(RecordCollectionOptions.class)
176+
public void upsertBatchAsync(RecordCollectionOptions options) {
177+
RedisVectorStoreRecordCollection<Hotel> recordCollection = buildrecordCollection(optionsMap.get(options), options.name());
151178

152179
List<Hotel> hotels = getHotels();
153-
List<String> keys = recordStore.upsertBatchAsync(hotels, null).block();
180+
List<String> keys = recordCollection.upsertBatchAsync(hotels, null).block();
154181
assertNotNull(keys);
155182

156-
List<Hotel> retrievedHotels = (List<Hotel>) recordStore.getBatchAsync(keys, null).block();
183+
List<Hotel> retrievedHotels = (List<Hotel>) recordCollection.getBatchAsync(keys, null).block();
157184

158185
assertNotNull(retrievedHotels);
159186
assertEquals(hotels.size(), retrievedHotels.size());
@@ -163,66 +190,68 @@ public void upsertBatchAsync(Options options) {
163190
}
164191

165192
@ParameterizedTest
166-
@EnumSource(Options.class)
167-
public void deleteAsync(Options options) {
168-
RedisVectorStoreRecordCollection<Hotel> recordStore = buildRecordStore(optionsMap.get(options), "deleteAsync");
193+
@EnumSource(RecordCollectionOptions.class)
194+
public void deleteAsync(RecordCollectionOptions options) {
195+
RedisVectorStoreRecordCollection<Hotel> recordCollection = buildrecordCollection(optionsMap.get(options), options.name());
169196

170197
List<Hotel> hotels = getHotels();
171-
recordStore.upsertBatchAsync(hotels, null).block();
198+
recordCollection.upsertBatchAsync(hotels, null).block();
172199

173200
for (Hotel hotel : hotels) {
174-
recordStore.deleteAsync(hotel.getId(), null).block();
175-
Hotel retrievedHotel = recordStore.getAsync(hotel.getId(), null).block();
201+
recordCollection.deleteAsync(hotel.getId(), null).block();
202+
Hotel retrievedHotel = recordCollection.getAsync(hotel.getId(), null).block();
176203
assertNull(retrievedHotel);
177204
}
178205
}
179206

180207
@ParameterizedTest
181-
@EnumSource(Options.class)
182-
public void deleteBatchAsync(Options options) {
183-
RedisVectorStoreRecordCollection<Hotel> recordStore = buildRecordStore(optionsMap.get(options), "deleteBatchAsync");
208+
@EnumSource(RecordCollectionOptions.class)
209+
public void deleteBatchAsync(RecordCollectionOptions options) {
210+
RedisVectorStoreRecordCollection<Hotel> recordCollection = buildrecordCollection(optionsMap.get(options), options.name());
184211

185212
List<Hotel> hotels = getHotels();
186-
recordStore.upsertBatchAsync(hotels, null).block();
213+
recordCollection.upsertBatchAsync(hotels, null).block();
187214

188215
List<String> ids = new ArrayList<>();
189216
hotels.forEach(hotel -> ids.add(hotel.getId()));
190217

191-
recordStore.deleteBatchAsync(ids, null).block();
218+
recordCollection.deleteBatchAsync(ids, null).block();
192219

193220
for (String id : ids) {
194-
Hotel retrievedHotel = recordStore.getAsync(id, null).block();
221+
Hotel retrievedHotel = recordCollection.getAsync(id, null).block();
195222
assertNull(retrievedHotel);
196223
}
197224
}
198225

199-
@Test
200-
public void getAsyncWithVectors() {
201-
RedisVectorStoreRecordCollection<Hotel> recordStore = buildRecordStore(optionsMap.get(Options.DEFAULT), "getAsyncWithVectors");
226+
@ParameterizedTest
227+
@EnumSource(RecordCollectionOptions.class)
228+
public void getAsyncWithVectors(RecordCollectionOptions options) {
229+
RedisVectorStoreRecordCollection<Hotel> recordCollection = buildrecordCollection(optionsMap.get(options), options.name());
202230

203231
List<Hotel> hotels = getHotels();
204-
recordStore.upsertBatchAsync(hotels, null).block();
232+
recordCollection.upsertBatchAsync(hotels, null).block();
205233

206234
for (Hotel hotel : hotels) {
207-
Hotel retrievedHotel = recordStore.getAsync(hotel.getId(), null).block();
235+
Hotel retrievedHotel = recordCollection.getAsync(hotel.getId(), null).block();
208236
assertNotNull(retrievedHotel);
209237
assertNotNull(retrievedHotel.getDescriptionEmbedding());
210238
assertEquals(hotel.getId(), retrievedHotel.getId());
211239
assertEquals(hotel.getDescription(), retrievedHotel.getDescription());
212240
}
213241
}
214242

215-
@Test
216-
public void getBatchAsyncWithVectors() {
217-
RedisVectorStoreRecordCollection<Hotel> recordStore = buildRecordStore(optionsMap.get(Options.DEFAULT), "getBatchAsyncWithVectors");
243+
@ParameterizedTest
244+
@EnumSource(RecordCollectionOptions.class)
245+
public void getBatchAsyncWithVectors(RecordCollectionOptions options) {
246+
RedisVectorStoreRecordCollection<Hotel> recordCollection = buildrecordCollection(optionsMap.get(options), options.name());
218247

219248
List<Hotel> hotels = getHotels();
220-
recordStore.upsertBatchAsync(hotels, null).block();
249+
recordCollection.upsertBatchAsync(hotels, null).block();
221250

222251
List<String> ids = new ArrayList<>();
223252
hotels.forEach(hotel -> ids.add(hotel.getId()));
224253

225-
List<Hotel> retrievedHotels = recordStore.getBatchAsync(ids, null).block();
254+
List<Hotel> retrievedHotels = recordCollection.getBatchAsync(ids, null).block();
226255

227256
assertNotNull(retrievedHotels);
228257
assertEquals(hotels.size(), retrievedHotels.size());
@@ -233,35 +262,37 @@ public void getBatchAsyncWithVectors() {
233262
}
234263
}
235264

236-
@Test
237-
public void getAsyncWithNoVectors() {
238-
RedisVectorStoreRecordCollection<Hotel> recordStore = buildRecordStore(optionsMap.get(Options.WITH_CUSTOM_DEFINITION), "getAsyncWithNoVectors");
265+
@ParameterizedTest
266+
@EnumSource(RecordCollectionOptions.class)
267+
public void getAsyncWithNoVectors(RecordCollectionOptions options) {
268+
RedisVectorStoreRecordCollection<Hotel> recordCollection = buildrecordCollection(optionsMap.get(options), options.name());
239269

240270
List<Hotel> hotels = getHotels();
241-
recordStore.upsertBatchAsync(hotels, null).block();
271+
recordCollection.upsertBatchAsync(hotels, null).block();
242272

243273
GetRecordOptions getRecordOptions = GetRecordOptions.builder().includeVectors(false).build();
244274
for (Hotel hotel : hotels) {
245-
Hotel retrievedHotel = recordStore.getAsync(hotel.getId(), getRecordOptions).block();
275+
Hotel retrievedHotel = recordCollection.getAsync(hotel.getId(), getRecordOptions).block();
246276
assertNotNull(retrievedHotel);
247277
assertNull(retrievedHotel.getDescriptionEmbedding());
248278
assertEquals(hotel.getId(), retrievedHotel.getId());
249279
assertEquals(hotel.getDescription(), retrievedHotel.getDescription());
250280
}
251281
}
252282

253-
@Test
254-
public void getBatchAsyncWithNoVectors() {
255-
RedisVectorStoreRecordCollection<Hotel> recordStore = buildRecordStore(optionsMap.get(Options.WITH_CUSTOM_DEFINITION), "getBatchAsyncWithNoVectors");
283+
@ParameterizedTest
284+
@EnumSource(RecordCollectionOptions.class)
285+
public void getBatchAsyncWithNoVectors(RecordCollectionOptions options) {
286+
RedisVectorStoreRecordCollection<Hotel> recordCollection = buildrecordCollection(optionsMap.get(options), options.name());
256287

257288
List<Hotel> hotels = getHotels();
258-
recordStore.upsertBatchAsync(hotels, null).block();
289+
recordCollection.upsertBatchAsync(hotels, null).block();
259290

260291
GetRecordOptions getRecordOptions = GetRecordOptions.builder().includeVectors(false).build();
261292
List<String> ids = new ArrayList<>();
262293
hotels.forEach(hotel -> ids.add(hotel.getId()));
263294

264-
List<Hotel> retrievedHotels = recordStore.getBatchAsync(ids, getRecordOptions).block();
295+
List<Hotel> retrievedHotels = recordCollection.getBatchAsync(ids, getRecordOptions).block();
265296

266297
assertNotNull(retrievedHotels);
267298
assertEquals(hotels.size(), retrievedHotels.size());

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

+50
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
package com.microsoft.semantickernel.tests.connectors.memory.redis;
2+
3+
import com.microsoft.semantickernel.connectors.memory.redis.RedisVectorStore;
4+
import com.microsoft.semantickernel.connectors.memory.redis.RedisVectorStoreOptions;
5+
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition;
6+
import com.microsoft.semantickernel.tests.connectors.memory.Hotel;
7+
import com.redis.testcontainers.RedisContainer;
8+
import org.junit.jupiter.api.BeforeAll;
9+
import org.junit.jupiter.api.Test;
10+
import org.testcontainers.junit.jupiter.Container;
11+
import org.testcontainers.junit.jupiter.Testcontainers;
12+
import reactor.core.publisher.Mono;
13+
import redis.clients.jedis.JedisPooled;
14+
15+
import java.util.ArrayList;
16+
import java.util.Arrays;
17+
import java.util.List;
18+
19+
import static org.junit.jupiter.api.Assertions.assertEquals;
20+
import static org.junit.jupiter.api.Assertions.assertNotNull;
21+
import static org.junit.jupiter.api.Assertions.assertTrue;
22+
23+
@Testcontainers
24+
public class RedisVectorStoreTest {
25+
@Container
26+
private static final RedisContainer redisContainer = new RedisContainer("redis/redis-stack:latest");
27+
private static JedisPooled jedis;
28+
29+
@BeforeAll
30+
public static void setUp() {
31+
jedis = new JedisPooled(redisContainer.getRedisURI());
32+
}
33+
34+
@Test
35+
public void getCollectionNamesAsync() {
36+
RedisVectorStore<Hotel> vectorStore = new RedisVectorStore<>(jedis, new RedisVectorStoreOptions<>(Hotel.class, null));
37+
List<String> collectionNames = Arrays.asList("collection1", "collection2", "collection3");
38+
39+
for (String collectionName : collectionNames) {
40+
vectorStore.getCollection(collectionName, VectorStoreRecordDefinition.fromRecordClass(Hotel.class)).createCollectionAsync().block();
41+
}
42+
43+
List<String> retrievedCollectionNames = vectorStore.getCollectionNamesAsync().block();
44+
assertNotNull(retrievedCollectionNames);
45+
assertEquals(collectionNames.size(), retrievedCollectionNames.size());
46+
for (String collectionName : collectionNames) {
47+
assertTrue(retrievedCollectionNames.contains(collectionName));
48+
}
49+
}
50+
}

‎semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/memory/azureaisearch/AzureAISearchVectorStoreRecordCollection.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,7 @@ public AzureAISearchVectorStoreRecordCollection(
9292
// Validate supported types
9393
VectorStoreRecordDefinition.validateSupportedTypes(
9494
this.options.getRecordClass(), this.recordDefinition, supportedKeyTypes,
95-
supportedDataTypes, supportedVectorTypes);
95+
supportedVectorTypes, supportedDataTypes);
9696

9797
// Add non-vector fields to the list
9898
nonVectorFields.add(this.recordDefinition.getKeyField().getName());

0 commit comments

Comments
 (0)