From 016f08fb268b0cba34fae12865899c4a0001416e Mon Sep 17 00:00:00 2001 From: John Oliver <1615532+johnoliver@users.noreply.github.com> Date: Thu, 8 Aug 2024 19:52:46 +0100 Subject: [PATCH] Memory updates for sample app --- .../JDBCVectorStoreRecordCollectionTest.java | 59 +++++----- .../memory/jdbc/JDBCVectorStoreTest.java | 18 ++- .../memory/AzureAISearch_DataStorage.java | 15 ++- .../memory/InMemory_DataStorage.java | 24 ++-- .../memory/JDBC_DataStorage.java | 104 +++++++++--------- .../memory/Redis_DataStorage.java | 4 +- .../AzureAISearchVectorStore.java | 39 +++++-- ...rchVectorStoreCollectionCreateMapping.java | 5 +- .../AzureAISearchVectorStoreOptions.java | 3 +- ...reAISearchVectorStoreRecordCollection.java | 34 +++--- ...rchVectorStoreRecordCollectionFactory.java | 5 +- ...rchVectorStoreRecordCollectionOptions.java | 3 +- .../connectors/data/jdbc/JDBCVectorStore.java | 29 ++++- .../jdbc/JDBCVectorStoreRecordCollection.java | 32 +++--- ...DBCVectorStoreRecordCollectionFactory.java | 1 - .../connectors/data/jdbc/SQLVectorStore.java | 5 +- .../data/redis/RedisVectorStore.java | 23 +++- .../RedisVectorStoreRecordCollection.java | 39 +++---- ...disVectorStoreRecordCollectionFactory.java | 7 +- .../redis/RedisVectorStoreRecordMapper.java | 4 +- .../semantickernel/data/VectorStore.java | 17 ++- .../data/VectorStoreRecordCollection.java | 20 ++-- .../data/VectorStoreRecordMapper.java | 5 +- .../data/VolatileVectorStore.java | 31 ++++-- .../VolatileVectorStoreRecordCollection.java | 25 +++-- ...latileVectorStoreRecordCollectionTest.java | 16 ++- .../data/VolatileVectorStoreTest.java | 12 +- 27 files changed, 319 insertions(+), 260 deletions(-) diff --git a/api-test/integration-tests/src/test/java/com/microsoft/semantickernel/tests/connectors/memory/jdbc/JDBCVectorStoreRecordCollectionTest.java b/api-test/integration-tests/src/test/java/com/microsoft/semantickernel/tests/connectors/memory/jdbc/JDBCVectorStoreRecordCollectionTest.java index 6e80e4ac..c873ea5b 100644 --- a/api-test/integration-tests/src/test/java/com/microsoft/semantickernel/tests/connectors/memory/jdbc/JDBCVectorStoreRecordCollectionTest.java +++ b/api-test/integration-tests/src/test/java/com/microsoft/semantickernel/tests/connectors/memory/jdbc/JDBCVectorStoreRecordCollectionTest.java @@ -1,37 +1,34 @@ package com.microsoft.semantickernel.tests.connectors.memory.jdbc; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; + import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreRecordCollection; import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreRecordCollectionOptions; import com.microsoft.semantickernel.connectors.data.jdbc.MySQLVectorStoreQueryProvider; import com.microsoft.semantickernel.data.recordoptions.GetRecordOptions; import com.microsoft.semantickernel.tests.connectors.memory.Hotel; import com.mysql.cj.jdbc.MysqlDataSource; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import javax.annotation.Nonnull; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import org.testcontainers.containers.MySQLContainer; import org.testcontainers.junit.jupiter.Container; import org.testcontainers.junit.jupiter.Testcontainers; -import javax.annotation.Nonnull; -import javax.sql.DataSource; -import java.sql.Connection; -import java.sql.DriverManager; -import java.sql.SQLException; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.List; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; - @Testcontainers public class JDBCVectorStoreRecordCollectionTest { + @Container private static final MySQLContainer CONTAINER = new MySQLContainer<>("mysql:5.7.34"); private static final String MYSQL_USER = "test"; private static final String MYSQL_PASSWORD = "test"; private static MysqlDataSource dataSource; + @BeforeAll static void setup() { dataSource = new MysqlDataSource(); @@ -40,16 +37,17 @@ static void setup() { dataSource.setPassword(MYSQL_PASSWORD); } - private JDBCVectorStoreRecordCollection buildRecordCollection(@Nonnull String collectionName) { - JDBCVectorStoreRecordCollection recordCollection = new JDBCVectorStoreRecordCollection<>( - dataSource, - collectionName, - JDBCVectorStoreRecordCollectionOptions.builder() - .withRecordClass(Hotel.class) - .withQueryProvider(MySQLVectorStoreQueryProvider.builder() - .withDataSource(dataSource) - .build()) - .build()); + private JDBCVectorStoreRecordCollection buildRecordCollection( + @Nonnull String collectionName) { + JDBCVectorStoreRecordCollection recordCollection = new JDBCVectorStoreRecordCollection<>( + dataSource, + collectionName, + JDBCVectorStoreRecordCollectionOptions.builder() + .withRecordClass(Hotel.class) + .withQueryProvider(MySQLVectorStoreQueryProvider.builder() + .withDataSource(dataSource) + .build()) + .build()); recordCollection.prepareAsync().block(); recordCollection.createCollectionIfNotExistsAsync().block(); @@ -63,11 +61,16 @@ public void buildRecordCollection() { private List getHotels() { return List.of( - new Hotel("id_1", "Hotel 1", 1, "Hotel 1 description", Arrays.asList(1.0f, 2.0f, 3.0f), 4.0), - new Hotel("id_2", "Hotel 2", 2, "Hotel 2 description", Arrays.asList(1.0f, 2.0f, 3.0f), 3.0), - new Hotel("id_3", "Hotel 3", 3, "Hotel 3 description", Arrays.asList(1.0f, 2.0f, 3.0f), 5.0), - new Hotel("id_4", "Hotel 4", 4, "Hotel 4 description", Arrays.asList(1.0f, 2.0f, 3.0f), 4.0), - new Hotel("id_5", "Hotel 5", 5, "Hotel 5 description", Arrays.asList(1.0f, 2.0f, 3.0f), 5.0) + new Hotel("id_1", "Hotel 1", 1, "Hotel 1 description", Arrays.asList(1.0f, 2.0f, 3.0f), + 4.0), + new Hotel("id_2", "Hotel 2", 2, "Hotel 2 description", Arrays.asList(1.0f, 2.0f, 3.0f), + 3.0), + new Hotel("id_3", "Hotel 3", 3, "Hotel 3 description", Arrays.asList(1.0f, 2.0f, 3.0f), + 5.0), + new Hotel("id_4", "Hotel 4", 4, "Hotel 4 description", Arrays.asList(1.0f, 2.0f, 3.0f), + 4.0), + new Hotel("id_5", "Hotel 5", 5, "Hotel 5 description", Arrays.asList(1.0f, 2.0f, 3.0f), + 5.0) ); } diff --git a/api-test/integration-tests/src/test/java/com/microsoft/semantickernel/tests/connectors/memory/jdbc/JDBCVectorStoreTest.java b/api-test/integration-tests/src/test/java/com/microsoft/semantickernel/tests/connectors/memory/jdbc/JDBCVectorStoreTest.java index eb134dd0..e0906a1c 100644 --- a/api-test/integration-tests/src/test/java/com/microsoft/semantickernel/tests/connectors/memory/jdbc/JDBCVectorStoreTest.java +++ b/api-test/integration-tests/src/test/java/com/microsoft/semantickernel/tests/connectors/memory/jdbc/JDBCVectorStoreTest.java @@ -1,26 +1,22 @@ package com.microsoft.semantickernel.tests.connectors.memory.jdbc; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStore; import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreOptions; import com.microsoft.semantickernel.connectors.data.jdbc.MySQLVectorStoreQueryProvider; import com.microsoft.semantickernel.tests.connectors.memory.Hotel; import com.mysql.cj.jdbc.MysqlDataSource; +import java.util.Arrays; +import java.util.List; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import org.testcontainers.containers.MySQLContainer; import org.testcontainers.junit.jupiter.Container; import org.testcontainers.junit.jupiter.Testcontainers; -import java.sql.Connection; -import java.sql.DriverManager; -import java.sql.SQLException; -import java.util.Arrays; -import java.util.List; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertTrue; - @Testcontainers public class JDBCVectorStoreTest { @Container @@ -43,7 +39,7 @@ public void getCollectionNamesAsync() { .withDataSource(dataSource) .build(); - JDBCVectorStore vectorStore = JDBCVectorStore.builder() + JDBCVectorStore vectorStore = JDBCVectorStore.builder() .withDataSource(dataSource) .withOptions( JDBCVectorStoreOptions.builder() diff --git a/samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/AzureAISearch_DataStorage.java b/samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/AzureAISearch_DataStorage.java index 92316385..03cc93e3 100644 --- a/samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/AzureAISearch_DataStorage.java +++ b/samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/AzureAISearch_DataStorage.java @@ -13,13 +13,10 @@ import com.microsoft.semantickernel.aiservices.openai.textembedding.OpenAITextEmbeddingGenerationService; import com.microsoft.semantickernel.connectors.data.azureaisearch.AzureAISearchVectorStore; import com.microsoft.semantickernel.connectors.data.azureaisearch.AzureAISearchVectorStoreOptions; -import com.microsoft.semantickernel.connectors.data.azureaisearch.AzureAISearchVectorStoreRecordCollection; +import com.microsoft.semantickernel.data.VectorStoreRecordCollection; import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordDataAttribute; import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordKeyAttribute; import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordVectorAttribute; -import reactor.core.publisher.Flux; -import reactor.core.publisher.Mono; - import java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.Base64; @@ -27,8 +24,11 @@ import java.util.List; import java.util.Map; import java.util.stream.Collectors; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; public class AzureAISearch_DataStorage { + private static final String CLIENT_KEY = System.getenv("CLIENT_KEY"); private static final String AZURE_CLIENT_KEY = System.getenv("AZURE_CLIENT_KEY"); @@ -45,6 +45,7 @@ public class AzureAISearch_DataStorage { private static final int EMBEDDING_DIMENSIONS = 1536; static class GitHubFile { + @VectorStoreRecordKeyAttribute() private final String id; @VectorStoreRecordDataAttribute(hasEmbedding = true, embeddingFieldName = "embedding") @@ -120,7 +121,9 @@ public static void dataStorageWithAzureAISearch( .build(); String collectionName = "skgithubfiles"; - var collection = azureAISearchVectorStore.getCollection(collectionName, GitHubFile.class, + var collection = azureAISearchVectorStore.getCollection( + collectionName, + GitHubFile.class, null); // Create collection if it does not exist and store data @@ -140,7 +143,7 @@ public static void dataStorageWithAzureAISearch( } private static Mono> storeData( - AzureAISearchVectorStoreRecordCollection recordStore, + VectorStoreRecordCollection recordStore, OpenAITextEmbeddingGenerationService embeddingGeneration, Map data) { diff --git a/samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/InMemory_DataStorage.java b/samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/InMemory_DataStorage.java index d92c1e06..c74d16f8 100644 --- a/samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/InMemory_DataStorage.java +++ b/samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/InMemory_DataStorage.java @@ -5,32 +5,22 @@ import com.azure.ai.openai.OpenAIClientBuilder; import com.azure.core.credential.AzureKeyCredential; import com.azure.core.credential.KeyCredential; -import com.azure.core.util.ClientOptions; -import com.azure.core.util.MetricsOptions; -import com.azure.core.util.TracingOptions; -import com.azure.search.documents.indexes.SearchIndexAsyncClient; -import com.azure.search.documents.indexes.SearchIndexClientBuilder; import com.microsoft.semantickernel.aiservices.openai.textembedding.OpenAITextEmbeddingGenerationService; -import com.microsoft.semantickernel.connectors.data.azureaisearch.AzureAISearchVectorStore; -import com.microsoft.semantickernel.connectors.data.azureaisearch.AzureAISearchVectorStoreOptions; -import com.microsoft.semantickernel.connectors.data.azureaisearch.AzureAISearchVectorStoreRecordCollection; +import com.microsoft.semantickernel.data.VectorStoreRecordCollection; import com.microsoft.semantickernel.data.VolatileVectorStore; -import com.microsoft.semantickernel.data.VolatileVectorStoreRecordCollection; import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordDataAttribute; import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordKeyAttribute; import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordVectorAttribute; -import reactor.core.publisher.Flux; -import reactor.core.publisher.Mono; - -import java.nio.charset.StandardCharsets; import java.util.Arrays; -import java.util.Base64; import java.util.Collections; import java.util.List; import java.util.Map; import java.util.stream.Collectors; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; public class InMemory_DataStorage { + private static final String CLIENT_KEY = System.getenv("CLIENT_KEY"); private static final String AZURE_CLIENT_KEY = System.getenv("AZURE_CLIENT_KEY"); @@ -43,6 +33,7 @@ public class InMemory_DataStorage { private static final int EMBEDDING_DIMENSIONS = 1536; static class GitHubFile { + @VectorStoreRecordKeyAttribute() private final String id; @VectorStoreRecordDataAttribute(hasEmbedding = true, embeddingFieldName = "embedding") @@ -72,8 +63,7 @@ public String getDescription() { } static String encodeId(String realId) { - byte[] bytes = Base64.getUrlEncoder().encode(realId.getBytes(StandardCharsets.UTF_8)); - return new String(bytes, StandardCharsets.UTF_8); + return AzureAISearch_DataStorage.GitHubFile.encodeId(realId); } } @@ -126,7 +116,7 @@ public static void inMemoryDataStorage( } private static Mono> storeData( - VolatileVectorStoreRecordCollection recordCollection, + VectorStoreRecordCollection recordCollection, OpenAITextEmbeddingGenerationService embeddingGeneration, Map data) { diff --git a/samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/JDBC_DataStorage.java b/samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/JDBC_DataStorage.java index 2379e572..8aa4bddb 100644 --- a/samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/JDBC_DataStorage.java +++ b/samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/JDBC_DataStorage.java @@ -1,3 +1,4 @@ +// Copyright (c) Microsoft. All rights reserved. package com.microsoft.semantickernel.samples.syntaxexamples.memory; import com.azure.ai.openai.OpenAIAsyncClient; @@ -13,13 +14,7 @@ import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordKeyAttribute; import com.microsoft.semantickernel.data.recordattributes.VectorStoreRecordVectorAttribute; import com.mysql.cj.jdbc.MysqlDataSource; -import reactor.core.publisher.Flux; -import reactor.core.publisher.Mono; - -import javax.sql.DataSource; import java.nio.charset.StandardCharsets; -import java.sql.Connection; -import java.sql.DriverManager; import java.sql.SQLException; import java.util.Arrays; import java.util.Base64; @@ -27,6 +22,9 @@ import java.util.List; import java.util.Map; import java.util.stream.Collectors; +import javax.sql.DataSource; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; public class JDBC_DataStorage { @@ -36,13 +34,14 @@ public class JDBC_DataStorage { // Only required if AZURE_CLIENT_KEY is set private static final String CLIENT_ENDPOINT = System.getenv("CLIENT_ENDPOINT"); private static final String MODEL_ID = System.getenv() - .getOrDefault("EMBEDDING_MODEL_ID", "text-embedding-3-large"); + .getOrDefault("EMBEDDING_MODEL_ID", "text-embedding-3-large"); private static final int EMBEDDING_DIMENSIONS = 1536; // Run a MySQL server with: // docker run -d --name mysql-container -e MYSQL_ROOT_PASSWORD=root -e MYSQL_DATABASE=sk -p 3306:3306 mysql:latest static class GitHubFile { + @VectorStoreRecordKeyAttribute() private final String id; @VectorStoreRecordDataAttribute(hasEmbedding = true, embeddingFieldName = "embedding") @@ -57,10 +56,10 @@ public GitHubFile() { } public GitHubFile( - String id, - String description, - String link, - List embedding) { + String id, + String description, + String link, + List embedding) { this.id = id; this.description = description; this.link = link; @@ -90,21 +89,21 @@ public static void main(String[] args) throws SQLException { if (AZURE_CLIENT_KEY != null) { client = new OpenAIClientBuilder() - .credential(new AzureKeyCredential(AZURE_CLIENT_KEY)) - .endpoint(CLIENT_ENDPOINT) - .buildAsyncClient(); + .credential(new AzureKeyCredential(AZURE_CLIENT_KEY)) + .endpoint(CLIENT_ENDPOINT) + .buildAsyncClient(); } else { client = new OpenAIClientBuilder() - .credential(new KeyCredential(CLIENT_KEY)) - .buildAsyncClient(); + .credential(new KeyCredential(CLIENT_KEY)) + .buildAsyncClient(); } var embeddingGeneration = OpenAITextEmbeddingGenerationService.builder() - .withOpenAIAsyncClient(client) - .withModelId(MODEL_ID) - .withDimensions(EMBEDDING_DIMENSIONS) - .build(); + .withOpenAIAsyncClient(client) + .withModelId(MODEL_ID) + .withDimensions(EMBEDDING_DIMENSIONS) + .build(); var dataSource = new MysqlDataSource(); dataSource.setUrl("jdbc:mysql://localhost:3306/sk"); @@ -115,31 +114,32 @@ public static void main(String[] args) throws SQLException { } public static void dataStorageWithMySQL( - DataSource dataSource, - OpenAITextEmbeddingGenerationService embeddingGeneration) { + DataSource dataSource, + OpenAITextEmbeddingGenerationService embeddingGeneration) { // Build a query provider var queryProvider = MySQLVectorStoreQueryProvider.builder() - .withDataSource(dataSource) - .build(); + .withDataSource(dataSource) + .build(); // Create a new vector store var jdbcVectorStore = JDBCVectorStore.builder() - .withDataSource(dataSource) - .withOptions(JDBCVectorStoreOptions.builder() - .withQueryProvider(queryProvider) - .build()) - .build(); + .withDataSource(dataSource) + .withOptions(JDBCVectorStoreOptions.builder() + .withQueryProvider(queryProvider) + .build()) + .build(); String collectionName = "skgithubfiles"; - var collection = jdbcVectorStore.getCollection(collectionName, GitHubFile.class, - null); + var collection = jdbcVectorStore.getCollection(collectionName, + GitHubFile.class, + null); // Create collection if it does not exist and store data List ids = collection - .createCollectionIfNotExistsAsync() - .then(storeData(collection, embeddingGeneration, sampleData())) - .block(); + .createCollectionIfNotExistsAsync() + .then(storeData(collection, embeddingGeneration, sampleData())) + .block(); List data = collection.getBatchAsync(ids, null).block(); @@ -147,26 +147,26 @@ public static void dataStorageWithMySQL( } private static Mono> storeData( - VectorStoreRecordCollection recordStore, - OpenAITextEmbeddingGenerationService embeddingGeneration, - Map data) { + VectorStoreRecordCollection recordStore, + OpenAITextEmbeddingGenerationService embeddingGeneration, + Map data) { return Flux.fromIterable(data.entrySet()) - .flatMap(entry -> { - System.out.println("Save '" + entry.getKey() + "' to memory."); - - return embeddingGeneration - .generateEmbeddingsAsync(Collections.singletonList(entry.getValue())) - .flatMap(embeddings -> { - GitHubFile gitHubFile = new GitHubFile( - GitHubFile.encodeId(entry.getKey()), - entry.getValue(), - entry.getKey(), - embeddings.get(0).getVector()); - return recordStore.upsertAsync(gitHubFile, null); - }); - }) - .collectList(); + .flatMap(entry -> { + System.out.println("Save '" + entry.getKey() + "' to memory."); + + return embeddingGeneration + .generateEmbeddingsAsync(Collections.singletonList(entry.getValue())) + .flatMap(embeddings -> { + GitHubFile gitHubFile = new GitHubFile( + GitHubFile.encodeId(entry.getKey()), + entry.getValue(), + entry.getKey(), + embeddings.get(0).getVector()); + return recordStore.upsertAsync(gitHubFile, null); + }); + }) + .collectList(); } private static Map sampleData() { diff --git a/samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/Redis_DataStorage.java b/samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/Redis_DataStorage.java index cde31a3a..6cd1db3d 100644 --- a/samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/Redis_DataStorage.java +++ b/samples/semantickernel-concepts/semantickernel-syntax-examples/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/memory/Redis_DataStorage.java @@ -22,6 +22,7 @@ import java.util.List; import java.util.Map; import java.util.stream.Collectors; +import com.microsoft.semantickernel.samples.syntaxexamples.memory.AzureAISearch_DataStorage.GitHubFile; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import redis.clients.jedis.JedisPooled; @@ -77,8 +78,7 @@ public String getDescription() { } static String encodeId(String realId) { - byte[] bytes = Base64.getUrlEncoder().encode(realId.getBytes(StandardCharsets.UTF_8)); - return new String(bytes, StandardCharsets.UTF_8); + return AzureAISearch_DataStorage.GitHubFile.encodeId(realId); } } diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStore.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStore.java index 39d13f75..a7b7a7c8 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStore.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStore.java @@ -4,16 +4,15 @@ import com.azure.search.documents.indexes.SearchIndexAsyncClient; import com.azure.search.documents.indexes.models.SearchIndex; import com.microsoft.semantickernel.data.VectorStore; +import com.microsoft.semantickernel.data.VectorStoreRecordCollection; import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; -import reactor.core.publisher.Mono; - +import java.util.List; import javax.annotation.Nonnull; import javax.annotation.Nullable; -import java.util.List; +import reactor.core.publisher.Mono; -public class AzureAISearchVectorStore - implements VectorStore> { +public class AzureAISearchVectorStore implements VectorStore { private final SearchIndexAsyncClient client; private final AzureAISearchVectorStoreOptions options; @@ -21,7 +20,7 @@ public class AzureAISearchVectorStore /** * Creates a new instance of {@link AzureAISearchVectorStore}. * - * @param client The Azure AI Search client. + * @param client The Azure AI Search client. * @param options The options for the vector store. */ @SuppressFBWarnings("EI_EXPOSE_REP2") @@ -34,17 +33,29 @@ public AzureAISearchVectorStore(@Nonnull SearchIndexAsyncClient client, /** * Gets a new instance of {@link AzureAISearchVectorStoreRecordCollection} * - * @param collectionName The name of the collection. - * @param recordClass The class type of the record. + * @param collectionName The name of the collection. + * @param recordClass The class type of the record. * @param recordDefinition The record definition. * @return The collection. */ @Override - public AzureAISearchVectorStoreRecordCollection getCollection( + public final VectorStoreRecordCollection getCollection( @Nonnull String collectionName, + @Nonnull Class keyClass, @Nonnull Class recordClass, - VectorStoreRecordDefinition recordDefinition) { + @Nullable VectorStoreRecordDefinition recordDefinition) { + if (!keyClass.equals(String.class)) { + throw new IllegalArgumentException("Azure AI Search only supports string keys"); + } + return (VectorStoreRecordCollection) getCollection( + collectionName, recordClass, recordDefinition); + } + + public AzureAISearchVectorStoreRecordCollection getCollection( + @Nonnull String collectionName, + @Nonnull Class recordClass, + @Nullable VectorStoreRecordDefinition recordDefinition) { if (options.getVectorStoreRecordCollectionFactory() != null) { return options.getVectorStoreRecordCollectionFactory() .createVectorStoreRecordCollection( @@ -56,7 +67,9 @@ public AzureAISearchVectorStoreRecordCollection getCollect .build()); } - return new AzureAISearchVectorStoreRecordCollection<>(client, collectionName, + return new AzureAISearchVectorStoreRecordCollection<>( + client, + collectionName, AzureAISearchVectorStoreRecordCollectionOptions.builder() .withRecordClass(recordClass) .withRecordDefinition(recordDefinition) @@ -86,6 +99,7 @@ public static Builder builder() { * Builder for {@link AzureAISearchVectorStore}. */ public static class Builder { + @Nullable private SearchIndexAsyncClient client; @Nullable @@ -109,7 +123,8 @@ public Builder withClient(@Nonnull SearchIndexAsyncClient client) { * @param options The options for the Azure AI Search vector store. * @return The updated builder instance. */ - public Builder withOptions(@Nonnull AzureAISearchVectorStoreOptions options) { + public Builder withOptions( + @Nonnull AzureAISearchVectorStoreOptions options) { this.options = options; return this; } diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreCollectionCreateMapping.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreCollectionCreateMapping.java index e077a510..c57e6f9d 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreCollectionCreateMapping.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreCollectionCreateMapping.java @@ -13,12 +13,9 @@ import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDataField; import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordKeyField; import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordVectorField; - -import javax.annotation.Nonnull; -import javax.annotation.Nullable; import java.time.OffsetDateTime; import java.util.List; -import java.util.Objects; +import javax.annotation.Nonnull; public class AzureAISearchVectorStoreCollectionCreateMapping { diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreOptions.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreOptions.java index d7bb0314..db3b7ab6 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreOptions.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreOptions.java @@ -1,10 +1,10 @@ // Copyright (c) Microsoft. All rights reserved. package com.microsoft.semantickernel.connectors.data.azureaisearch; -import javax.annotation.Nonnull; import javax.annotation.Nullable; public class AzureAISearchVectorStoreOptions { + @Nullable private final AzureAISearchVectorStoreRecordCollectionFactory vectorStoreRecordCollectionFactory; @@ -49,6 +49,7 @@ public AzureAISearchVectorStoreRecordCollectionFactory getVectorStoreRecordColle * */ public static class Builder { + @Nullable private AzureAISearchVectorStoreRecordCollectionFactory vectorStoreRecordCollectionFactory; diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreRecordCollection.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreRecordCollection.java index 5155299d..5937d928 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreRecordCollection.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreRecordCollection.java @@ -11,22 +11,18 @@ import com.azure.search.documents.indexes.models.VectorSearchProfile; import com.azure.search.documents.models.IndexDocumentsResult; import com.azure.search.documents.models.IndexingResult; -import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordField; -import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordKeyField; -import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordVectorField; -import com.microsoft.semantickernel.exceptions.SKException; import com.microsoft.semantickernel.data.VectorStoreRecordCollection; import com.microsoft.semantickernel.data.VectorStoreRecordMapper; import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDataField; import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition; +import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordField; +import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordKeyField; +import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordVectorField; import com.microsoft.semantickernel.data.recordoptions.DeleteRecordOptions; import com.microsoft.semantickernel.data.recordoptions.GetRecordOptions; import com.microsoft.semantickernel.data.recordoptions.UpsertRecordOptions; +import com.microsoft.semantickernel.exceptions.SKException; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; -import reactor.core.publisher.Flux; -import reactor.core.publisher.Mono; - -import javax.annotation.Nonnull; import java.time.OffsetDateTime; import java.util.ArrayList; import java.util.Arrays; @@ -38,9 +34,12 @@ import java.util.Map; import java.util.concurrent.ConcurrentHashMap; import java.util.stream.Collectors; +import javax.annotation.Nonnull; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; -public class AzureAISearchVectorStoreRecordCollection - implements VectorStoreRecordCollection { +public class AzureAISearchVectorStoreRecordCollection implements + VectorStoreRecordCollection { private static final HashSet> supportedKeyTypes = new HashSet<>( Collections.singletonList( @@ -124,7 +123,7 @@ public Mono collectionExistsAsync() { } @Override - public Mono createCollectionAsync() { + public Mono> createCollectionAsync() { List searchFields = new ArrayList<>(); List algorithms = new ArrayList<>(); List profiles = new ArrayList<>(); @@ -151,18 +150,19 @@ public Mono createCollectionAsync() { .setAlgorithms(algorithms) .setProfiles(profiles)); - return client.createIndex(newIndex).then(); + return client.createIndex(newIndex).then(Mono.just(this)); } @Override - public Mono createCollectionIfNotExistsAsync() { + public Mono> createCollectionIfNotExistsAsync() { return collectionExistsAsync().flatMap( exists -> { if (!exists) { return createCollectionAsync(); } return Mono.empty(); - }); + }) + .then(Mono.just(this)); } @Override @@ -191,11 +191,11 @@ public Mono getAsync( } return client.getDocumentWithResponse(key, this.options.getRecordClass(), selectedFields) - .map(response -> { + .flatMap(response -> { if (response.getStatusCode() == 404) { - throw new SKException("Record not found: " + key); + return Mono.error(new SKException("Record not found: " + key)); } - return response.getValue(); + return Mono.just(response.getValue()); }); } diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreRecordCollectionFactory.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreRecordCollectionFactory.java index c5041284..ec08ba03 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreRecordCollectionFactory.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreRecordCollectionFactory.java @@ -5,16 +5,15 @@ /** * Factory for creating Azure AI Search vector store record collections. - * */ public interface AzureAISearchVectorStoreRecordCollectionFactory { /** * Creates a new Azure AI Search vector store record collection. * - * @param client The Azure AI Search client. + * @param client The Azure AI Search client. * @param collectionName The name of the collection. - * @param options The options for the collection. + * @param options The options for the collection. * @return The new Azure AI Search vector store record collection. */ AzureAISearchVectorStoreRecordCollection createVectorStoreRecordCollection( diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreRecordCollectionOptions.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreRecordCollectionOptions.java index 45fb410c..7275dd43 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreRecordCollectionOptions.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreRecordCollectionOptions.java @@ -4,7 +4,6 @@ import com.azure.search.documents.SearchDocument; import com.microsoft.semantickernel.data.VectorStoreRecordMapper; import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition; - import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -14,6 +13,7 @@ * @param the record type */ public class AzureAISearchVectorStoreRecordCollectionOptions { + private final Class recordClass; @Nullable private final VectorStoreRecordMapper vectorStoreRecordMapper; @@ -75,6 +75,7 @@ private AzureAISearchVectorStoreRecordCollectionOptions( * @param the record type */ public static class Builder { + @Nullable private VectorStoreRecordMapper vectorStoreRecordMapper; @Nullable diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/JDBCVectorStore.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/JDBCVectorStore.java index 5e497176..66ad995b 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/JDBCVectorStore.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/JDBCVectorStore.java @@ -1,6 +1,8 @@ // Copyright (c) Microsoft. All rights reserved. package com.microsoft.semantickernel.connectors.data.jdbc; +import com.microsoft.semantickernel.connectors.data.redis.RedisVectorStoreRecordCollection; +import com.microsoft.semantickernel.data.VectorStoreRecordCollection; import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; import reactor.core.publisher.Mono; @@ -14,7 +16,7 @@ /** * A JDBC vector store. */ -public class JDBCVectorStore implements SQLVectorStore> { +public class JDBCVectorStore implements SQLVectorStore { private final DataSource dataSource; private final JDBCVectorStoreOptions options; private final JDBCVectorStoreQueryProvider queryProvider; @@ -59,11 +61,32 @@ public static Builder builder() { * @return The collection. */ @Override - public JDBCVectorStoreRecordCollection getCollection( - @Nonnull String collectionName, + public VectorStoreRecordCollection getCollection( + @Nonnull String collectionName, @Nonnull Class keyClass, @Nonnull Class recordClass, @Nullable VectorStoreRecordDefinition recordDefinition) { + if (keyClass != String.class) { + throw new IllegalArgumentException("Redis only supports string keys"); + } + + return (VectorStoreRecordCollection) getCollection( + collectionName, + recordClass, + recordDefinition); + } + /** + * Gets a collection from the vector store. + * + * @param collectionName The name of the collection. + * @param recordClass The class type of the record. + * @param recordDefinition The record definition. + * @return The collection. + */ + public JDBCVectorStoreRecordCollection getCollection( + @Nonnull String collectionName, + @Nonnull Class recordClass, + @Nullable VectorStoreRecordDefinition recordDefinition) { if (this.options != null && this.options.getVectorStoreRecordCollectionFactory() != null) { return this.options.getVectorStoreRecordCollectionFactory() .createVectorStoreRecordCollection( diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/JDBCVectorStoreRecordCollection.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/JDBCVectorStoreRecordCollection.java index b9c0bd3c..6135b284 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/JDBCVectorStoreRecordCollection.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/JDBCVectorStoreRecordCollection.java @@ -2,24 +2,25 @@ package com.microsoft.semantickernel.connectors.data.jdbc; import com.microsoft.semantickernel.builders.SemanticKernelBuilder; +import com.microsoft.semantickernel.data.VectorStoreRecordCollection; import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition; import com.microsoft.semantickernel.data.recordoptions.DeleteRecordOptions; import com.microsoft.semantickernel.data.recordoptions.GetRecordOptions; import com.microsoft.semantickernel.data.recordoptions.UpsertRecordOptions; import com.microsoft.semantickernel.exceptions.SKException; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; -import reactor.core.publisher.Mono; -import reactor.core.scheduler.Schedulers; - -import javax.annotation.Nonnull; -import javax.sql.DataSource; import java.lang.reflect.Field; import java.util.Collections; import java.util.List; import java.util.stream.Collectors; +import javax.annotation.Nonnull; +import javax.sql.DataSource; +import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; public class JDBCVectorStoreRecordCollection implements SQLVectorStoreRecordCollection { + private final String collectionName; private final VectorStoreRecordDefinition recordDefinition; private final JDBCVectorStoreRecordCollectionOptions options; @@ -29,9 +30,9 @@ public class JDBCVectorStoreRecordCollection /** * Creates a new instance of the {@link JDBCVectorStoreRecordCollection}. * - * @param dataSource the data source + * @param dataSource the data source * @param collectionName the name of the collection - * @param options the options + * @param options the options */ @SuppressFBWarnings("EI_EXPOSE_REP2") // DataSource is not exposed public JDBCVectorStoreRecordCollection( @@ -99,12 +100,12 @@ public Mono collectionExistsAsync() { * @throws SKException if the operation fails */ @Override - public Mono createCollectionAsync() { + public Mono> createCollectionAsync() { return Mono.fromRunnable( () -> queryProvider.createCollection(this.collectionName, options.getRecordClass(), recordDefinition)) .subscribeOn(Schedulers.boundedElastic()) - .then(); + .then(Mono.just(this)); } /** @@ -114,7 +115,7 @@ public Mono createCollectionAsync() { * @throws SKException if the operation fails */ @Override - public Mono createCollectionIfNotExistsAsync() { + public Mono> createCollectionIfNotExistsAsync() { return collectionExistsAsync().map( exists -> { if (!exists) { @@ -123,7 +124,7 @@ public Mono createCollectionIfNotExistsAsync() { return Mono.empty(); }) .flatMap(mono -> mono) - .then(); + .then(Mono.just(this)); } /** @@ -143,7 +144,7 @@ public Mono deleteCollectionAsync() { /** * Gets a record from the store. * - * @param key The key of the record to get. + * @param key The key of the record to get. * @param options The options for getting the record. * @return A Mono emitting the record. * @throws SKException if the operation fails @@ -162,7 +163,7 @@ public Mono getAsync(String key, GetRecordOptions options) { /** * Gets a batch of records from the store. * - * @param keys The keys of the records to get. + * @param keys The keys of the records to get. * @param options The options for getting the records. * @return A Mono emitting a collection of records. * @throws SKException if the operation fails @@ -227,7 +228,7 @@ public Mono> upsertBatchAsync(List data, UpsertRecordOption /** * Deletes a record from the store. * - * @param key The key of the record to delete. + * @param key The key of the record to delete. * @param options The options for deleting the record. * @return A Mono representing the completion of the deletion operation. * @throws SKException if the operation fails @@ -240,7 +241,7 @@ public Mono deleteAsync(String key, DeleteRecordOptions options) { /** * Deletes a batch of records from the store. * - * @param keys The keys of the records to delete. + * @param keys The keys of the records to delete. * @param options The options for deleting the records. * @return A Mono representing the completion of the deletion operation. * @throws SKException if the operation fails @@ -267,6 +268,7 @@ public Mono prepareAsync() { public static class Builder implements SemanticKernelBuilder> { + private DataSource dataSource; private String collectionName; private JDBCVectorStoreRecordCollectionOptions options; diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/JDBCVectorStoreRecordCollectionFactory.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/JDBCVectorStoreRecordCollectionFactory.java index 70b62a7e..6cfcdcad 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/JDBCVectorStoreRecordCollectionFactory.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/JDBCVectorStoreRecordCollectionFactory.java @@ -2,7 +2,6 @@ package com.microsoft.semantickernel.connectors.data.jdbc; import javax.sql.DataSource; -import java.sql.Connection; /** * Factory for creating JDBC vector store record collections. diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/SQLVectorStore.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/SQLVectorStore.java index 10e4d2ef..046f9941 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/SQLVectorStore.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/SQLVectorStore.java @@ -2,11 +2,10 @@ package com.microsoft.semantickernel.connectors.data.jdbc; import com.microsoft.semantickernel.data.VectorStore; -import com.microsoft.semantickernel.data.VectorStoreRecordCollection; import reactor.core.publisher.Mono; -public interface SQLVectorStore> - extends VectorStore { +public interface SQLVectorStore + extends VectorStore { /** * Prepares the vector store. diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/redis/RedisVectorStore.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/redis/RedisVectorStore.java index 7e561e43..02f728f0 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/redis/RedisVectorStore.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/redis/RedisVectorStore.java @@ -3,6 +3,7 @@ import com.microsoft.semantickernel.builders.SemanticKernelBuilder; import com.microsoft.semantickernel.data.VectorStore; +import com.microsoft.semantickernel.data.VectorStoreRecordCollection; import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; import java.util.ArrayList; @@ -12,7 +13,7 @@ import reactor.core.publisher.Mono; import redis.clients.jedis.JedisPooled; -public class RedisVectorStore implements VectorStore> { +public class RedisVectorStore implements VectorStore { private final JedisPooled client; private final RedisVectorStoreOptions options; @@ -30,6 +31,22 @@ public RedisVectorStore(@Nonnull JedisPooled client, this.options = options; } + @Override + public VectorStoreRecordCollection getCollection( + @Nonnull String collectionName, + @Nonnull Class keyClass, + @Nonnull Class recordClass, + @Nullable VectorStoreRecordDefinition recordDefinition) { + if (keyClass != String.class) { + throw new IllegalArgumentException("Redis only supports string keys"); + } + + return (VectorStoreRecordCollection) getCollection( + collectionName, + recordClass, + recordDefinition); + } + /** * Gets a collection from the vector store. * @@ -38,8 +55,7 @@ public RedisVectorStore(@Nonnull JedisPooled client, * @param recordDefinition The record definition. * @return The collection. */ - @Override - public RedisVectorStoreRecordCollection getCollection( + public RedisVectorStoreRecordCollection getCollection( @Nonnull String collectionName, @Nonnull Class recordClass, @Nullable VectorStoreRecordDefinition recordDefinition) { @@ -74,7 +90,6 @@ public Mono> getCollectionNamesAsync() { /** * Builder for the Redis vector store. - * */ public static Builder builder() { return new Builder(); diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/redis/RedisVectorStoreRecordCollection.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/redis/RedisVectorStoreRecordCollection.java index 52d30bf8..b0f8858b 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/redis/RedisVectorStoreRecordCollection.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/redis/RedisVectorStoreRecordCollection.java @@ -12,6 +12,17 @@ import com.microsoft.semantickernel.data.recordoptions.GetRecordOptions; import com.microsoft.semantickernel.data.recordoptions.UpsertRecordOptions; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; +import java.util.AbstractMap.SimpleEntry; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Map.Entry; +import java.util.stream.Collectors; +import javax.annotation.Nonnull; import org.json.JSONArray; import org.json.JSONObject; import reactor.core.publisher.Mono; @@ -25,18 +36,6 @@ import redis.clients.jedis.search.IndexOptions; import redis.clients.jedis.search.Schema; -import javax.annotation.Nonnull; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Collection; -import java.util.Collections; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Map.Entry; -import java.util.AbstractMap.SimpleEntry; -import java.util.stream.Collectors; - public class RedisVectorStoreRecordCollection implements VectorStoreRecordCollection { @@ -143,7 +142,7 @@ public Mono collectionExistsAsync() { * @return A Mono representing the completion of the creation operation. */ @Override - public Mono createCollectionAsync() { + public Mono> createCollectionAsync() { return Mono.fromRunnable(() -> { Schema schema = RedisVectorStoreCollectionCreateMapping .mapToSchema(recordDefinition.getAllFields()); @@ -155,17 +154,19 @@ public Mono createCollectionAsync() { collectionName, IndexOptions.defaultOptions().setDefinition(indexDefinition), schema); - }).subscribeOn(Schedulers.boundedElastic()).then(); + }) + .subscribeOn(Schedulers.boundedElastic()) + .then(Mono.just(this)); } @Override - public Mono createCollectionIfNotExistsAsync() { + public Mono> createCollectionIfNotExistsAsync() { return collectionExistsAsync().flatMap(exists -> { if (!exists) { return createCollectionAsync(); } - return Mono.empty(); + return Mono.just(this); }); } @@ -203,7 +204,7 @@ private JsonNode removeRedisPathPrefix(JSONObject object) { /** * Gets a record from the store. * - * @param key The key of the record to get. + * @param key The key of the record to get. * @param options The options for getting the record. * @return A Mono emitting the record. */ @@ -243,7 +244,7 @@ public Mono getAsync(String key, GetRecordOptions options) { /** * Gets a batch of records from the store. * - * @param keys The keys of the records to get. + * @param keys The keys of the records to get. * @param options The options for getting the records. * @return A Mono emitting a list of records. */ @@ -336,7 +337,7 @@ public Mono> upsertBatchAsync(List data, UpsertRecordOption /** * Deletes a record from the store. * - * @param key The key of the record to delete. + * @param key The key of the record to delete. * @param options The options for deleting the record. * @return A Mono representing the completion of the deletion operation. */ diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/redis/RedisVectorStoreRecordCollectionFactory.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/redis/RedisVectorStoreRecordCollectionFactory.java index 45417980..df1e7544 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/redis/RedisVectorStoreRecordCollectionFactory.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/redis/RedisVectorStoreRecordCollectionFactory.java @@ -1,19 +1,20 @@ // Copyright (c) Microsoft. All rights reserved. package com.microsoft.semantickernel.connectors.data.redis; +import com.azure.search.documents.indexes.SearchIndexAsyncClient; import redis.clients.jedis.JedisPooled; /** * Factory for creating Redis vector store record collections. - * */ public interface RedisVectorStoreRecordCollectionFactory { + /** * Creates a new vector store record collection. * - * @param client The Redis client. + * @param client The Redis client. * @param collectionName The name of the collection. - * @param options The options for the collection. + * @param options The options for the collection. * @return The collection. */ RedisVectorStoreRecordCollection createVectorStoreRecordCollection( diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/redis/RedisVectorStoreRecordMapper.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/redis/RedisVectorStoreRecordMapper.java index a4f5f798..f33a6b1d 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/redis/RedisVectorStoreRecordMapper.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/redis/RedisVectorStoreRecordMapper.java @@ -1,17 +1,15 @@ // Copyright (c) Microsoft. All rights reserved. package com.microsoft.semantickernel.connectors.data.redis; -import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.node.ObjectNode; import com.microsoft.semantickernel.builders.SemanticKernelBuilder; import com.microsoft.semantickernel.data.VectorStoreRecordMapper; import com.microsoft.semantickernel.exceptions.SKException; - -import javax.annotation.Nullable; import java.util.AbstractMap; import java.util.Map.Entry; import java.util.function.Function; +import javax.annotation.Nullable; public class RedisVectorStoreRecordMapper extends VectorStoreRecordMapper> { diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VectorStore.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VectorStore.java index d778829a..a51f044c 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VectorStore.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VectorStore.java @@ -1,31 +1,28 @@ // Copyright (c) Microsoft. All rights reserved. package com.microsoft.semantickernel.data; -import com.microsoft.semantickernel.data.VectorStoreRecordCollection; import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition; -import reactor.core.publisher.Mono; - +import java.util.List; import javax.annotation.Nonnull; import javax.annotation.Nullable; -import java.util.List; +import reactor.core.publisher.Mono; /** * Represents a vector store. - * - * @param The type of the record collection. */ -public interface VectorStore> { +public interface VectorStore { /** * Gets a collection from the vector store. * - * @param collectionName The name of the collection. - * @param recordClass The class type of the record. + * @param collectionName The name of the collection. + * @param recordClass The class type of the record. * @param recordDefinition The record definition. * @return The collection. */ - RecordCollection getCollection( + VectorStoreRecordCollection getCollection( @Nonnull String collectionName, + @Nonnull Class keyClass, @Nonnull Class recordClass, @Nullable VectorStoreRecordDefinition recordDefinition); diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VectorStoreRecordCollection.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VectorStoreRecordCollection.java index 1466ac35..0b7319c0 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VectorStoreRecordCollection.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VectorStoreRecordCollection.java @@ -4,11 +4,11 @@ import com.microsoft.semantickernel.data.recordoptions.DeleteRecordOptions; import com.microsoft.semantickernel.data.recordoptions.GetRecordOptions; import com.microsoft.semantickernel.data.recordoptions.UpsertRecordOptions; -import reactor.core.publisher.Mono; - import java.util.List; +import reactor.core.publisher.Mono; public interface VectorStoreRecordCollection { + /** * Gets the name of the collection. * @@ -28,14 +28,14 @@ public interface VectorStoreRecordCollection { * * @return A Mono representing the completion of the creation operation. */ - Mono createCollectionAsync(); + Mono> createCollectionAsync(); /** * Creates the collection in the store if it does not exist. * * @return A Mono representing the completion of the creation operation. */ - Mono createCollectionIfNotExistsAsync(); + Mono> createCollectionIfNotExistsAsync(); /** * Deletes the collection from the store. @@ -47,7 +47,7 @@ public interface VectorStoreRecordCollection { /** * Gets a record from the store. * - * @param key The key of the record to get. + * @param key The key of the record to get. * @param options The options for getting the record. * @return A Mono emitting the record. */ @@ -56,7 +56,7 @@ public interface VectorStoreRecordCollection { /** * Gets a batch of records from the store. * - * @param keys The keys of the records to get. + * @param keys The keys of the records to get. * @param options The options for getting the records. * @return A Mono emitting a list of records. */ @@ -65,7 +65,7 @@ public interface VectorStoreRecordCollection { /** * Inserts or updates a record in the store. * - * @param data The record to upsert. + * @param data The record to upsert. * @param options The options for upserting the record. * @return A Mono emitting the key of the upserted record. */ @@ -74,7 +74,7 @@ public interface VectorStoreRecordCollection { /** * Inserts or updates a batch of records in the store. * - * @param data The records to upsert. + * @param data The records to upsert. * @param options The options for upserting the records. * @return A Mono emitting a list of keys of the upserted records. */ @@ -83,7 +83,7 @@ public interface VectorStoreRecordCollection { /** * Deletes a record from the store. * - * @param key The key of the record to delete. + * @param key The key of the record to delete. * @param options The options for deleting the record. * @return A Mono representing the completion of the deletion operation. */ @@ -92,7 +92,7 @@ public interface VectorStoreRecordCollection { /** * Deletes a batch of records from the store. * - * @param keys The keys of the records to delete. + * @param keys The keys of the records to delete. * @param options The options for deleting the records. * @return A Mono representing the completion of the deletion operation. */ diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VectorStoreRecordMapper.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VectorStoreRecordMapper.java index 09420cf8..100c7e6d 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VectorStoreRecordMapper.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VectorStoreRecordMapper.java @@ -3,6 +3,7 @@ import com.microsoft.semantickernel.builders.SemanticKernelBuilder; +import javax.annotation.Nullable; import java.util.function.Function; /** @@ -12,6 +13,7 @@ * @param the storage model type */ public class VectorStoreRecordMapper { + @Nullable private final Function recordToStorageModelMapper; private final Function storageModelToRecordMapper; @@ -22,7 +24,7 @@ public class VectorStoreRecordMapper { * @param storageModelToRecordMapper the function to convert a storage model to a record */ protected VectorStoreRecordMapper( - Function recordToStorageModelMapper, + @Nullable Function recordToStorageModelMapper, Function storageModelToRecordMapper) { this.recordToStorageModelMapper = recordToStorageModelMapper; this.storageModelToRecordMapper = storageModelToRecordMapper; @@ -33,6 +35,7 @@ protected VectorStoreRecordMapper( * * @return the function to convert a record to a storage model */ + @Nullable public Function getRecordToStorageModelMapper() { return recordToStorageModelMapper; } diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VolatileVectorStore.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VolatileVectorStore.java index ef074247..25e90ad9 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VolatileVectorStore.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VolatileVectorStore.java @@ -2,15 +2,16 @@ package com.microsoft.semantickernel.data; import com.microsoft.semantickernel.data.recorddefinition.VectorStoreRecordDefinition; -import reactor.core.publisher.Mono; - -import javax.annotation.Nonnull; import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import reactor.core.publisher.Mono; + +public class VolatileVectorStore implements VectorStore { -public class VolatileVectorStore implements VectorStore> { private final Map> collections; public VolatileVectorStore() { @@ -25,9 +26,25 @@ public VolatileVectorStore() { * @return The collection. */ @Override - public VolatileVectorStoreRecordCollection getCollection( - @Nonnull String collectionName, @Nonnull Class recordClass, - VectorStoreRecordDefinition recordDefinition) { + public VectorStoreRecordCollection getCollection( + @Nonnull String collectionName, + @Nonnull Class keyClass, + @Nonnull Class recordClass, + @Nullable VectorStoreRecordDefinition recordDefinition) { + if (keyClass != String.class) { + throw new IllegalArgumentException("Volatile only supports string keys"); + } + + return (VectorStoreRecordCollection) getCollection( + collectionName, + recordClass, + recordDefinition); + } + + public VectorStoreRecordCollection getCollection( + @Nonnull String collectionName, + @Nonnull Class recordClass, + @Nullable VectorStoreRecordDefinition recordDefinition) { return new VolatileVectorStoreRecordCollection<>( collectionName, collections, diff --git a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VolatileVectorStoreRecordCollection.java b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VolatileVectorStoreRecordCollection.java index 7a433dbb..07d55cdd 100644 --- a/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VolatileVectorStoreRecordCollection.java +++ b/semantickernel-experimental/src/main/java/com/microsoft/semantickernel/data/VolatileVectorStoreRecordCollection.java @@ -8,18 +8,17 @@ import com.microsoft.semantickernel.data.recordoptions.GetRecordOptions; import com.microsoft.semantickernel.data.recordoptions.UpsertRecordOptions; import com.microsoft.semantickernel.exceptions.SKException; -import reactor.core.publisher.Mono; - import java.util.Collections; import java.util.HashSet; import java.util.List; import java.util.Map; -import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.stream.Collectors; +import reactor.core.publisher.Mono; + +public class VolatileVectorStoreRecordCollection implements + VectorStoreRecordCollection { -public class VolatileVectorStoreRecordCollection - implements VectorStoreRecordCollection { private static final HashSet> supportedKeyTypes = new HashSet<>( Collections.singletonList(String.class)); private Map> collections; @@ -82,8 +81,9 @@ public Mono collectionExistsAsync() { * @return A Mono representing the completion of the creation operation. */ @Override - public Mono createCollectionAsync() { - return Mono.fromRunnable(() -> collections.put(collectionName, new ConcurrentHashMap<>())); + public Mono> createCollectionAsync() { + return Mono.fromRunnable(() -> collections.put(collectionName, new ConcurrentHashMap<>())) + .then(Mono.just(this)); } /** @@ -92,9 +92,10 @@ public Mono createCollectionAsync() { * @return A Mono representing the completion of the creation operation. */ @Override - public Mono createCollectionIfNotExistsAsync() { + public Mono> createCollectionIfNotExistsAsync() { return Mono - .fromRunnable(() -> collections.putIfAbsent(collectionName, new ConcurrentHashMap<>())); + .fromRunnable(() -> collections.putIfAbsent(collectionName, new ConcurrentHashMap<>())) + .then(Mono.just(this)); } /** @@ -110,7 +111,7 @@ public Mono deleteCollectionAsync() { /** * Gets a record from the store. * - * @param key The key of the record to get. + * @param key The key of the record to get. * @param options The options for getting the record. * @return A Mono emitting the record. */ @@ -122,7 +123,7 @@ public Mono getAsync(String key, GetRecordOptions options) { /** * Gets a batch of records from the store. * - * @param keys The keys of the records to get. + * @param keys The keys of the records to get. * @param options The options for getting the records. * @return A Mono emitting a list of records. */ @@ -188,7 +189,7 @@ public Mono> upsertBatchAsync(List data, UpsertRecordOption /** * Deletes a record from the store. * - * @param key The key of the record to delete. + * @param key The key of the record to delete. * @param options The options for deleting the record. * @return A Mono representing the completion of the deletion operation. */ diff --git a/semantickernel-experimental/src/test/java/com/microsoft/semantickernel/data/VolatileVectorStoreRecordCollectionTest.java b/semantickernel-experimental/src/test/java/com/microsoft/semantickernel/data/VolatileVectorStoreRecordCollectionTest.java index 9a087adb..915b2166 100644 --- a/semantickernel-experimental/src/test/java/com/microsoft/semantickernel/data/VolatileVectorStoreRecordCollectionTest.java +++ b/semantickernel-experimental/src/test/java/com/microsoft/semantickernel/data/VolatileVectorStoreRecordCollectionTest.java @@ -1,20 +1,18 @@ // Copyright (c) Microsoft. All rights reserved. package com.microsoft.semantickernel.data; -import org.junit.jupiter.api.BeforeAll; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Order; -import org.junit.jupiter.api.Test; - -import java.util.Arrays; -import java.util.List; -import java.util.stream.Collectors; - import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; +import java.util.Arrays; +import java.util.List; +import java.util.stream.Collectors; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + public class VolatileVectorStoreRecordCollectionTest { private static VolatileVectorStoreRecordCollection recordCollection; diff --git a/semantickernel-experimental/src/test/java/com/microsoft/semantickernel/data/VolatileVectorStoreTest.java b/semantickernel-experimental/src/test/java/com/microsoft/semantickernel/data/VolatileVectorStoreTest.java index cfd52757..99c643e5 100644 --- a/semantickernel-experimental/src/test/java/com/microsoft/semantickernel/data/VolatileVectorStoreTest.java +++ b/semantickernel-experimental/src/test/java/com/microsoft/semantickernel/data/VolatileVectorStoreTest.java @@ -1,17 +1,17 @@ // Copyright (c) Microsoft. All rights reserved. package com.microsoft.semantickernel.data; -import org.junit.jupiter.api.BeforeAll; -import org.junit.jupiter.api.Test; - -import java.util.Arrays; -import java.util.List; - import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; +import java.util.Arrays; +import java.util.List; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + public class VolatileVectorStoreTest { + private static VolatileVectorStore vectorStore; @BeforeAll