Skip to content

Commit 0fd0ce5

Browse files
johnoliverjohnoliver
authored
Merge pull request #202 from milderhc/vector-search
Add API classes for vector search and Azure AI implementation
2 parents 5955e31 + de36946 commit 0fd0ce5

75 files changed

Lines changed: 1016 additions & 321 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,9 @@
22

33
import com.fasterxml.jackson.annotation.JsonCreator;
44
import com.fasterxml.jackson.annotation.JsonProperty;
5-
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordDataAttribute;
6-
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordKeyAttribute;
7-
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordVectorAttribute;
5+
import com.microsoft.semantickernel.data.vectorstorage.attributes.VectorStoreRecordDataAttribute;
6+
import com.microsoft.semantickernel.data.vectorstorage.attributes.VectorStoreRecordKeyAttribute;
7+
import com.microsoft.semantickernel.data.vectorstorage.attributes.VectorStoreRecordVectorAttribute;
88

99
import java.util.List;
1010

@@ -16,7 +16,7 @@ public class Hotel {
1616
@VectorStoreRecordDataAttribute
1717
private final int code;
1818
@JsonProperty("summary")
19-
@VectorStoreRecordDataAttribute(hasEmbedding = true, embeddingFieldName = "descriptionEmbedding")
19+
@VectorStoreRecordDataAttribute()
2020
private final String description;
2121
@JsonProperty("summaryEmbedding")
2222
@VectorStoreRecordVectorAttribute(dimensions = 3)

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreRecordCollectionOptions;
66
import com.microsoft.semantickernel.connectors.data.mysql.MySQLVectorStoreQueryProvider;
77
import com.microsoft.semantickernel.connectors.data.postgres.PostgreSQLVectorStoreQueryProvider;
8-
import com.microsoft.semantickernel.data.recordoptions.GetRecordOptions;
8+
import com.microsoft.semantickernel.data.vectorstorage.options.GetRecordOptions;
99
import com.microsoft.semantickernel.tests.connectors.memory.Hotel;
1010
import com.mysql.cj.jdbc.MysqlDataSource;
1111
import org.junit.jupiter.params.ParameterizedTest;

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

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,12 @@
22

33
import com.microsoft.semantickernel.connectors.data.redis.RedisHashSetVectorStoreRecordCollection;
44
import com.microsoft.semantickernel.connectors.data.redis.RedisHashSetVectorStoreRecordCollectionOptions;
5-
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDataField;
6-
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition;
7-
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordField;
8-
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordKeyField;
9-
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordVectorField;
10-
import com.microsoft.semantickernel.data.recordoptions.GetRecordOptions;
5+
import com.microsoft.semantickernel.data.vectorstorage.definition.VectorStoreRecordDataField;
6+
import com.microsoft.semantickernel.data.vectorstorage.definition.VectorStoreRecordDefinition;
7+
import com.microsoft.semantickernel.data.vectorstorage.definition.VectorStoreRecordField;
8+
import com.microsoft.semantickernel.data.vectorstorage.definition.VectorStoreRecordKeyField;
9+
import com.microsoft.semantickernel.data.vectorstorage.definition.VectorStoreRecordVectorField;
10+
import com.microsoft.semantickernel.data.vectorstorage.options.GetRecordOptions;
1111
import com.microsoft.semantickernel.tests.connectors.memory.Hotel;
1212
import com.redis.testcontainers.RedisContainer;
1313
import org.junit.jupiter.api.BeforeAll;
@@ -67,8 +67,6 @@ static void setup() {
6767
.withName("description")
6868
.withStorageName("summary")
6969
.withFieldType(String.class)
70-
.withHasEmbedding(true)
71-
.withEmbeddingFieldName("descriptionEmbedding")
7270
.build());
7371
fields.add(VectorStoreRecordVectorField.builder()
7472
.withName("descriptionEmbedding")

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

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,12 @@
22

33
import com.microsoft.semantickernel.connectors.data.redis.RedisJsonVectorStoreRecordCollection;
44
import com.microsoft.semantickernel.connectors.data.redis.RedisJsonVectorStoreRecordCollectionOptions;
5-
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDataField;
6-
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition;
7-
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordField;
8-
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordKeyField;
9-
import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordVectorField;
10-
import com.microsoft.semantickernel.data.recordoptions.GetRecordOptions;
5+
import com.microsoft.semantickernel.data.vectorstorage.definition.VectorStoreRecordDataField;
6+
import com.microsoft.semantickernel.data.vectorstorage.definition.VectorStoreRecordDefinition;
7+
import com.microsoft.semantickernel.data.vectorstorage.definition.VectorStoreRecordField;
8+
import com.microsoft.semantickernel.data.vectorstorage.definition.VectorStoreRecordKeyField;
9+
import com.microsoft.semantickernel.data.vectorstorage.definition.VectorStoreRecordVectorField;
10+
import com.microsoft.semantickernel.data.vectorstorage.options.GetRecordOptions;
1111
import com.microsoft.semantickernel.tests.connectors.memory.Hotel;
1212
import com.redis.testcontainers.RedisContainer;
1313
import org.junit.jupiter.api.BeforeAll;
@@ -68,8 +68,6 @@ static void setup() {
6868
.withName("description")
6969
.withStorageName("summary")
7070
.withFieldType(String.class)
71-
.withHasEmbedding(true)
72-
.withEmbeddingFieldName("descriptionEmbedding")
7371
.build());
7472
fields.add(VectorStoreRecordVectorField.builder()
7573
.withName("descriptionEmbedding")

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

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,9 @@
55
import com.microsoft.semantickernel.connectors.data.redis.RedisStorageType;
66
import com.microsoft.semantickernel.connectors.data.redis.RedisVectorStore;
77
import com.microsoft.semantickernel.connectors.data.redis.RedisVectorStoreOptions;
8-
import com.microsoft.semantickernel.data.VectorStoreRecordCollectionOptions;
8+
import com.microsoft.semantickernel.data.vectorstorage.VectorStoreRecordCollectionOptions;
99
import com.microsoft.semantickernel.tests.connectors.memory.Hotel;
1010
import com.redis.testcontainers.RedisContainer;
11-
import org.junit.jupiter.api.BeforeAll;
12-
import org.junit.jupiter.api.BeforeEach;
13-
import org.junit.jupiter.api.Test;
1411
import org.junit.jupiter.params.ParameterizedTest;
1512
import org.junit.jupiter.params.provider.EnumSource;
1613
import org.testcontainers.junit.jupiter.Container;

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
import com.microsoft.semantickernel.samples.syntaxexamples.functions.Example59_OpenAIFunctionCalling;
1818
import com.microsoft.semantickernel.samples.syntaxexamples.functions.Example60_AdvancedMethodFunctions;
1919
import com.microsoft.semantickernel.samples.syntaxexamples.java.KernelFunctionYaml_Example;
20-
import com.microsoft.semantickernel.samples.syntaxexamples.memory.AzureAISearch_DataStorage;
20+
import com.microsoft.semantickernel.samples.syntaxexamples.memory.AzureAISearchVectorStore;
2121
import com.microsoft.semantickernel.samples.syntaxexamples.plugins.Example10_DescribeAllPluginsAndFunctions;
2222
import com.microsoft.semantickernel.samples.syntaxexamples.plugins.Example13_ConversationSummaryPlugin;
2323
import com.microsoft.semantickernel.samples.syntaxexamples.template.Example06_TemplateLanguage;
@@ -38,7 +38,7 @@ public class RunAll {
3838

3939
public static void main(String[] args) {
4040
List<MainMethod> mains = Arrays.asList(
41-
AzureAISearch_DataStorage::main,
41+
AzureAISearchVectorStore::main,
4242
Example01_NativeFunctions::main,
4343
Example03_Arguments::main,
4444
Example05_InlineFunctionDefinition::main,

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

Lines changed: 33 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -10,27 +10,27 @@
1010
import com.azure.core.util.TracingOptions;
1111
import com.azure.search.documents.indexes.SearchIndexAsyncClient;
1212
import com.azure.search.documents.indexes.SearchIndexClientBuilder;
13-
import com.fasterxml.jackson.annotation.JsonCreator;
1413
import com.fasterxml.jackson.annotation.JsonProperty;
1514
import com.microsoft.semantickernel.aiservices.openai.textembedding.OpenAITextEmbeddingGenerationService;
16-
import com.microsoft.semantickernel.connectors.data.azureaisearch.AzureAISearchVectorStore;
1715
import com.microsoft.semantickernel.connectors.data.azureaisearch.AzureAISearchVectorStoreOptions;
16+
import com.microsoft.semantickernel.connectors.data.azureaisearch.AzureAISearchVectorStoreRecordCollection;
1817
import com.microsoft.semantickernel.connectors.data.azureaisearch.AzureAISearchVectorStoreRecordCollectionOptions;
19-
import com.microsoft.semantickernel.data.VectorStoreRecordCollection;
20-
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordDataAttribute;
21-
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordKeyAttribute;
22-
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordVectorAttribute;
18+
import com.microsoft.semantickernel.data.vectorsearch.VectorSearchResult;
19+
import com.microsoft.semantickernel.data.vectorstorage.attributes.VectorStoreRecordDataAttribute;
20+
import com.microsoft.semantickernel.data.vectorstorage.attributes.VectorStoreRecordKeyAttribute;
21+
import com.microsoft.semantickernel.data.vectorstorage.attributes.VectorStoreRecordVectorAttribute;
2322
import java.nio.charset.StandardCharsets;
2423
import java.util.Arrays;
2524
import java.util.Base64;
2625
import java.util.Collections;
2726
import java.util.List;
2827
import java.util.Map;
2928
import java.util.stream.Collectors;
29+
3030
import reactor.core.publisher.Flux;
3131
import reactor.core.publisher.Mono;
3232

33-
public class AzureAISearch_DataStorage {
33+
public class AzureAISearchVectorStore {
3434

3535
private static final String CLIENT_KEY = System.getenv("CLIENT_KEY");
3636
private static final String AZURE_CLIENT_KEY = System.getenv("AZURE_CLIENT_KEY");
@@ -51,7 +51,7 @@ static class GitHubFile {
5151
@JsonProperty("fileId") // Set a different name for the storage field if needed
5252
@VectorStoreRecordKeyAttribute()
5353
private final String id;
54-
@VectorStoreRecordDataAttribute(hasEmbedding = true, embeddingFieldName = "embedding")
54+
@VectorStoreRecordDataAttribute()
5555
private final String description;
5656
@VectorStoreRecordDataAttribute
5757
private final String link;
@@ -118,13 +118,13 @@ public static void dataStorageWithAzureAISearch(
118118
OpenAITextEmbeddingGenerationService embeddingGeneration) {
119119

120120
// Create a new Azure AI Search vector store
121-
var azureAISearchVectorStore = AzureAISearchVectorStore.builder()
122-
.withClient(searchClient)
121+
var azureAISearchVectorStore = com.microsoft.semantickernel.connectors.data.azureaisearch.AzureAISearchVectorStore.builder()
122+
.withSearchIndexAsyncClient(searchClient)
123123
.withOptions(new AzureAISearchVectorStoreOptions())
124124
.build();
125125

126126
String collectionName = "skgithubfiles";
127-
var collection = azureAISearchVectorStore.getCollection(
127+
var collection = (AzureAISearchVectorStoreRecordCollection<GitHubFile>) azureAISearchVectorStore.getCollection(
128128
collectionName,
129129
AzureAISearchVectorStoreRecordCollectionOptions.<GitHubFile>builder()
130130
.withRecordClass(GitHubFile.class)
@@ -136,18 +136,31 @@ public static void dataStorageWithAzureAISearch(
136136
.then(storeData(collection, embeddingGeneration, sampleData()))
137137
.block();
138138

139-
// Query the Azure AI Search client for results
140-
// This might take a few seconds to return the best result
141-
var result = searchClient.getSearchAsyncClient(collectionName)
142-
.search("How to get started with the Semantic Kernel?")
143-
.blockFirst();
139+
// Search for results
140+
// Might need to wait for the data to be indexed
141+
var results = search("How to get started", collection, embeddingGeneration).block();
142+
143+
if (results == null || results.isEmpty()) {
144+
System.out.println("No search results found.");
145+
return;
146+
}
147+
var searchResult = results.get(0);
148+
System.out.printf("Search result with score: %f.%n Link: %s, Description: %s%n",
149+
searchResult.getScore(), searchResult.getRecord().link, searchResult.getRecord().description);
150+
}
151+
152+
153+
private static Mono<List<VectorSearchResult<GitHubFile>>> search(
154+
String searchText,
155+
AzureAISearchVectorStoreRecordCollection<GitHubFile> recordCollection,
156+
OpenAITextEmbeddingGenerationService embeddingGeneration) {
144157

145-
GitHubFile gitHubFile = result.getDocument(GitHubFile.class);
146-
System.out.println("Best result: " + gitHubFile.description + ". Link: " + gitHubFile.link);
158+
return embeddingGeneration.generateEmbeddingsAsync(Collections.singletonList(searchText))
159+
.flatMap(r -> recordCollection.searchAsync(r.get(0).getVector(), null));
147160
}
148161

149162
private static Mono<List<String>> storeData(
150-
VectorStoreRecordCollection<String, GitHubFile> recordStore,
163+
AzureAISearchVectorStoreRecordCollection<GitHubFile> recordCollection,
151164
OpenAITextEmbeddingGenerationService embeddingGeneration,
152165
Map<String, String> data) {
153166

@@ -163,7 +176,7 @@ private static Mono<List<String>> storeData(
163176
entry.getValue(),
164177
entry.getKey(),
165178
embeddings.get(0).getVector());
166-
return recordStore.upsertAsync(gitHubFile, null);
179+
return recordCollection.upsertAsync(gitHubFile, null);
167180
});
168181
})
169182
.collectList();

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,12 @@
66
import com.azure.core.credential.AzureKeyCredential;
77
import com.azure.core.credential.KeyCredential;
88
import com.microsoft.semantickernel.aiservices.openai.textembedding.OpenAITextEmbeddingGenerationService;
9-
import com.microsoft.semantickernel.data.VectorStoreRecordCollection;
9+
import com.microsoft.semantickernel.data.vectorstorage.VectorStoreRecordCollection;
1010
import com.microsoft.semantickernel.data.VolatileVectorStore;
1111
import com.microsoft.semantickernel.data.VolatileVectorStoreRecordCollectionOptions;
12-
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordDataAttribute;
13-
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordKeyAttribute;
14-
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordVectorAttribute;
12+
import com.microsoft.semantickernel.data.vectorstorage.attributes.VectorStoreRecordDataAttribute;
13+
import com.microsoft.semantickernel.data.vectorstorage.attributes.VectorStoreRecordKeyAttribute;
14+
import com.microsoft.semantickernel.data.vectorstorage.attributes.VectorStoreRecordVectorAttribute;
1515
import java.util.Arrays;
1616
import java.util.Collections;
1717
import java.util.List;
@@ -37,7 +37,7 @@ static class GitHubFile {
3737

3838
@VectorStoreRecordKeyAttribute()
3939
private final String id;
40-
@VectorStoreRecordDataAttribute(hasEmbedding = true, embeddingFieldName = "embedding")
40+
@VectorStoreRecordDataAttribute()
4141
private final String description;
4242
@VectorStoreRecordDataAttribute
4343
private final String link;
@@ -64,7 +64,7 @@ public String getDescription() {
6464
}
6565

6666
static String encodeId(String realId) {
67-
return AzureAISearch_DataStorage.GitHubFile.encodeId(realId);
67+
return AzureAISearchVectorStore.GitHubFile.encodeId(realId);
6868
}
6969
}
7070

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,10 @@
1010
import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreOptions;
1111
import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreRecordCollectionOptions;
1212
import com.microsoft.semantickernel.connectors.data.mysql.MySQLVectorStoreQueryProvider;
13-
import com.microsoft.semantickernel.data.VectorStoreRecordCollection;
14-
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordDataAttribute;
15-
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordKeyAttribute;
16-
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordVectorAttribute;
13+
import com.microsoft.semantickernel.data.vectorstorage.VectorStoreRecordCollection;
14+
import com.microsoft.semantickernel.data.vectorstorage.attributes.VectorStoreRecordDataAttribute;
15+
import com.microsoft.semantickernel.data.vectorstorage.attributes.VectorStoreRecordKeyAttribute;
16+
import com.microsoft.semantickernel.data.vectorstorage.attributes.VectorStoreRecordVectorAttribute;
1717
import com.mysql.cj.jdbc.MysqlDataSource;
1818
import java.nio.charset.StandardCharsets;
1919
import java.sql.SQLException;
@@ -45,7 +45,7 @@ static class GitHubFile {
4545

4646
@VectorStoreRecordKeyAttribute()
4747
private final String id;
48-
@VectorStoreRecordDataAttribute(hasEmbedding = true, embeddingFieldName = "embedding")
48+
@VectorStoreRecordDataAttribute()
4949
private final String description;
5050
@VectorStoreRecordDataAttribute
5151
private final String link;

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,10 @@
1212
import com.microsoft.semantickernel.connectors.data.redis.RedisJsonVectorStoreRecordCollectionOptions;
1313
import com.microsoft.semantickernel.connectors.data.redis.RedisVectorStore;
1414
import com.microsoft.semantickernel.connectors.data.redis.RedisVectorStoreOptions;
15-
import com.microsoft.semantickernel.data.VectorStoreRecordCollection;
16-
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordDataAttribute;
17-
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordKeyAttribute;
18-
import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordVectorAttribute;
15+
import com.microsoft.semantickernel.data.vectorstorage.VectorStoreRecordCollection;
16+
import com.microsoft.semantickernel.data.vectorstorage.attributes.VectorStoreRecordDataAttribute;
17+
import com.microsoft.semantickernel.data.vectorstorage.attributes.VectorStoreRecordKeyAttribute;
18+
import com.microsoft.semantickernel.data.vectorstorage.attributes.VectorStoreRecordVectorAttribute;
1919

2020
import java.util.Arrays;
2121
import java.util.Collections;
@@ -47,7 +47,7 @@ public static class GitHubFile {
4747

4848
@VectorStoreRecordKeyAttribute()
4949
private final String id;
50-
@VectorStoreRecordDataAttribute(hasEmbedding = true, embeddingFieldName = "embedding")
50+
@VectorStoreRecordDataAttribute()
5151
private final String description;
5252
@VectorStoreRecordDataAttribute
5353
private final String link;
@@ -78,7 +78,7 @@ public String getDescription() {
7878
}
7979

8080
static String encodeId(String realId) {
81-
return AzureAISearch_DataStorage.GitHubFile.encodeId(realId);
81+
return AzureAISearchVectorStore.GitHubFile.encodeId(realId);
8282
}
8383
}
8484

0 commit comments

Comments
 (0)