Skip to content

Commit 38c0b70

Browse files
author
David Grieve
authored
Merge pull request #195 from milderhc/get-collection
Add VectorStoreRecordCollectionOptions interface and make getCollecti…
2 parents 9bf0833 + c213a08 commit 38c0b70

18 files changed

Lines changed: 208 additions & 135 deletions

File tree

‎api-test/integration-tests/src/test/java/com/microsoft/semantickernel/tests/connectors/memory/jdbc/JDBCVectorStoreTest.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStore;
44
import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreOptions;
55
import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreQueryProvider;
6+
import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreRecordCollectionOptions;
67
import com.microsoft.semantickernel.connectors.data.mysql.MySQLVectorStoreQueryProvider;
78
import com.microsoft.semantickernel.connectors.data.postgres.PostgreSQLVectorStoreQueryProvider;
89
import com.microsoft.semantickernel.tests.connectors.memory.Hotel;
@@ -88,7 +89,10 @@ public void getCollectionNamesAsync(QueryProvider provider) {
8889
List<String> collectionNames = Arrays.asList("collection1", "collection2", "collection3");
8990

9091
for (String collectionName : collectionNames) {
91-
vectorStore.getCollection(collectionName, Hotel.class, null).createCollectionAsync().block();
92+
vectorStore.getCollection(collectionName,
93+
JDBCVectorStoreRecordCollectionOptions.<Hotel>builder()
94+
.withRecordClass(Hotel.class)
95+
.build()).createCollectionAsync().block();
9296
}
9397

9498
List<String> retrievedCollectionNames = vectorStore.getCollectionNamesAsync().block();

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

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
11
package com.microsoft.semantickernel.tests.connectors.memory.redis;
22

3+
import com.microsoft.semantickernel.connectors.data.redis.RedisHashSetVectorStoreRecordCollectionOptions;
4+
import com.microsoft.semantickernel.connectors.data.redis.RedisJsonVectorStoreRecordCollectionOptions;
35
import com.microsoft.semantickernel.connectors.data.redis.RedisStorageType;
46
import com.microsoft.semantickernel.connectors.data.redis.RedisVectorStore;
57
import com.microsoft.semantickernel.connectors.data.redis.RedisVectorStoreOptions;
8+
import com.microsoft.semantickernel.data.VectorStoreRecordCollectionOptions;
69
import com.microsoft.semantickernel.tests.connectors.memory.Hotel;
710
import com.redis.testcontainers.RedisContainer;
811
import org.junit.jupiter.api.BeforeAll;
@@ -36,6 +39,18 @@ public static JedisPooled buildClient(RedisStorageType storageType) {
3639
}
3740
}
3841

