Skip to content

Commit 734a40e

Browse files
author
Milder Hernandez
authored
Merge pull request #93 from milderhc/azure-ai-search
Add AzureAISearchVectorStore and supported types validation
2 parents ce68d94 + 6994f27 commit 734a40e

11 files changed

Lines changed: 405 additions & 129 deletions

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ static void setup() {
7272
.withName("rating")
7373
.withFieldType(Double.class)
7474
.build());
75-
VectorStoreRecordDefinition recordDefinition = VectorStoreRecordDefinition.create(fields);
75+
VectorStoreRecordDefinition recordDefinition = VectorStoreRecordDefinition.fromFields(fields);
7676

7777
optionsMap.put(Options.WITH_CUSTOM_DEFINITION, RedisVectorStoreOptions.<Hotel>builder()
7878
.withRecordClass(Hotel.class)

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

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,9 @@
1111
import com.azure.search.documents.indexes.SearchIndexAsyncClient;
1212
import com.azure.search.documents.indexes.SearchIndexClientBuilder;
1313
import com.microsoft.semantickernel.aiservices.openai.textembedding.OpenAITextEmbeddingGenerationService;
14+
import com.microsoft.semantickernel.connectors.memory.azureaisearch.AzureAISearchVectorStore;
1415
import com.microsoft.semantickernel.connectors.memory.azureaisearch.AzureAISearchVectorStoreOptions;
16+
import com.microsoft.semantickernel.connectors.memory.azureaisearch.AzureAISearchVectorStoreRecordCollectionOptions;
1517
import com.microsoft.semantickernel.connectors.memory.azureaisearch.AzureAISearchVectorStoreRecordCollection;
1618
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordDataAttribute;
1719
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordKeyAttribute;
@@ -41,7 +43,6 @@ public class AzureAISearch {
4143
private static final String AZURE_AISEARCH_KEY = System.getenv("AZURE_AISEARCH_KEY");
4244
private static final String MODEL_ID = System.getenv()
4345
.getOrDefault("EMBEDDING_MODEL_ID", "text-embedding-3-large");
44-
private static final String COLLECTION_NAME = "skgithubtest";
4546
private static final int EMBEDDING_DIMENSIONS = 1536;
4647

4748
static class GitHubFile {
@@ -113,20 +114,24 @@ public static void documentSearchWithAzureAISearch(
113114
SearchIndexAsyncClient searchClient,
114115
OpenAITextEmbeddingGenerationService embeddingGeneration) {
115116

116-
var azureAISearchVectorStoreCollection = new AzureAISearchVectorStoreRecordCollection<GitHubFile>(
117-
searchClient,
118-
COLLECTION_NAME,
119-
AzureAISearchVectorStoreOptions.<GitHubFile>builder()
117+
// Create a new Azure AI Search vector store
118+
var azureAISearchVectorStore = new AzureAISearchVectorStore<>(searchClient,
119+
AzureAISearchVectorStoreOptions.<GitHubFile>builder()
120120
.withRecordClass(GitHubFile.class)
121121
.build());
122122

123-
azureAISearchVectorStoreCollection.createCollectionIfNotExistsAsync()
124-
.then(storeData(azureAISearchVectorStoreCollection, embeddingGeneration, sampleData()))
123+
String collectionName = "skgithubfiles";
124+
var collection = azureAISearchVectorStore.getCollection(collectionName, null);
125+
126+
// Create collection if it does not exist and store data
127+
collection
128+
.createCollectionIfNotExistsAsync()
129+
.then(storeData(collection, embeddingGeneration, sampleData()))
125130
.block();
126131

127132
// Query the Azure AI Search client for results
128133
// This might take a few seconds to return the best result
129-
var result = searchClient.getSearchAsyncClient(COLLECTION_NAME)
134+
var result = searchClient.getSearchAsyncClient(collectionName)
130135
.search("How to get started with the Semantic Kernel?")
131136
.blockFirst();
132137

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

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
package com.microsoft.semantickernel.connectors.memory.azureaisearch;
3+
4+
import com.azure.search.documents.indexes.SearchIndexAsyncClient;
5+
import com.azure.search.documents.indexes.models.SearchIndex;
6+
import com.microsoft.semantickernel.data.VectorStore;
7+
import com.microsoft.semantickernel.data.VectorStoreRecordCollection;
8+
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition;
9+
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
10+
import reactor.core.publisher.Mono;
11+
12+
import javax.annotation.Nonnull;
13+
import java.util.List;
14+
15+
public class AzureAISearchVectorStore<Record>
16+
implements VectorStore<String, Record, AzureAISearchVectorStoreRecordCollection<Record>> {
17+
18+
private final SearchIndexAsyncClient client;
19+
private final AzureAISearchVectorStoreOptions<Record> options;
20+
21+
/**
22+
* Creates a new instance of {@link AzureAISearchVectorStore}.
23+
*
24+
* @param client The Azure AI Search client.
25+
* @param options The options for the vector store.
26+
*/
27+
@SuppressFBWarnings("EI_EXPOSE_REP2")
28+
public AzureAISearchVectorStore(@Nonnull SearchIndexAsyncClient client,
29+
@Nonnull AzureAISearchVectorStoreOptions<Record> options) {
30+
this.client = client;
31+
this.options = options;
32+
}
33+
34+
/**
35+
* Gets a new instance of {@link AzureAISearchVectorStoreRecordCollection}
36+
*
37+
* @param collectionName The name of the collection.
38+
* @param recordDefinition The record definition.
39+
* @return The collection.
40+
*/
41+
@Override
42+
public AzureAISearchVectorStoreRecordCollection<Record> getCollection(
43+
@Nonnull String collectionName,
44+
VectorStoreRecordDefinition recordDefinition) {
45+
46+
if (options.getVectorStoreRecordCollectionFactory() != null) {
47+
return options.getVectorStoreRecordCollectionFactory()
48+
.createVectorStoreRecordCollection(
49+
client,
50+
collectionName,
51+
AzureAISearchVectorStoreRecordCollectionOptions.<Record>builder()
52+
.withRecordClass(options.getRecordClass())
53+
.withRecordDefinition(recordDefinition)
54+
.build());
55+
}
56+
57+
return new AzureAISearchVectorStoreRecordCollection<>(client, collectionName,
58+
AzureAISearchVectorStoreRecordCollectionOptions.<Record>builder()
59+
.withRecordClass(options.getRecordClass())
60+
.withRecordDefinition(recordDefinition)
61+
.build());
62+
}
63+
64+
/**
65+
* Gets the names of all collections in the Azure AI Search vector store.
66+
*
67+
* @return A list of collection names.
68+
*/
69+
@Override
70+
public Mono<List<String>> getCollectionNamesAsync() {
71+
return client.listIndexes().map(SearchIndex::getName).collectList();
72+
}
73+
}

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

Lines changed: 34 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,21 @@
11
// Copyright (c) Microsoft. All rights reserved.
22
package com.microsoft.semantickernel.connectors.memory.azureaisearch;
33

4-
import com.azure.search.documents.SearchDocument;
5-
import com.microsoft.semantickernel.data.VectorStoreRecordMapper;
6-
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition;
7-
8-
/**
9-
* Options for an Azure AI Search vector store.
10-
*
11-
* @param <Record> the record type
12-
*/
134
public class AzureAISearchVectorStoreOptions<Record> {
145
private final Class<Record> recordClass;
15-
private final VectorStoreRecordMapper<Record, SearchDocument> vectorStoreRecordMapper;
16-
private final VectorStoreRecordDefinition recordDefinition;
6+
private final AzureAISearchVectorStoreRecordCollectionFactory<Record> vectorStoreRecordCollectionFactory;
7+
8+
/**
9+
* Creates a new instance of the Azure AI Search vector store options.
10+
*
11+
* @param recordClass The record class.
12+
* @param vectorStoreRecordCollectionFactory The vector store record collection factory.
13+
*/
14+
public AzureAISearchVectorStoreOptions(Class<Record> recordClass,
15+
AzureAISearchVectorStoreRecordCollectionFactory<Record> vectorStoreRecordCollectionFactory) {
16+
this.recordClass = recordClass;
17+
this.vectorStoreRecordCollectionFactory = vectorStoreRecordCollectionFactory;
18+
}
1719

1820
/**
1921
* Creates a new builder.
@@ -35,84 +37,54 @@ public Class<Record> getRecordClass() {
3537
}
3638

3739
/**
38-
* Gets the record definition.
40+
* Gets the vector store record collection factory.
3941
*
40-
* @return the record definition
42+
* @return the vector store record collection factory
4143
*/
42-
public VectorStoreRecordDefinition getRecordDefinition() {
43-
return recordDefinition;
44+
public AzureAISearchVectorStoreRecordCollectionFactory<Record> getVectorStoreRecordCollectionFactory() {
45+
return vectorStoreRecordCollectionFactory;
4446
}
4547

4648
/**
47-
* Gets the vector store record mapper.
48-
*
49-
* @return the vector store record mapper
50-
*/
51-
public VectorStoreRecordMapper<Record, SearchDocument> getVectorStoreRecordMapper() {
52-
return vectorStoreRecordMapper;
53-
}
54-
55-
private AzureAISearchVectorStoreOptions(
56-
Class<Record> recordClass,
57-
VectorStoreRecordMapper<Record, SearchDocument> vectorStoreRecordMapper,
58-
VectorStoreRecordDefinition recordDefinition) {
59-
this.recordClass = recordClass;
60-
this.vectorStoreRecordMapper = vectorStoreRecordMapper;
61-
this.recordDefinition = recordDefinition;
62-
}
63-
64-
/**
65-
* Builder for {@link AzureAISearchVectorStoreOptions}.
49+
* Builder for Azure AI Search vector store options.
6650
*
6751
* @param <Record> the record type
6852
*/
6953
public static class Builder<Record> {
70-
private VectorStoreRecordMapper<Record, SearchDocument> vectorStoreRecordMapper;
7154
private Class<Record> recordClass;
72-
private VectorStoreRecordDefinition recordDefinition;
73-
74-
public Builder<Record> withRecordClass(Class<Record> recordClass) {
75-
this.recordClass = recordClass;
76-
return this;
77-
}
55+
private AzureAISearchVectorStoreRecordCollectionFactory<Record> vectorStoreRecordCollectionFactory;
7856

7957
/**
80-
* Sets the vector store record mapper.
58+
* Sets the record class.
8159
*
82-
* @param vectorStoreRecordMapper the vector store record mapper
83-
* @return the builder
60+
* @param recordClass The record class.
61+
* @return The updated builder instance.
8462
*/
85-
public Builder<Record> withVectorStoreRecordMapper(
86-
VectorStoreRecordMapper<Record, SearchDocument> vectorStoreRecordMapper) {
87-
this.vectorStoreRecordMapper = vectorStoreRecordMapper;
63+
public Builder<Record> withRecordClass(Class<Record> recordClass) {
64+
this.recordClass = recordClass;
8865
return this;
8966
}
9067

9168
/**
92-
* Sets the record definition.
69+
* Sets the vector store record collection factory.
9370
*
94-
* @param recordDefinition the record definition
95-
* @return the builder
71+
* @param vectorStoreRecordCollectionFactory The vector store record collection factory.
72+
* @return The updated builder instance.
9673
*/
97-
public Builder<Record> withRecordDefinition(VectorStoreRecordDefinition recordDefinition) {
98-
this.recordDefinition = recordDefinition;
74+
public Builder<Record> withVectorStoreRecordCollectionFactory(
75+
AzureAISearchVectorStoreRecordCollectionFactory<Record> vectorStoreRecordCollectionFactory) {
76+
this.vectorStoreRecordCollectionFactory = vectorStoreRecordCollectionFactory;
9977
return this;
10078
}
10179

10280
/**
103-
* Builds the options.
81+
* Builds the Azure AI Search vector store options.
10482
*
105-
* @return the options
83+
* @return The Azure AI Search vector store options.
10684
*/
10785
public AzureAISearchVectorStoreOptions<Record> build() {
108-
if (recordClass == null) {
109-
throw new IllegalArgumentException("recordClass must be provided");
110-
}
111-
112-
return new AzureAISearchVectorStoreOptions<>(
113-
recordClass,
114-
vectorStoreRecordMapper,
115-
recordDefinition);
86+
return new AzureAISearchVectorStoreOptions<>(recordClass,
87+
vectorStoreRecordCollectionFactory);
11688
}
11789
}
11890
}

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

Lines changed: 41 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,10 @@
44
import com.azure.search.documents.SearchAsyncClient;
55
import com.azure.search.documents.SearchDocument;
66
import com.azure.search.documents.indexes.SearchIndexAsyncClient;
7-
import com.azure.search.documents.indexes.models.HnswAlgorithmConfiguration;
8-
import com.azure.search.documents.indexes.models.HnswParameters;
97
import com.azure.search.documents.indexes.models.SearchField;
108
import com.azure.search.documents.indexes.models.SearchIndex;
119
import com.azure.search.documents.indexes.models.VectorSearch;
1210
import com.azure.search.documents.indexes.models.VectorSearchAlgorithmConfiguration;
13-
import com.azure.search.documents.indexes.models.VectorSearchAlgorithmMetric;
1411
import com.azure.search.documents.indexes.models.VectorSearchProfile;
1512
import com.azure.search.documents.models.IndexDocumentsResult;
1613
import com.azure.search.documents.models.IndexingResult;
@@ -45,56 +42,61 @@
4542
public class AzureAISearchVectorStoreRecordCollection<Record>
4643
implements VectorStoreRecordCollection<String, Record> {
4744

48-
// TODO: Check if the types are supported
49-
//
50-
// private HashSet<Class<?>> supportedDataTypes = new HashSet<>(
51-
// Arrays.asList(
52-
// String.class,
53-
// Integer.class,
54-
// int.class,
55-
// Long.class,
56-
// long.class,
57-
// Float.class,
58-
// float.class,
59-
// Double.class,
60-
// double.class,
61-
// Boolean.class,
62-
// boolean.class,
63-
// OffsetDateTime.class));
64-
//
65-
// private HashSet<Class<?>> supportedVectorTypes = new HashSet<>(
66-
// Arrays.asList(
67-
// List.class,
68-
// Collection.class));
45+
private static final HashSet<Class<?>> supportedKeyTypes = new HashSet<>(
46+
Collections.singletonList(
47+
String.class));
48+
49+
private static final HashSet<Class<?>> supportedDataTypes = new HashSet<>(
50+
Arrays.asList(
51+
String.class,
52+
Integer.class,
53+
int.class,
54+
Long.class,
55+
long.class,
56+
Float.class,
57+
float.class,
58+
Double.class,
59+
double.class,
60+
Boolean.class,
61+
boolean.class,
62+
OffsetDateTime.class));
63+
64+
private static final HashSet<Class<?>> supportedVectorTypes = new HashSet<>(
65+
Arrays.asList(
66+
List.class,
67+
Collection.class));
6968

7069
private final SearchIndexAsyncClient client;
7170
private final String collectionName;
7271
private final Map<String, SearchAsyncClient> clientsByIndex = new ConcurrentHashMap<>();
73-
private final AzureAISearchVectorStoreOptions<Record> options;
72+
private final AzureAISearchVectorStoreRecordCollectionOptions<Record> options;
73+
private final VectorStoreRecordDefinition recordDefinition;
74+
75+
// List of non-vector fields. Used to fetch only non-vector fields when vectors are not requested
7476
private final List<String> nonVectorFields = new ArrayList<>();
7577

7678
@SuppressFBWarnings("EI_EXPOSE_REP2")
7779
public AzureAISearchVectorStoreRecordCollection(
7880
@Nonnull SearchIndexAsyncClient client,
7981
@Nonnull String collectionName,
80-
@Nonnull AzureAISearchVectorStoreOptions<Record> options) {
82+
@Nonnull AzureAISearchVectorStoreRecordCollectionOptions<Record> options) {
8183
this.client = client;
8284
this.collectionName = collectionName;
85+
this.options = options;
8386

8487
// If record definition is not provided, create one from the record class
85-
if (options.getRecordDefinition() == null) {
86-
this.options = AzureAISearchVectorStoreOptions.<Record>builder()
87-
.withRecordClass(options.getRecordClass())
88-
.withVectorStoreRecordMapper(options.getVectorStoreRecordMapper())
89-
.withRecordDefinition(VectorStoreRecordDefinition.create(options.getRecordClass()))
90-
.build();
91-
} else {
92-
this.options = options;
93-
}
88+
this.recordDefinition = options.getRecordDefinition() == null
89+
? VectorStoreRecordDefinition.fromRecordClass(options.getRecordClass())
90+
: options.getRecordDefinition();
91+
92+
// Validate supported types
93+
VectorStoreRecordDefinition.validateSupportedTypes(
94+
this.options.getRecordClass(), this.recordDefinition, supportedKeyTypes,
95+
supportedDataTypes, supportedVectorTypes);
9496

9597
// Add non-vector fields to the list
96-
nonVectorFields.add(this.options.getRecordDefinition().getKeyField().getName());
97-
nonVectorFields.addAll(this.options.getRecordDefinition().getDataFields().stream()
98+
nonVectorFields.add(this.recordDefinition.getKeyField().getName());
99+
nonVectorFields.addAll(this.recordDefinition.getDataFields().stream()
98100
.map(VectorStoreRecordDataField::getName)
99101
.collect(Collectors.toList()));
100102
}
@@ -120,7 +122,7 @@ public Mono<Void> createCollectionAsync() {
120122
List<VectorSearchAlgorithmConfiguration> algorithms = new ArrayList<>();
121123
List<VectorSearchProfile> profiles = new ArrayList<>();
122124

123-
for (VectorStoreRecordField field : this.options.getRecordDefinition().getAllFields()) {
125+
for (VectorStoreRecordField field : this.recordDefinition.getAllFields()) {
124126
if (field instanceof VectorStoreRecordKeyField) {
125127
searchFields.add(AzureAISearchVectorStoreCollectionCreateMapping
126128
.mapKeyField((VectorStoreRecordKeyField) field));
@@ -247,7 +249,7 @@ public Mono<Void> deleteBatchAsync(List<String> keys, DeleteRecordOptions option
247249

248250
return client.deleteDocuments(keys.stream().map(key -> {
249251
SearchDocument document = new SearchDocument();
250-
document.put(this.options.getRecordDefinition().getKeyField().getName(), key);
252+
document.put(this.recordDefinition.getKeyField().getName(), key);
251253
return document;
252254
}).collect(Collectors.toList())).then();
253255
}

0 commit comments

Comments
 (0)