42+
private static VectorStoreRecordCollectionOptions<String, Hotel> getRecordCollectionOptions(RedisStorageType storageType) {
43+
if (storageType == RedisStorageType.JSON) {
44+
return RedisJsonVectorStoreRecordCollectionOptions.<Hotel>builder()
45+
.withRecordClass(Hotel.class)
46+
.build();
47+
} else {
48+
return RedisHashSetVectorStoreRecordCollectionOptions.<Hotel>builder()
49+
.withRecordClass(Hotel.class)
50+
.build();
51+
}
52+
}
53+
3954
@ParameterizedTest
4055
@EnumSource(RedisStorageType.class)
4156
public void getCollectionNamesAsync(RedisStorageType storageType) {
@@ -46,7 +61,7 @@ public void getCollectionNamesAsync(RedisStorageType storageType) {
4661
List<String> collectionNames = Arrays.asList("collection1", "collection2", "collection3");
4762

4863
for (String collectionName : collectionNames) {
49-
vectorStore.getCollection(collectionName, Hotel.class, null).createCollectionAsync().block();
64+
vectorStore.getCollection(collectionName, getRecordCollectionOptions(storageType)).createCollectionAsync().block();
5065
}
5166

5267
List<String> retrievedCollectionNames = vectorStore.getCollectionNamesAsync().block();

‎samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/AzureAISearch_DataStorage.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import com.microsoft.semantickernel.aiservices.openai.textembedding.OpenAITextEmbeddingGenerationService;
1616
import com.microsoft.semantickernel.connectors.data.azureaisearch.AzureAISearchVectorStore;
1717
import com.microsoft.semantickernel.connectors.data.azureaisearch.AzureAISearchVectorStoreOptions;
18+
import com.microsoft.semantickernel.connectors.data.azureaisearch.AzureAISearchVectorStoreRecordCollectionOptions;
1819
import com.microsoft.semantickernel.data.VectorStoreRecordCollection;
1920
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordDataAttribute;
2021
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordKeyAttribute;
@@ -125,8 +126,9 @@ public static void dataStorageWithAzureAISearch(
125126
String collectionName = "skgithubfiles";
126127
var collection = azureAISearchVectorStore.getCollection(
127128
collectionName,
128-
GitHubFile.class,
129-
null);
129+
AzureAISearchVectorStoreRecordCollectionOptions.<GitHubFile>builder()
130+
.withRecordClass(GitHubFile.class)
131+
.build());
130132

131133
// Create collection if it does not exist and store data
132134
collection

‎samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/InMemory_DataStorage.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import com.microsoft.semantickernel.aiservices.openai.textembedding.OpenAITextEmbeddingGenerationService;
99
import com.microsoft.semantickernel.data.VectorStoreRecordCollection;
1010
import com.microsoft.semantickernel.data.VolatileVectorStore;
11+
import com.microsoft.semantickernel.data.VolatileVectorStoreRecordCollectionOptions;
1112
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordDataAttribute;
1213
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordKeyAttribute;
1314
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordVectorAttribute;
@@ -101,7 +102,10 @@ public static void inMemoryDataStorage(
101102
var volatileVectorStore = new VolatileVectorStore();
102103

103104
String collectionName = "skgithubfiles";
104-
var collection = volatileVectorStore.getCollection(collectionName, GitHubFile.class, null);
105+
var collection = volatileVectorStore.getCollection(collectionName,
106+
VolatileVectorStoreRecordCollectionOptions.<GitHubFile>builder()
107+
.withRecordClass(GitHubFile.class)
108+
.build());
105109

106110
// Create collection if it does not exist and store data
107111
List<String> ids = collection

‎samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/JDBC_DataStorage.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import com.microsoft.semantickernel.aiservices.openai.textembedding.OpenAITextEmbeddingGenerationService;
99
import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStore;
1010
import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreOptions;
11+
import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreRecordCollectionOptions;
1112
import com.microsoft.semantickernel.connectors.data.mysql.MySQLVectorStoreQueryProvider;
1213
import com.microsoft.semantickernel.data.VectorStoreRecordCollection;
1314
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordDataAttribute;
@@ -132,8 +133,9 @@ public static void dataStorageWithMySQL(
132133

133134
String collectionName = "skgithubfiles";
134135
var collection = jdbcVectorStore.getCollection(collectionName,
135-
GitHubFile.class,
136-
null);
136+
JDBCVectorStoreRecordCollectionOptions.<GitHubFile>builder()
137+
.withRecordClass(GitHubFile.class)
138+
.build());
137139

138140
// Create collection if it does not exist and store data
139141
List<String> ids = collection

‎samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/Redis_DataStorage.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import com.azure.core.util.MetricsOptions;
1010
import com.azure.core.util.TracingOptions;
1111
import com.microsoft.semantickernel.aiservices.openai.textembedding.OpenAITextEmbeddingGenerationService;
12+
import com.microsoft.semantickernel.connectors.data.redis.RedisJsonVectorStoreRecordCollectionOptions;
1213
import com.microsoft.semantickernel.connectors.data.redis.RedisVectorStore;
1314
import com.microsoft.semantickernel.connectors.data.redis.RedisVectorStoreOptions;
1415
import com.microsoft.semantickernel.data.VectorStoreRecordCollection;
@@ -120,7 +121,10 @@ public static void dataStorageWithRedis(
120121
.build();
121122

122123
String collectionName = "skgithubfiles";
123-
var collection = vectorStore.getCollection(collectionName, GitHubFile.class, null);
124+
var collection = vectorStore.getCollection(collectionName,
125+
RedisJsonVectorStoreRecordCollectionOptions.<GitHubFile>builder()
126+
.withRecordClass(GitHubFile.class)
127+
.build());
124128

125129
// Create collection if it does not exist and store data
126130
List<String> ids = collection

‎semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStore.java‎

Lines changed: 16 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,9 @@
55
import com.azure.search.documents.indexes.models.SearchIndex;
66
import com.microsoft.semantickernel.data.VectorStore;
77
import com.microsoft.semantickernel.data.VectorStoreRecordCollection;
8+
import com.microsoft.semantickernel.data.VectorStoreRecordCollectionOptions;
89
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition;
10+
import com.microsoft.semantickernel.exceptions.SKException;
911
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
1012
import java.util.List;
1113
import javax.annotation.Nonnull;
@@ -34,44 +36,34 @@ public AzureAISearchVectorStore(@Nonnull SearchIndexAsyncClient client,
3436
* Gets a new instance of {@link AzureAISearchVectorStoreRecordCollection}
3537
*
3638
* @param collectionName The name of the collection.
37-
* @param recordClass The class type of the record.
38-
* @param recordDefinition The record definition.
39+
* @param options The options for the collection.
3940
* @return The collection.
4041
*/
4142
@Override
4243
public final <Key, Record> VectorStoreRecordCollection<Key, Record> getCollection(
4344
@Nonnull String collectionName,
44-
@Nonnull Class<Key> keyClass,
45-
@Nonnull Class<Record> recordClass,
46-
@Nullable VectorStoreRecordDefinition recordDefinition) {
47-
if (!keyClass.equals(String.class)) {
48-
throw new IllegalArgumentException("Azure AI Search only supports string keys");
45+
@Nonnull VectorStoreRecordCollectionOptions<Key, Record> options) {
46+
if (!options.getKeyClass().equals(String.class)) {
47+
throw new SKException("Azure AI Search only supports string keys");
48+
}
49+
if (options.getRecordClass() == null) {
50+
throw new SKException("Record class is required");
4951
}
5052

51-
return (VectorStoreRecordCollection<Key, Record>) getCollection(
52-
collectionName, recordClass, recordDefinition);
53-
}
54-
55-
public <Record> AzureAISearchVectorStoreRecordCollection<Record> getCollection(
56-
@Nonnull String collectionName,
57-
@Nonnull Class<Record> recordClass,
58-
@Nullable VectorStoreRecordDefinition recordDefinition) {
59-
if (options.getVectorStoreRecordCollectionFactory() != null) {
60-
return options.getVectorStoreRecordCollectionFactory()
53+
if (this.options.getVectorStoreRecordCollectionFactory() != null) {
54+
return (VectorStoreRecordCollection<Key, Record>) this.options
55+
.getVectorStoreRecordCollectionFactory()
6156
.createVectorStoreRecordCollection(
6257
client,
6358
collectionName,
64-
recordClass,
65-
recordDefinition);
59+
options.getRecordClass(),
60+
options.getRecordDefinition());
6661
}
6762

68-
return new AzureAISearchVectorStoreRecordCollection<>(
63+
return (VectorStoreRecordCollection<Key, Record>) new AzureAISearchVectorStoreRecordCollection<>(
6964
client,
7065
collectionName,
71-
AzureAISearchVectorStoreRecordCollectionOptions.<Record>builder()
72-
.withRecordClass(recordClass)
73-
.withRecordDefinition(recordDefinition)
74-
.build());
66+
(AzureAISearchVectorStoreRecordCollectionOptions<Record>) options);
7567
}
7668

7769
/**

‎semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreRecordCollectionOptions.java‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
package com.microsoft.semantickernel.connectors.data.azureaisearch;
33

44
import com.azure.search.documents.SearchDocument;
5+
import com.microsoft.semantickernel.data.VectorStoreRecordCollectionOptions;
56
import com.microsoft.semantickernel.data.VectorStoreRecordMapper;
67
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition;
78
import javax.annotation.Nonnull;
@@ -12,12 +13,11 @@
1213
*
1314
* @param <Record> the record type
1415
*/
15-
public class AzureAISearchVectorStoreRecordCollectionOptions<Record> {
16-
16+
public class AzureAISearchVectorStoreRecordCollectionOptions<Record>
17+
implements VectorStoreRecordCollectionOptions<String, Record> {
1718
private final Class<Record> recordClass;
1819
@Nullable
1920
private final VectorStoreRecordMapper<Record, SearchDocument> vectorStoreRecordMapper;
20-
2121
@Nullable
2222
private final VectorStoreRecordDefinition recordDefinition;
2323

@@ -31,6 +31,16 @@ public static <Record> Builder<Record> builder() {
3131
return new Builder<>();
3232
}
3333

34+
/**
35+
* Gets the key class.
36+
*
37+
* @return the key class
38+
*/
39+
@Override
40+
public Class<String> getKeyClass() {
41+
return String.class;
42+
}
43+
3444
/**
3545
* Gets the record class.
3646
*

‎semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/JDBCVectorStore.java‎

Lines changed: 23 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,9 @@
22
package com.microsoft.semantickernel.connectors.data.jdbc;
33

44
import com.microsoft.semantickernel.data.VectorStoreRecordCollection;
5+
import com.microsoft.semantickernel.data.VectorStoreRecordCollectionOptions;
56
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition;
7+
import com.microsoft.semantickernel.exceptions.SKException;
68
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
79
import reactor.core.publisher.Mono;
810
import reactor.core.scheduler.Schedulers;
@@ -55,53 +57,42 @@ public static Builder builder() {
5557
* Gets a collection from the vector store.
5658
*
5759
* @param collectionName The name of the collection.
58-
* @param recordClass The class type of the record.
59-
* @param recordDefinition The record definition.
60+
* @param options The options for the collection.
6061
* @return The collection.
6162
*/
6263
@Override
6364
public <Key, Record> VectorStoreRecordCollection<Key, Record> getCollection(
64-
@Nonnull String collectionName, @Nonnull Class<Key> keyClass,
65-
@Nonnull Class<Record> recordClass,
66-
@Nullable VectorStoreRecordDefinition recordDefinition) {
67-
if (keyClass != String.class) {
68-
throw new IllegalArgumentException("Redis only supports string keys");
65+
@Nonnull String collectionName,
66+
@Nonnull VectorStoreRecordCollectionOptions<Key, Record> options) {
67+
if (!options.getKeyClass().equals(String.class)) {
68+
throw new SKException("JDBC only supports string keys");
69+
}
70+
if (options.getRecordClass() == null) {
71+
throw new SKException("Record class is required");
6972
}
7073

71-
return (VectorStoreRecordCollection<Key, Record>) getCollection(
72-
collectionName,
73-
recordClass,
74-
recordDefinition);
75-
}
76-
77-
/**
78-
* Gets a collection from the vector store.
79-
*
80-
* @param collectionName The name of the collection.
81-
* @param recordClass The class type of the record.
82-
* @param recordDefinition The record definition.
83-
* @return The collection.
84-
*/
85-
public <Record> JDBCVectorStoreRecordCollection<Record> getCollection(
86-
@Nonnull String collectionName,
87-
@Nonnull Class<Record> recordClass,
88-
@Nullable VectorStoreRecordDefinition recordDefinition) {
8974
if (this.options != null && this.options.getVectorStoreRecordCollectionFactory() != null) {
90-
return this.options.getVectorStoreRecordCollectionFactory()
75+
return (VectorStoreRecordCollection<Key, Record>) this.options
76+
.getVectorStoreRecordCollectionFactory()
9177
.createVectorStoreRecordCollection(
9278
dataSource,
9379
collectionName,
94-
recordClass,
95-
recordDefinition);
80+
options.getRecordClass(),
81+
options.getRecordDefinition());
9682
}
9783

98-
return new JDBCVectorStoreRecordCollection<>(
84+
JDBCVectorStoreRecordCollectionOptions<Record> jdbcOptions = (JDBCVectorStoreRecordCollectionOptions<Record>) options;
85+
return (VectorStoreRecordCollection<Key, Record>) new JDBCVectorStoreRecordCollection<>(
9986
dataSource,
10087
collectionName,
10188
JDBCVectorStoreRecordCollectionOptions.<Record>builder()
102-
.withRecordClass(recordClass)
103-
.withRecordDefinition(recordDefinition)
104-
.withQueryProvider(this.queryProvider)
89+
.withCollectionsTableName(jdbcOptions.getCollectionsTableName())
90+
.withPrefixForCollectionTables(jdbcOptions.getPrefixForCollectionTables())
91+
.withQueryProvider(jdbcOptions.getQueryProvider() == null ? queryProvider
92+
: jdbcOptions.getQueryProvider())
93+
.withRecordClass(jdbcOptions.getRecordClass())
94+
.withRecordDefinition(jdbcOptions.getRecordDefinition())
95+
.withVectorStoreRecordMapper(jdbcOptions.getVectorStoreRecordMapper())
10596
.build());
10697
}
10798

‎semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/JDBCVectorStoreRecordCollectionOptions.java‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
// Copyright (c) Microsoft. All rights reserved.
22
package com.microsoft.semantickernel.connectors.data.jdbc;
33

4+
import com.microsoft.semantickernel.data.VectorStoreRecordCollectionOptions;
45
import com.microsoft.semantickernel.data.VectorStoreRecordMapper;
56
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition;
67
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
@@ -11,7 +12,8 @@
1112
import static com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreQueryProvider.DEFAULT_COLLECTIONS_TABLE;
1213
import static com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreQueryProvider.DEFAULT_PREFIX_FOR_COLLECTION_TABLES;
1314

14-
public class JDBCVectorStoreRecordCollectionOptions<Record> {
15+
public class JDBCVectorStoreRecordCollectionOptions<Record>
16+
implements VectorStoreRecordCollectionOptions<String, Record> {
1517
private final Class<Record> recordClass;
1618
private final VectorStoreRecordMapper<Record, ResultSet> vectorStoreRecordMapper;
1719
private final VectorStoreRecordDefinition recordDefinition;
@@ -43,6 +45,16 @@ public static <Record> Builder<Record> builder() {
4345
return new Builder<>();
4446
}
4547

48+
/**
49+
* Gets the key class.
50+
*
51+
* @return the key class
52+
*/
53+
@Override
54+
public Class<String> getKeyClass() {
55+
return String.class;
56+
}
57+
4658
/**
4759
* Gets the record class.
4860
* @return the record class

0 commit comments

Comments
 (0)