>> getNearestMatchesAsync(
+ @Nonnull String collectionName,
+ @Nonnull Embedding embedding,
+ int limit,
+ float minRelevanceScore,
+ boolean withEmbedding) {
+ return getNearestMatchesAsync(
+ collectionName,
+ embedding,
+ limit,
+ minRelevanceScore,
+ withEmbedding,
+ memoryRecordMapper);
+ }
+
+ /**
+ * Gets the nearest matches to the {@link Embedding} of type {@code Float}. Does not guarantee
+ * that the collection exists.
+ *
+ * If a memoryRecordMapper is provided this overrides any mapper provided when constructing
+ * this memory store.
+ *
+ * @param collectionName The name associated with a collection of embeddings.
+ * @param embedding The {@link Embedding} to compare the collection's embeddings with.
+ * @param limit The maximum number of similarity results to return.
+ * @param minRelevanceScore The minimum relevance threshold for returned results.
+ * @param withEmbedding If true, the embeddings will be returned in the memory records.
+ * @param memoryRecordMapper a mapper that controls how to map the search document to a memory
+ * record
+ * @return A collection of tuples where item1 is a {@link MemoryRecord} and item2 is its
+ * similarity score as a {@code Float}.
+ */
+ public Mono>> getNearestMatchesAsync(
+ @Nonnull String collectionName,
+ @Nonnull Embedding embedding,
+ int limit,
+ float minRelevanceScore,
+ boolean withEmbedding,
+ @Nullable Function memoryRecordMapper) {
+ collectionName = normalizeIndexName(collectionName);
+ SearchAsyncClient client = getSearchClient(collectionName);
+
+ SearchQueryVector searchVector =
+ new SearchQueryVector()
+ .setKNearestNeighborsCount(limit)
+ .setFields(AzureCognitiveSearchMemoryRecord.EMBEDDING)
+ .setValue(embedding.getVector());
+
+ SearchOptions searchOptions = new SearchOptions().setVectors(searchVector);
+
+ return client.search(null, searchOptions)
+ .filter(result -> (double) minRelevanceScore <= result.getScore())
+ .map(
+ result -> {
+ MemoryRecord memoryRecord;
+ if (memoryRecordMapper != null) {
+ SearchDocument doc = result.getDocument(SearchDocument.class);
+ memoryRecord = memoryRecordMapper.apply(doc);
+ } else {
+ AzureCognitiveSearchMemoryRecord doc =
+ result.getDocument(AzureCognitiveSearchMemoryRecord.class);
+ memoryRecord = doc.toMemoryRecord(withEmbedding);
+ }
+ float score = (float) result.getScore();
+ return Tuples.of(memoryRecord, score);
+ })
+ .collect(Collectors.toList());
+ }
+
+ @Override
+ public Mono removeAsync(@Nonnull String collectionName, @Nonnull String key) {
+ return removeBatchAsync(collectionName, Collections.singleton(key));
+ }
+
+ @Override
+ public Mono removeBatchAsync(
+ @Nonnull String collectionName, @Nonnull Collection keys) {
+ collectionName = normalizeIndexName(collectionName);
+ SearchAsyncClient client = this.getSearchClient(collectionName);
+
+ List documents =
+ keys.stream()
+ .map(AzureCognitiveSearchMemoryRecord::encodeId)
+ .map(
+ key -> {
+ SearchDocument searchDocument = new SearchDocument();
+ searchDocument.put("Id", key); // TODO: Id should be a constant
+ return searchDocument;
+ })
+ .collect(Collectors.toList());
+
+ return client.deleteDocuments(documents).then();
+ }
+
+ static List searchFields(String configName, int embeddingSize) {
+ return Arrays.asList(
+ new SearchField(AzureCognitiveSearchMemoryRecord.ID, SearchFieldDataType.STRING)
+ .setKey(true)
+ .setFilterable(false),
+ new SearchField(
+ AzureCognitiveSearchMemoryRecord.EMBEDDING,
+ SearchFieldDataType.collection(SearchFieldDataType.SINGLE))
+ .setSearchable(true)
+ .setVectorSearchDimensions(embeddingSize)
+ .setVectorSearchConfiguration(configName),
+ new SearchField(AzureCognitiveSearchMemoryRecord.TEXT, SearchFieldDataType.STRING)
+ .setSearchable(true)
+ .setFacetable(true)
+ .setAnalyzerName(LexicalAnalyzerName.EN_LUCENE),
+ new SearchField(
+ AzureCognitiveSearchMemoryRecord.DESCRIPTION,
+ SearchFieldDataType.STRING)
+ .setSearchable(true)
+ .setFacetable(true)
+ .setAnalyzerName(LexicalAnalyzerName.EN_LUCENE),
+ new SearchField(
+ AzureCognitiveSearchMemoryRecord.ADDITIONAL_METADATA,
+ SearchFieldDataType.STRING)
+ .setSearchable(true)
+ .setFacetable(true)
+ .setAnalyzerName(LexicalAnalyzerName.EN_LUCENE),
+ new SearchField(
+ AzureCognitiveSearchMemoryRecord.EXTERNAL_SOURCE_NAME,
+ SearchFieldDataType.STRING)
+ .setSearchable(true)
+ .setFacetable(true)
+ .setAnalyzerName(LexicalAnalyzerName.EN_LUCENE),
+ new SearchField(
+ AzureCognitiveSearchMemoryRecord.IS_REFERENCE,
+ SearchFieldDataType.BOOLEAN)
+ .setFilterable(true)
+ .setFacetable(true));
+ }
+
+ /**
+ * Create a new search index.
+ *
+ * @param indexName Index name
+ * @param embeddingSize Size of the embedding vector
+ * @return A Mono that completes when the index is created
+ */
+ private Mono> createIndexAsync(
+ @Nonnull String indexName, int embeddingSize) {
+ if (embeddingSize < 1) {
+ throw new AzureCognitiveSearchMemoryException(
+ ErrorCodes.INVALID_EMBEDDING_SIZE, "the value must be greater than zero");
+ }
+
+ String normalizedIndexName = normalizeIndexName(indexName);
+
+ VectorSearch vectorSearch =
+ new VectorSearch()
+ .setAlgorithmConfigurations(
+ Collections.singletonList(
+ new HnswVectorSearchAlgorithmConfiguration(
+ SEARCH_CONFIG_NAME)
+ .setParameters(
+ new HnswParameters()
+ .setMetric(
+ VectorSearchAlgorithmMetric
+ .COSINE))));
+
+ SearchIndex newIndex =
+ new SearchIndex(normalizedIndexName)
+ .setFields(searchFields(SEARCH_CONFIG_NAME, embeddingSize))
+ .setVectorSearch(vectorSearch);
+
+ return _adminClient.createIndexWithResponse(newIndex);
+ }
+
+ private Mono> getIndexesAsync() {
+ return _adminClient.listIndexes().map(SearchIndex::getName).collect(Collectors.toList());
+ }
+
+ private Mono upsertRecordAsync(
+ @Nonnull String indexName, @Nonnull AzureCognitiveSearchMemoryRecord record) {
+ return upsertBatchAsync(indexName, Collections.singletonList(record))
+ .map(Collection::iterator)
+ .map(Iterator::next);
+ }
+
+ private Mono> upsertBatchAsync(
+ @Nonnull String indexName, @Nonnull List records) {
+
+ if (records.isEmpty()) {
+ return Mono.just(Collections.emptyList());
+ }
+
+ List documents =
+ records.stream()
+ .map(
+ record -> {
+ SearchDocument searchDocument = new SearchDocument();
+ searchDocument.put(
+ AzureCognitiveSearchMemoryRecord.ID,
+ AzureCognitiveSearchMemoryRecord.encodeId(
+ record.getId()));
+ searchDocument.put(
+ AzureCognitiveSearchMemoryRecord.TEXT,
+ record.getText());
+ searchDocument.put(
+ AzureCognitiveSearchMemoryRecord.DESCRIPTION,
+ record.getDescription());
+ searchDocument.put(
+ AzureCognitiveSearchMemoryRecord.EMBEDDING,
+ record.getEmbedding());
+ searchDocument.put(
+ AzureCognitiveSearchMemoryRecord.ADDITIONAL_METADATA,
+ record.getAdditionalMetadata());
+ searchDocument.put(
+ AzureCognitiveSearchMemoryRecord.EXTERNAL_SOURCE_NAME,
+ record.getExternalSourceName());
+ searchDocument.put(
+ AzureCognitiveSearchMemoryRecord.IS_REFERENCE,
+ record.isReference());
+ return searchDocument;
+ })
+ .collect(Collectors.toList());
+
+ IndexDocumentsOptions options = new IndexDocumentsOptions().setThrowOnAnyError(true);
+
+ int embeddingSize =
+ records.stream()
+ .map(AzureCognitiveSearchMemoryRecord::getEmbedding)
+ .map(List::size)
+ .max(Integer::compareTo)
+ .orElse(0);
+
+ String normalizedIndexName = normalizeIndexName(indexName);
+ SearchAsyncClient client = this.getSearchClient(normalizedIndexName);
+
+ return client.uploadDocumentsWithResponse(documents, options)
+ .map(
+ response -> {
+ if (response.getStatusCode() == 404) {
+ throw new AzureCognitiveSearchMemoryException(
+ ErrorCodes.INVALID_INDEX_NAME);
+ }
+ return response;
+ })
+ .onErrorResume(
+ e ->
+ createIndexAsync(normalizedIndexName, embeddingSize)
+ .then(
+ client.uploadDocumentsWithResponse(
+ documents, options)))
+ .map(
+ response ->
+ response.getValue().getResults().stream()
+ .map(IndexingResult::getKey)
+ .collect(Collectors.toList()));
+ }
+
+ /**
+ * Get a search client for the index specified. Note: the index might not exist, but we avoid
+ * checking everytime and the extra latency.
+ *
+ * @param indexName Index name
+ * @return Search client ready to read/write
+ */
+ protected SearchAsyncClient getSearchClient(@Nonnull String indexName) {
+ String normalizedIndexName = normalizeIndexName(indexName);
+ return _clientsByIndex.computeIfAbsent(
+ normalizedIndexName, _adminClient::getSearchAsyncClient);
+ }
+
+ ///
+ /// Options used by the Azure Cognitive Search client, e.g. User Agent.
+ /// See also
+ // https://github.com/Azure/azure-sdk-for-net/blob/main/sdk/core/Azure.Core/src/DiagnosticsOptions.cs
+ ///
+ private static ClientOptions clientOptions() {
+ return SemanticKernelHttpSettings.getUserAgent()
+ .setTracingOptions(new TracingOptions())
+ .setMetricsOptions(new MetricsOptions());
+ }
+
+ // Index names cannot contain special chars. We use this rule to replace a few common ones
+ // with an underscore and reduce the chance of errors. If other special chars are used, we leave
+ // it
+ // to the service to throw an error.
+ // Note:
+ // - replacing chars introduces a small chance of conflicts, e.g. "the-user" and "the_user".
+ // - we should consider whether making this optional and leave it to the developer to handle.
+ //
+ // Java 8 does not support \s for whitespace, so we have to list them all.
+ private static final String whitespace = " |\\t|\\n|\\x0B|\\f|\\r";
+ private static final String s_replaceIndexNameSymbolsRegex = "[" + whitespace + "|\\|/|.|_|:]";
+
+ // Normalize index name to match ACS rules.
+ // See https://learn.microsoft.com/en-us/rest/api/searchservice/Create-Index
+ // The method doesn't handle all the error scenarios, leaving it to the service
+ // to throw an error for edge cases not handled locally.
+ // Value to normalize
+ // Normalized name
+ private static String normalizeIndexName(String indexName) {
+ if (indexName.length() > 128) {
+ throw new IllegalArgumentException("The collection name cannot exceed 128 chars");
+ }
+
+ indexName = indexName.toLowerCase(Locale.ROOT);
+ return indexName.replaceAll(s_replaceIndexNameSymbolsRegex, "-");
+ }
+}
diff --git a/java/connectors/semantickernel-connectors-memory-azurecognitivesearch/src/test/java/com/microsoft/semantickernel/connectors/memory/azurecognitivesearch/AzureCognitiveSearchMemoryTests.java b/java/connectors/semantickernel-connectors-memory-azurecognitivesearch/src/test/java/com/microsoft/semantickernel/connectors/memory/azurecognitivesearch/AzureCognitiveSearchMemoryTests.java
new file mode 100644
index 0000000000..e1d081c828
--- /dev/null
+++ b/java/connectors/semantickernel-connectors-memory-azurecognitivesearch/src/test/java/com/microsoft/semantickernel/connectors/memory/azurecognitivesearch/AzureCognitiveSearchMemoryTests.java
@@ -0,0 +1,434 @@
+// Copyright (c) Microsoft. All rights reserved.
+package com.microsoft.semantickernel.connectors.memory.azurecognitivesearch;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import com.azure.core.http.HttpHeaders;
+import com.azure.core.http.HttpPipeline;
+import com.azure.core.http.HttpRequest;
+import com.azure.core.http.HttpResponse;
+import com.azure.core.http.rest.Response;
+import com.azure.core.util.Context;
+import com.azure.core.util.tracing.Tracer;
+import com.azure.search.documents.SearchAsyncClient;
+import com.azure.search.documents.SearchDocument;
+import com.azure.search.documents.indexes.SearchIndexAsyncClient;
+import com.azure.search.documents.models.SearchResult;
+import com.azure.search.documents.util.SearchPagedFlux;
+import com.microsoft.semantickernel.Kernel;
+import com.microsoft.semantickernel.SKBuilders;
+import com.microsoft.semantickernel.ai.embeddings.Embedding;
+import com.microsoft.semantickernel.ai.embeddings.EmbeddingGeneration;
+import com.microsoft.semantickernel.memory.MemoryQueryResult;
+import com.microsoft.semantickernel.memory.MemoryRecord;
+import com.microsoft.semantickernel.memory.MemoryRecordMetadata;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Base64;
+import java.util.List;
+import java.util.function.Function;
+import javax.annotation.Nullable;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+import org.mockito.invocation.InvocationOnMock;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+class AzureCognitiveSearchMemoryTests {
+
+ // com.azure.core.implementation.http.rest.RestProxyBase.RestProxyBase
+ // expects a non-null tracer. There is a TracerProvider in the Azure SDK
+ // that can be used to create a tracer, but it has some ceremony around
+ // it. We create a simple no-op tracer instead.
+ private static class NoOpTracer implements Tracer {
+
+ @Override
+ public Context start(String methodName, Context context) {
+ return context;
+ }
+
+ @Override
+ public void end(String errorMessage, Throwable throwable, Context context) {}
+
+ @Override
+ public void setAttribute(String key, String value, Context context) {}
+
+ @Override
+ public boolean isEnabled() {
+ // we could also set a property to disable tracing, but this is simpler.
+ return false;
+ }
+ }
+
+ // Kernel is configured with a real AzureCognitiveSearchMemoryStore which is backed
+ // by a mocked HttpPipeline. The HttpPipeline is mocked to return a mocked
+ // HttpResponse whose body and status code are configured by the test.
+ Kernel kernel;
+
+ //
+ // A test may generate more than one HttpRequest. This function
+ // can be set in the test to return the appropriate response body
+ // for each request.
+ //
+ // A good place to set a breakpoint for debugging exceptions in
+ // deserializing the response body is:
+ // com.azure.search.documents.implementation.util.MappingUtils#exceptionMapper
+ //
+ Function responseBody;
+
+ //
+ // A test may generate more than one HttpRequest. This function
+ // can be set in the test to return the appropriate HTTP status
+ // code for each request.
+ //
+ Function responseStatusCode;
+
+ @BeforeEach
+ void setUp() {
+ kernel = buildKernel(null);
+ }
+
+ private Kernel buildKernel(
+ @Nullable Function memoryRecordMapper) {
+ SearchAsyncClient searchAsyncClient = Mockito.mock(SearchAsyncClient.class);
+ SearchIndexAsyncClient searchIndexAsyncClient = Mockito.mock(SearchIndexAsyncClient.class);
+ Response response = Mockito.mock(Response.class);
+ AzureCognitiveSearchMemoryRecord azureCognitiveSearchMemoryRecord =
+ Mockito.mock(AzureCognitiveSearchMemoryRecord.class);
+ MemoryRecord memoryRecord = Mockito.mock(MemoryRecord.class);
+ SearchPagedFlux searchPagedFlux = Mockito.mock(SearchPagedFlux.class);
+ SearchResult searchResult = Mockito.mock(SearchResult.class);
+
+ Mockito.when(searchResult.getDocument(Mockito.any())).thenReturn(new SearchDocument());
+ Mockito.when(searchIndexAsyncClient.getSearchAsyncClient(Mockito.any()))
+ .thenReturn(searchAsyncClient);
+
+ Mockito.when(searchPagedFlux.filter(Mockito.any())).thenReturn(Flux.just(searchResult));
+
+ Mockito.when(
+ searchAsyncClient.getDocumentWithResponse(
+ Mockito.any(), Mockito.any(), Mockito.any()))
+ .thenReturn(Mono.just(response));
+
+ Mockito.when(searchAsyncClient.search(Mockito.any(), Mockito.any()))
+ .thenReturn(searchPagedFlux);
+ Mockito.when(response.getValue()).thenReturn(azureCognitiveSearchMemoryRecord);
+ Mockito.when(azureCognitiveSearchMemoryRecord.toMemoryRecord(Mockito.anyBoolean()))
+ .thenReturn(memoryRecord);
+
+ EmbeddingGeneration embeddingGeneration =
+ new EmbeddingGeneration() {
+ /*
+ new SearchIndexClientBuilder()
+ .endpoint("https://fake-random-test-host/fake-path")
+ .credential(new AzureKeyCredential(apiKey))
+ .pipeline(pipeline)
+ .httpLogOptions(
+ new HttpLogOptions()
+ .setLogLevel(HttpLogDetailLevel.BODY_AND_HEADERS))
+ .buildAsyncClient();
+
+ */
+
+ List embeddings = new ArrayList<>();
+
+ {
+ embeddings.add(new Embedding(Arrays.asList(1.0f, 2.0f, 3.0f)));
+ }
+
+ @Override
+ public Mono> generateEmbeddingsAsync(List data) {
+ return Mono.just(embeddings);
+ }
+ };
+
+ return SKBuilders.kernel()
+ .withMemoryStorage(
+ new AzureCognitiveSearchMemoryStore(
+ searchIndexAsyncClient, memoryRecordMapper))
+ .withDefaultAIService(embeddingGeneration)
+ .build();
+ }
+
+ private HttpPipeline setUpHttpPipeline() {
+
+ HttpPipeline pipeline = Mockito.mock(HttpPipeline.class);
+
+ Tracer tracer = new NoOpTracer();
+ Mockito.when(pipeline.getTracer()).thenReturn(tracer);
+
+ Mockito.when(pipeline.send(Mockito.any(HttpRequest.class), Mockito.any(Context.class)))
+ .thenAnswer(this::answerHttpRequest);
+
+ return pipeline;
+ }
+
+ private Mono answerHttpRequest(InvocationOnMock invocation) {
+ HttpRequest request = invocation.getArgument(0);
+
+ String query = request.getUrl().getQuery();
+ int offset = query.indexOf("api-version");
+ assertTrue(offset >= 0);
+ // HttpResponse is based on this api-version
+ assertEquals("api-version=2023-07-01-Preview", query.substring(offset));
+
+ // System.err.println(request.getHttpMethod() + " " + request.getUrl() + "\n\t" +
+ // request.getHeaders());
+ byte[] body = responseBody.apply(request).getBytes(StandardCharsets.UTF_8);
+ int statusCode = responseStatusCode.apply(request);
+ HttpHeaders httpHeaders =
+ new HttpHeaders().add("Content-Type", "application/json; charset=utf-8");
+
+ HttpResponse httpResponse = Mockito.mock(HttpResponse.class);
+ Mockito.when(httpResponse.getRequest()).thenReturn(request);
+ Mockito.when(httpResponse.getHeaders()).thenReturn(httpHeaders);
+ Mockito.when(httpResponse.getBodyAsByteArray()).thenReturn(Mono.just(body));
+ Mockito.when(httpResponse.getStatusCode()).thenReturn(statusCode);
+ return Mono.just(httpResponse);
+ }
+
+ @Disabled("Pending https://github.com/microsoft/semantic-kernel/issues/1797")
+ @Test
+ void azureCognitiveSearchMemoryStoreShouldBeProperlyInitialized() {
+ // Arrange
+ responseBody = httpRequest -> "{\"value\": [{\"name\": \"fake-index\"}]}";
+ responseStatusCode = httpResponse -> 200;
+
+ // Act
+ // This call triggers a subsequent call to Azure Cognitive Search Memory store.
+ List collections = kernel.getMemory().getCollectionsAsync().block();
+
+ // Assert
+ assertTrue(collections.contains("fake-index"));
+ }
+
+ @Disabled("Pending https://github.com/microsoft/semantic-kernel/issues/1797")
+ @Test
+ void azureCognitiveSearchMemoryStoreCanSaveReferenceAsync() {
+ // Arrange
+ responseBody =
+ httpRequest ->
+ "{\"value\": ["
+ + " {"
+ + " \"key\": \"fake-key\","
+ + " \"status\": true,"
+ + " \"errorMessage\": null,"
+ + " \"statusCode\": 200"
+ + " }"
+ + " ]"
+ + "}";
+
+ responseStatusCode = httpRequest -> 200;
+
+ // Act
+ String key =
+ kernel.getMemory()
+ .saveReferenceAsync(
+ "fake-index",
+ "fake-text",
+ "fake-externalId",
+ "fake-externalSourceName",
+ "fake-description",
+ "fake-additionalMetadata")
+ .block();
+
+ // Assert
+ assertEquals("fake-key", key);
+ }
+
+ @Disabled("Pending https://github.com/microsoft/semantic-kernel/issues/1797")
+ @Test
+ void azureCognitiveSearchMemoryStoreCanSaveInformationAsync() {
+ // Arrange
+ responseBody =
+ httpRequest ->
+ "{\"value\": ["
+ + " {"
+ + " \"key\": \"fake-key\","
+ + " \"status\": true,"
+ + " \"errorMessage\": null,"
+ + " \"statusCode\": 200"
+ + " }"
+ + " ]"
+ + "}";
+
+ responseStatusCode = httpRequest -> 200;
+
+ // Act
+ String key =
+ kernel.getMemory()
+ .saveInformationAsync(
+ "fake-index",
+ "fake-text",
+ "fake-externalId",
+ "fake-description",
+ "fake-additionalMetadata")
+ .block();
+
+ // Assert
+ assertEquals("fake-key", key);
+ }
+
+ @Disabled("Pending https://github.com/microsoft/semantic-kernel/issues/1797")
+ @Test
+ void azureCognitiveSearchMemoryStoreGetAsync() {
+ // Arrange
+ String key = "fake-key";
+ String base64EncodedKey =
+ new String(
+ Base64.getUrlEncoder().encode("fake-key".getBytes(StandardCharsets.UTF_8)));
+
+ responseBody =
+ httpRequest ->
+ "{"
+ + " \"Id\": \""
+ + base64EncodedKey
+ + "\","
+ + " \"Text\": \"fake-text\","
+ + " \"Description\": \"fake-description\","
+ + " \"AdditionalMetadata\": \"fake-additionalmetadata\","
+ + " \"ExternalSourceName\": \"fake-externalsourcename\","
+ + " \"Reference\": false"
+ + "}";
+
+ responseStatusCode = httpRequest -> 200;
+
+ // Act
+ MemoryQueryResult result = kernel.getMemory().getAsync("fake-index", key, false).block();
+
+ // Assert
+ assertEquals(key, result.getMetadata().getId());
+ }
+
+ @Disabled("Pending https://github.com/microsoft/semantic-kernel/issues/1797")
+ @Test
+ void azureCognitiveSearchMemoryStoreRemoveAsync() {
+ // Arrange
+ String key = "fake-key";
+ String base64EncodedKey =
+ new String(
+ Base64.getUrlEncoder().encode("fake-key".getBytes(StandardCharsets.UTF_8)));
+
+ responseBody =
+ httpRequest ->
+ "{\"value\": ["
+ + " {"
+ + " \"key\": \""
+ + base64EncodedKey
+ + "\","
+ + " \"status\": true,"
+ + " \"errorMessage\": null,"
+ + " \"statusCode\": 200"
+ + " }"
+ + " ]"
+ + "}";
+ ;
+
+ responseStatusCode = httpRequest -> 200;
+
+ // Act
+ kernel.getMemory().removeAsync("fake-index", key).block();
+
+ // Assert
+ }
+
+ @Disabled("Pending https://github.com/microsoft/semantic-kernel/issues/1797")
+ @Test
+ void azureCognitiveSearchMemoryStoreSearchAsync() {
+ // Arrange
+ String key = "fake-key";
+ String base64EncodedKey =
+ new String(
+ Base64.getUrlEncoder().encode("fake-key".getBytes(StandardCharsets.UTF_8)));
+
+ responseBody =
+ httpRequest ->
+ "{\"value\": ["
+ + "{"
+ + " \"@search.score\": 1,"
+ + " \"@search.highlights\": {"
+ + " \"Text\": [\"It was a dark and stormy night\"]"
+ + " },"
+ + " \"@search.features\": {"
+ + " \"Text\": {"
+ + " \"uniqueTokenMatches\": 1,"
+ + " \"similarityScore\": 0.8,"
+ + " \"termFrequency\": 0.99"
+ + " }"
+ + " },"
+ + " \"Id\": \""
+ + base64EncodedKey
+ + "\","
+ + " \"Text\": \"fake-text\","
+ + " \"Description\": \"fake-description\","
+ + " \"AdditionalMetadata\": \"fake-additionalmetadata\","
+ + " \"ExternalSourceName\": \"fake-externalsourcename\","
+ + " \"Reference\": false"
+ + "}"
+ + "]}";
+
+ responseStatusCode = httpRequest -> 200;
+
+ // Act
+ List results =
+ kernel.getMemory().searchAsync("fake-index", "fake-query", 1, 0.5f, false).block();
+
+ // Assert
+ assertEquals(1, results.size());
+ assertEquals(key, results.get(0).getMetadata().getId());
+ }
+
+ @Test
+ void azureCognitiveSearchMemoryStoreGetAsyncWithCustomMapping() {
+ // Do not convert to lambda, it will break mockito
+ Function transform =
+ Mockito.spy(
+ new Function() {
+ @Override
+ public MemoryRecord apply(SearchDocument searchDocument) {
+ return new MemoryRecord(
+ new MemoryRecordMetadata(
+ false,
+ "an-id",
+ "fake-text",
+ "fake-description",
+ "fake-additionalmetadata",
+ "fake-externalsourcename"),
+ new Embedding(Arrays.asList(1.0f, 2.0f, 3.0f)),
+ "an-id",
+ null);
+ }
+ });
+
+ kernel = buildKernel(transform);
+ // Arrange
+ String key = "fake-key";
+ String base64EncodedKey =
+ new String(
+ Base64.getUrlEncoder().encode("fake-key".getBytes(StandardCharsets.UTF_8)));
+
+ responseBody =
+ httpRequest ->
+ "{"
+ + " \"Id\": \""
+ + base64EncodedKey
+ + "\","
+ + " \"Text\": \"fake-text\","
+ + " \"Description\": \"fake-description\","
+ + " \"AdditionalMetadata\": \"fake-additionalmetadata\","
+ + " \"ExternalSourceName\": \"fake-externalsourcename\","
+ + " \"Reference\": false"
+ + "}";
+
+ responseStatusCode = httpRequest -> 200;
+
+ // Act
+ kernel.getMemory().searchAsync("fake-index", "a-query", 1, 0.0f, false).block();
+
+ Mockito.verify(transform, Mockito.atLeastOnce()).apply(Mockito.any());
+ }
+}
diff --git a/java/connectors/semantickernel-connectors-memory-jdbc/pom.xml b/java/connectors/semantickernel-connectors-memory-jdbc/pom.xml
new file mode 100644
index 0000000000..eaa692e290
--- /dev/null
+++ b/java/connectors/semantickernel-connectors-memory-jdbc/pom.xml
@@ -0,0 +1,56 @@
+
+
+ 4.0.0
+
+ com.microsoft.semantic-kernel
+ semantickernel-parent
+ 0.2.10-alpha-SNAPSHOT
+ ../../pom.xml
+
+
+ semantickernel-connectors-memory-jdbc
+ Semantic Kernel JDBC Memory Connector
+
+
+
+ com.fasterxml.jackson.core
+ jackson-core
+ compile
+
+
+ com.fasterxml.jackson.core
+ jackson-databind
+ compile
+
+
+
+ com.microsoft.semantic-kernel
+ semantickernel-api
+
+
+ org.mockito
+ mockito-core
+
+
+ org.junit.jupiter
+ junit-jupiter
+
+
+ org.apache.logging.log4j
+ log4j-api
+ test
+
+
+ org.apache.logging.log4j
+ log4j-core
+ test
+
+
+ com.microsoft.semantic-kernel
+ semantickernel-api
+
+
+
+
+
+
\ No newline at end of file
diff --git a/java/connectors/semantickernel-connectors-memory-jdbc/semantickernel-connectors-memory-jdbc.iml b/java/connectors/semantickernel-connectors-memory-jdbc/semantickernel-connectors-memory-jdbc.iml
new file mode 100644
index 0000000000..224fc7d9a8
--- /dev/null
+++ b/java/connectors/semantickernel-connectors-memory-jdbc/semantickernel-connectors-memory-jdbc.iml
@@ -0,0 +1,10 @@
+
+
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/DatabaseEntry.java b/java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/DatabaseEntry.java
new file mode 100644
index 0000000000..05941fdc01
--- /dev/null
+++ b/java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/DatabaseEntry.java
@@ -0,0 +1,43 @@
+// Copyright (c) Microsoft. All rights reserved.
+package com.microsoft.semantickernel.connectors.memory.jdbc;
+
+import com.microsoft.semantickernel.memory.DataEntryBase;
+import java.time.ZonedDateTime;
+
+/** Represents an entry in the Semantic Kernel Memory Table. */
+public class DatabaseEntry extends DataEntryBase {
+ private final String embedding; // JSON with the embedding information
+ private final String metadata; // JSON with the metadata associated
+
+ /**
+ * Creates a new instance of DatabaseEntry.
+ *
+ * @param key The key identifying the entry.
+ * @param metadata The metadata associated.
+ * @param embedding The embedding information associated.
+ * @param timestamp The timestamp indicating when the entry was created or modified.
+ */
+ public DatabaseEntry(String key, String metadata, String embedding, ZonedDateTime timestamp) {
+ super(key, timestamp);
+ this.metadata = metadata;
+ this.embedding = embedding;
+ }
+
+ /**
+ * Gets the metadata associated with the entry.
+ *
+ * @return The metadata associated with the entry.
+ */
+ public String getMetadata() {
+ return metadata;
+ }
+
+ /**
+ * Gets the embedding information associated with the entry.
+ *
+ * @return The embedding information associated with the entry.
+ */
+ public String getEmbedding() {
+ return embedding;
+ }
+}
diff --git a/java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/JDBCConnector.java b/java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/JDBCConnector.java
new file mode 100644
index 0000000000..c16922048a
--- /dev/null
+++ b/java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/JDBCConnector.java
@@ -0,0 +1,495 @@
+// Copyright (c) Microsoft. All rights reserved.
+package com.microsoft.semantickernel.connectors.memory.jdbc;
+
+import com.microsoft.semantickernel.memory.MemoryException;
+import com.microsoft.semantickernel.memory.MemoryException.ErrorCodes;
+import java.io.Closeable;
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.time.ZonedDateTime;
+import java.time.format.DateTimeFormatter;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import javax.annotation.Nullable;
+import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Schedulers;
+
+public class JDBCConnector implements SQLConnector, Closeable {
+ protected final Connection connection;
+
+ public JDBCConnector(Connection connection) {
+ this.connection = connection;
+ }
+
+ // Convenience method to format a ZonedDateTime in a format acceptable to SQL
+ protected static String formatDatetime(@Nullable ZonedDateTime datetime) {
+ if (datetime == null) return "";
+ return datetime.format(DateTimeFormatter.ISO_OFFSET_DATE_TIME);
+ }
+
+ // Convenience method to parse a SQL datetime string into a ZonedDateTime
+ protected static ZonedDateTime parseDatetime(@Nullable String datetime) {
+ if (datetime == null || datetime.isEmpty()) return null;
+ return ZonedDateTime.parse(datetime, DateTimeFormatter.ISO_OFFSET_DATE_TIME);
+ }
+
+ protected static final String COLLECTIONS_TABLE_NAME = "SKCollectionTable";
+ protected static final String TABLE_NAME = "SKMemoryTable";
+ protected static final String INDEX_NAME = "SKMemoryIndex";
+
+ /**
+ * Returns the name of the table that stores the collection names.
+ *
+ * @return The name of the table that stores the collection names.
+ */
+ protected static String DEFAULT_COLLECTIONS_TABLE_NAME() {
+ return COLLECTIONS_TABLE_NAME;
+ }
+
+ /**
+ * Returns the name of the Semantic Kernel table.
+ *
+ * @return The name of the Semantic Kernel table.
+ */
+ protected static String DEFAULT_TABLE_NAME() {
+ return TABLE_NAME;
+ }
+
+ /**
+ * Returns the name of the index on the Semantic Kernel table.
+ *
+ * @return The name of the index on the Semantic Kernel table.
+ */
+ protected static String DEFAULT_INDEX_NAME() {
+ return INDEX_NAME;
+ }
+
+ public Mono createTableAsync() {
+ return Mono.fromRunnable(
+ () -> {
+ String createCollectionKeyTable =
+ "CREATE TABLE IF NOT EXISTS "
+ + COLLECTIONS_TABLE_NAME
+ + " ("
+ + "id TEXT PRIMARY KEY"
+ + " )";
+
+ String createSKMemoryTable =
+ "CREATE TABLE IF NOT EXISTS "
+ + TABLE_NAME
+ + " ("
+ + "collection TEXT NOT NULL, "
+ + "key TEXT NOT NULL, "
+ + "metadata TEXT, "
+ + "embedding TEXT, "
+ + "timestamp TEXT, "
+ + "PRIMARY KEY (collection, key), "
+ + "FOREIGN KEY (collection) REFERENCES "
+ + COLLECTIONS_TABLE_NAME
+ + "(id)"
+ + " )";
+
+ String createIndex =
+ "CREATE INDEX IF NOT EXISTS "
+ + INDEX_NAME
+ + " ON "
+ + TABLE_NAME
+ + "(collection)";
+
+ try (Statement statement = this.connection.createStatement()) {
+ statement.addBatch(createCollectionKeyTable);
+ statement.addBatch(createSKMemoryTable);
+ statement.addBatch(createIndex);
+ statement.executeBatch();
+ } catch (SQLException e) {
+ throw new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "\"CREATE TABLE\" failed",
+ e);
+ }
+ })
+ .subscribeOn(Schedulers.boundedElastic())
+ .then();
+ }
+
+ public Mono createCollectionAsync(String collectionName) {
+ return Mono.fromRunnable(
+ () -> {
+ String query =
+ "INSERT INTO " + COLLECTIONS_TABLE_NAME + " (id) VALUES (?)";
+ try (PreparedStatement statement =
+ this.connection.prepareStatement(query)) {
+ statement.setString(1, collectionName);
+ statement.executeUpdate();
+ } catch (SQLException e) {
+ throw new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "\"INSERT INTO\" failed",
+ e);
+ }
+ })
+ .subscribeOn(Schedulers.boundedElastic())
+ .then();
+ }
+
+ public Mono upsertAsync(
+ String collection,
+ String key,
+ String metadata,
+ String embedding,
+ ZonedDateTime timestamp) {
+ return Mono.fromRunnable(
+ () -> {
+ String query =
+ "INSERT OR REPLACE INTO "
+ + TABLE_NAME
+ + " (collection, key, metadata, embedding, timestamp)"
+ + " VALUES (?, ?, ?, ?, ?)";
+ try (PreparedStatement statement =
+ this.connection.prepareStatement(query)) {
+ statement.setString(1, collection);
+ statement.setString(2, key);
+ statement.setString(3, metadata != null ? metadata : "");
+ statement.setString(4, embedding != null ? embedding : "");
+ statement.setString(5, formatDatetime(timestamp));
+ statement.executeUpdate();
+ } catch (SQLException e) {
+ throw new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "\"INSERT OR REPLACE INTO\" failed",
+ e);
+ }
+ })
+ .subscribeOn(Schedulers.boundedElastic())
+ .thenReturn(key);
+ }
+
+ @Override
+ public Mono> upsertBatchAsync(
+ String collection, Collection records) {
+ Collection keys = new ArrayList<>();
+ return Mono.fromRunnable(
+ () -> {
+ String query =
+ "INSERT OR REPLACE INTO "
+ + TABLE_NAME
+ + " (collection, key, metadata, embedding, timestamp)"
+ + " VALUES (?, ?, ?, ?, ?)";
+ try (PreparedStatement statement =
+ this.connection.prepareStatement(query)) {
+ for (DatabaseEntry entry : records) {
+ statement.setString(1, collection);
+ statement.setString(2, entry.getKey());
+ statement.setString(
+ 3,
+ entry.getMetadata() != null ? entry.getMetadata() : "");
+ statement.setString(
+ 4,
+ entry.getEmbedding() != null
+ ? entry.getEmbedding()
+ : "");
+ statement.setString(5, formatDatetime(entry.getTimestamp()));
+ statement.addBatch();
+ keys.add(entry.getKey());
+ }
+
+ statement.executeBatch();
+ } catch (SQLException e) {
+ throw new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "\"INSERT OR REPLACE INTO\" failed",
+ e);
+ }
+ })
+ .subscribeOn(Schedulers.boundedElastic())
+ .thenReturn(keys);
+ }
+
+ public Mono doesCollectionExistsAsync(String collectionName) {
+ return Mono.fromCallable(() -> doesCollectionExists(collectionName))
+ .subscribeOn(Schedulers.boundedElastic());
+ }
+
+ private boolean doesCollectionExists(String collectionName) throws SQLException {
+ String query = "SELECT id FROM " + COLLECTIONS_TABLE_NAME + " WHERE id = ?";
+ try (PreparedStatement statement = this.connection.prepareStatement(query)) {
+ statement.setString(1, collectionName);
+ ResultSet resultSet = statement.executeQuery();
+ return resultSet.next();
+ }
+ }
+
+ public Mono> getCollectionsAsync() {
+ return Mono.defer(
+ () -> {
+ List collections = new ArrayList<>();
+ String query = "SELECT id FROM " + COLLECTIONS_TABLE_NAME;
+ try (PreparedStatement statement =
+ this.connection.prepareStatement(query)) {
+ ResultSet resultSet = statement.executeQuery();
+ while (resultSet.next()) {
+ String collection = resultSet.getString("id");
+ collections.add(collection);
+ }
+ } catch (SQLException e) {
+ return Mono.error(
+ new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "\"SELECT\" failed",
+ e));
+ }
+ return Mono.just(collections);
+ })
+ .subscribeOn(Schedulers.boundedElastic());
+ }
+
+ public Mono> readAllAsync(String collectionName) {
+ return Mono.defer(
+ () -> {
+ List entries = new ArrayList<>();
+ String query = "SELECT * FROM " + TABLE_NAME + " WHERE collection = ?";
+ try (PreparedStatement statement =
+ this.connection.prepareStatement(query)) {
+ statement.setString(1, collectionName);
+ ResultSet resultSet = statement.executeQuery();
+ while (resultSet.next()) {
+ String key = resultSet.getString("key");
+ String metadata = resultSet.getString("metadata");
+ String embedding = resultSet.getString("embedding");
+ String timestamp = resultSet.getString("timestamp");
+ ZonedDateTime zonedDateTime = parseDatetime(timestamp);
+ entries.add(
+ new DatabaseEntry(
+ key, metadata, embedding, zonedDateTime));
+ }
+ } catch (SQLException e) {
+ return Mono.error(
+ new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "\"SELECT * FROM\" failed",
+ e));
+ }
+ return Mono.just(entries);
+ })
+ .subscribeOn(Schedulers.boundedElastic());
+ }
+
+ public Mono readAsync(String collectionName, String key) {
+ return Mono.defer(
+ () -> {
+ String query =
+ "SELECT * FROM "
+ + TABLE_NAME
+ + " WHERE collection = ?"
+ + " AND key = ?";
+ try (PreparedStatement statement =
+ this.connection.prepareStatement(query)) {
+ statement.setString(1, collectionName);
+ statement.setString(2, key != null && !key.isEmpty() ? key : null);
+ ResultSet resultSet = statement.executeQuery();
+ if (resultSet.next()) {
+ String metadata = resultSet.getString("metadata");
+ String embedding = resultSet.getString("embedding");
+ String timestamp = resultSet.getString("timestamp");
+ ZonedDateTime zonedDateTime = parseDatetime(timestamp);
+ return Mono.just(
+ new DatabaseEntry(
+ key, metadata, embedding, zonedDateTime));
+ }
+ } catch (SQLException e) {
+ return Mono.error(
+ new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "\"SELECT * FROM\" failed",
+ e));
+ }
+ return Mono.empty();
+ })
+ .subscribeOn(Schedulers.boundedElastic());
+ }
+
+ protected enum BatchOperation {
+ SELECT,
+ DELETE
+ }
+
+ protected String batchQuery(BatchOperation operation, Collection keys) {
+ String queryPrefix;
+ switch (operation) {
+ case SELECT:
+ queryPrefix = "SELECT * FROM " + TABLE_NAME + " WHERE collection = ?";
+ break;
+ case DELETE:
+ queryPrefix = "DELETE FROM " + TABLE_NAME + " WHERE collection = ?";
+ break;
+ default:
+ throw new IllegalArgumentException("Invalid batch operation");
+ }
+ StringBuilder queryBuilder = new StringBuilder(queryPrefix);
+ queryBuilder.append(" AND key IN (");
+
+ // Add placeholders for each key
+ for (int i = 0; i < keys.size(); i++) {
+ queryBuilder.append("?");
+ if (i < keys.size() - 1) {
+ queryBuilder.append(",");
+ }
+ }
+
+ queryBuilder.append(")");
+ return queryBuilder.toString();
+ }
+
+ @Override
+ public Mono> readBatchAsync(
+ String collectionName, Collection keys) {
+ return Mono.defer(
+ () -> {
+ Collection entries = new ArrayList<>();
+ String query = batchQuery(BatchOperation.SELECT, keys);
+
+ try (PreparedStatement statement =
+ this.connection.prepareStatement(query)) {
+ statement.setString(1, collectionName);
+ int index = 2;
+ for (String key : keys) {
+ statement.setString(index++, key);
+ }
+ ResultSet resultSet = statement.executeQuery();
+
+ while (resultSet.next()) {
+ String key = resultSet.getString("key");
+ String metadata = resultSet.getString("metadata");
+ String embedding = resultSet.getString("embedding");
+ String timestamp = resultSet.getString("timestamp");
+ ZonedDateTime zonedDateTime = parseDatetime(timestamp);
+ entries.add(
+ new DatabaseEntry(
+ key, metadata, embedding, zonedDateTime));
+ }
+ } catch (SQLException e) {
+ return Mono.error(
+ new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "\"SELECT * FROM\" failed",
+ e));
+ }
+ return Mono.just(entries);
+ })
+ .subscribeOn(Schedulers.boundedElastic());
+ }
+
+ public Mono deleteCollectionAsync(String collectionName) {
+ return Mono.fromRunnable(
+ () -> {
+ String query1 = "DELETE FROM " + TABLE_NAME + " WHERE collection = ?";
+ String query2 =
+ "DELETE FROM " + COLLECTIONS_TABLE_NAME + " WHERE id = ?";
+ try (PreparedStatement statement =
+ this.connection.prepareStatement(query1);
+ PreparedStatement statement2 =
+ this.connection.prepareStatement(query2)) {
+ statement.setString(1, collectionName);
+ statement.executeUpdate();
+ statement2.setString(1, collectionName);
+ if (statement2.executeUpdate() == 0) {
+ throw new MemoryException(
+ ErrorCodes.ATTEMPTED_TO_ACCESS_NONEXISTENT_COLLECTION);
+ }
+ } catch (SQLException e) {
+ throw new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "\"DELETE FROM\" failed",
+ e);
+ }
+ })
+ .subscribeOn(Schedulers.boundedElastic())
+ .then();
+ }
+
+ public Mono deleteAsync(String collectionName, String key) {
+ return Mono.fromRunnable(
+ () -> {
+ String query =
+ "DELETE FROM "
+ + TABLE_NAME
+ + " WHERE collection = ?"
+ + " AND key = ?";
+ try (PreparedStatement statement =
+ this.connection.prepareStatement(query)) {
+ statement.setString(1, collectionName);
+ statement.setString(2, key != null && !key.isEmpty() ? key : null);
+ statement.executeUpdate();
+ } catch (SQLException e) {
+ throw new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "\"DELETE FROM\" failed",
+ e);
+ }
+ })
+ .subscribeOn(Schedulers.boundedElastic())
+ .then();
+ }
+
+ public Mono deleteBatchAsync(String collectionName, Collection keys) {
+ return Mono.fromRunnable(
+ () -> {
+ String query = batchQuery(BatchOperation.DELETE, keys);
+ try (PreparedStatement statement =
+ this.connection.prepareStatement(query)) {
+ statement.setString(1, collectionName);
+ int index = 2;
+ for (String key : keys) {
+ statement.setString(index++, key);
+ }
+ statement.executeUpdate();
+ } catch (SQLException e) {
+ throw new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "\"DELETE FROM\" failed",
+ e);
+ }
+ })
+ .subscribeOn(Schedulers.boundedElastic())
+ .then();
+ }
+
+ public Mono deleteEmptyAsync(String collectionName) {
+ return Mono.fromRunnable(
+ () -> {
+ String query =
+ "DELETE FROM "
+ + TABLE_NAME
+ + " WHERE collection = ?"
+ + " AND key is NULL";
+ try (PreparedStatement statement =
+ this.connection.prepareStatement(query)) {
+ statement.setString(1, collectionName);
+ statement.executeUpdate();
+ } catch (SQLException e) {
+ throw new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "\"DELETE FROM\" failed",
+ e);
+ }
+ })
+ .subscribeOn(Schedulers.boundedElastic())
+ .then();
+ }
+
+ @Override
+ public void close() {
+ try {
+ this.connection.close();
+ } catch (SQLException e) {
+ throw new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "Database access error while closing connection",
+ e);
+ }
+ }
+}
diff --git a/java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/JDBCMemoryStore.java b/java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/JDBCMemoryStore.java
new file mode 100644
index 0000000000..9c25486d8d
--- /dev/null
+++ b/java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/JDBCMemoryStore.java
@@ -0,0 +1,316 @@
+// Copyright (c) Microsoft. All rights reserved.
+package com.microsoft.semantickernel.connectors.memory.jdbc;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.microsoft.semantickernel.ai.embeddings.Embedding;
+import com.microsoft.semantickernel.memory.MemoryException;
+import com.microsoft.semantickernel.memory.MemoryException.ErrorCodes;
+import com.microsoft.semantickernel.memory.MemoryRecord;
+import java.sql.Connection;
+import java.util.*;
+import java.util.stream.Collectors;
+import javax.annotation.CheckReturnValue;
+import javax.annotation.Nonnull;
+import reactor.core.publisher.Mono;
+import reactor.util.function.Tuple2;
+import reactor.util.function.Tuples;
+
+public class JDBCMemoryStore implements SQLMemoryStore {
+ protected final SQLConnector dbConnector;
+
+ protected JDBCMemoryStore(SQLConnector connector) {
+ this.dbConnector = connector;
+ }
+
+ @Override
+ public Mono createCollectionAsync(@Nonnull String collectionName) {
+ Objects.requireNonNull(collectionName);
+ return this.dbConnector.createCollectionAsync(collectionName);
+ }
+
+ @Override
+ public Mono> getCollectionsAsync() {
+ return this.dbConnector.getCollectionsAsync();
+ }
+
+ @Override
+ public Mono doesCollectionExistAsync(@Nonnull String collectionName) {
+ Objects.requireNonNull(collectionName);
+ return this.dbConnector.doesCollectionExistsAsync(collectionName);
+ }
+
+ @Override
+ public Mono deleteCollectionAsync(@Nonnull String collectionName) {
+ Objects.requireNonNull(collectionName);
+ return this.dbConnector.deleteCollectionAsync(collectionName);
+ }
+
+ @Override
+ public Mono upsertAsync(@Nonnull String collectionName, @Nonnull MemoryRecord record) {
+ Objects.requireNonNull(collectionName);
+ Objects.requireNonNull(record);
+ return doesCollectionExistAsync(collectionName)
+ .handle(
+ (exists, sink) -> {
+ if (!exists) {
+ sink.error(
+ new MemoryException(
+ ErrorCodes
+ .ATTEMPTED_TO_ACCESS_NONEXISTENT_COLLECTION,
+ collectionName));
+ return;
+ }
+ sink.next(exists);
+ })
+ .then(internalUpsertAsync(collectionName, record));
+ }
+
+ protected DatabaseEntry memoryRecordToDatabaseEntry(MemoryRecord record) {
+ try {
+ return new DatabaseEntry(
+ record.getMetadata().getId(),
+ record.getSerializedMetadata(),
+ record.getSerializedEmbedding(),
+ record.getTimestamp());
+ } catch (JsonProcessingException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ protected MemoryRecord databaseEntryToMemoryRecord(DatabaseEntry entry, boolean withEmbedding) {
+ try {
+ Embedding embedding =
+ withEmbedding
+ ? new ObjectMapper().readValue(entry.getEmbedding(), Embedding.class)
+ : Embedding.empty();
+ return MemoryRecord.fromJsonMetadata(
+ entry.getMetadata(), embedding, entry.getKey(), entry.getTimestamp());
+ } catch (JsonProcessingException e) {
+ throw new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "Error deserializing database entry",
+ e);
+ }
+ }
+
+ private Mono internalUpsertAsync(String collectionName, MemoryRecord record) {
+ try {
+ return this.dbConnector.upsertAsync(
+ collectionName,
+ record.getMetadata().getId(),
+ record.getSerializedMetadata(),
+ record.getSerializedEmbedding(),
+ record.getTimestamp());
+ } catch (JsonProcessingException e) {
+ throw new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "Error serializing MemoryRecord",
+ e);
+ }
+ }
+
+ @Override
+ public Mono> upsertBatchAsync(
+ @Nonnull String collectionName, @Nonnull Collection records) {
+ Objects.requireNonNull(collectionName);
+ Objects.requireNonNull(records);
+ return doesCollectionExistAsync(collectionName)
+ .handle(
+ (exists, sink) -> {
+ if (!exists) {
+ sink.error(
+ new MemoryException(
+ ErrorCodes
+ .ATTEMPTED_TO_ACCESS_NONEXISTENT_COLLECTION,
+ collectionName));
+ return;
+ }
+ sink.next(exists);
+ })
+ .then(
+ dbConnector.upsertBatchAsync(
+ collectionName,
+ records.stream()
+ .map(this::memoryRecordToDatabaseEntry)
+ .collect(Collectors.toList())));
+ }
+
+ @Override
+ public Mono getAsync(
+ @Nonnull String collectionName, @Nonnull String key, boolean withEmbedding) {
+ Objects.requireNonNull(collectionName);
+ Objects.requireNonNull(key);
+ return this.internalGetAsync(collectionName, key, withEmbedding);
+ }
+
+ private Mono internalGetAsync(
+ String collectionName, String key, boolean withEmbedding) {
+ Objects.requireNonNull(collectionName);
+ Objects.requireNonNull(key);
+ Mono entry = this.dbConnector.readAsync(collectionName, key);
+
+ return entry.hasElement()
+ .flatMap(
+ hasElement -> {
+ if (!hasElement) {
+ return Mono.empty();
+ }
+
+ return entry.map(
+ databaseEntry ->
+ databaseEntryToMemoryRecord(
+ databaseEntry, withEmbedding));
+ });
+ }
+
+ @Override
+ public Mono> getBatchAsync(
+ @Nonnull String collectionName,
+ @Nonnull Collection keys,
+ boolean withEmbeddings) {
+ Objects.requireNonNull(collectionName);
+ Objects.requireNonNull(keys);
+ return this.dbConnector
+ .readBatchAsync(collectionName, keys)
+ .flatMap(
+ databaseEntries -> {
+ List records = new ArrayList<>();
+ for (DatabaseEntry entry : databaseEntries) {
+ records.add(databaseEntryToMemoryRecord(entry, withEmbeddings));
+ }
+ return Mono.just(records);
+ });
+ }
+
+ @Override
+ public Mono removeAsync(@Nonnull String collectionName, @Nonnull String key) {
+ Objects.requireNonNull(collectionName);
+ Objects.requireNonNull(key);
+ return this.dbConnector.deleteAsync(collectionName, key);
+ }
+
+ @Override
+ public Mono removeBatchAsync(
+ @Nonnull String collectionName, @Nonnull Collection keys) {
+ Objects.requireNonNull(collectionName);
+ Objects.requireNonNull(keys);
+ return this.dbConnector.deleteBatchAsync(collectionName, keys);
+ }
+
+ @Override
+ public Mono>> getNearestMatchesAsync(
+ @Nonnull String collectionName,
+ @Nonnull Embedding embedding,
+ int limit,
+ float minRelevanceScore,
+ boolean withEmbeddings) {
+ Objects.requireNonNull(collectionName);
+ Objects.requireNonNull(embedding);
+ Mono> entries = this.dbConnector.readAllAsync(collectionName);
+
+ return entries.flatMap(
+ databaseEntries -> {
+ List> nearestMatches = new ArrayList<>();
+ for (DatabaseEntry entry : databaseEntries) {
+ if (entry.getEmbedding() == null || entry.getEmbedding().isEmpty()) {
+ continue;
+ }
+ try {
+ Embedding recordEmbedding =
+ new ObjectMapper()
+ .readValue(entry.getEmbedding(), Embedding.class);
+ float similarity = embedding.cosineSimilarity(recordEmbedding);
+ if (similarity >= (float) minRelevanceScore) {
+ MemoryRecord record =
+ MemoryRecord.fromJsonMetadata(
+ entry.getMetadata(),
+ withEmbeddings ? recordEmbedding : null,
+ entry.getKey(),
+ entry.getTimestamp());
+ nearestMatches.add(Tuples.of(record, similarity));
+ }
+ } catch (JsonProcessingException e) {
+ throw new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "Error deserializing database entry",
+ e);
+ }
+ }
+ List> results =
+ nearestMatches.stream()
+ .sorted(
+ Comparator.comparing(
+ Tuple2::getT2, (a, b) -> Float.compare(b, a)))
+ .limit(limit)
+ .collect(Collectors.toList());
+
+ return Mono.just(results);
+ });
+ }
+
+ @Override
+ public Mono> getNearestMatchAsync(
+ @Nonnull String collectionName,
+ @Nonnull Embedding embedding,
+ float minRelevanceScore,
+ boolean withEmbedding) {
+ Objects.requireNonNull(collectionName);
+ Objects.requireNonNull(embedding);
+ return getNearestMatchesAsync(
+ collectionName, embedding, 1, minRelevanceScore, withEmbedding)
+ .flatMap(
+ nearestMatches -> {
+ if (nearestMatches.isEmpty()) {
+ return Mono.empty();
+ }
+ return Mono.just(nearestMatches.iterator().next());
+ });
+ }
+
+ /** Builds a JDBCMemoryStore. */
+ public static class Builder implements SQLMemoryStore.Builder {
+ private Connection connection;
+
+ /**
+ * Builds and returns a JDBCMemoryStore instance with the specified database connection. The
+ * build process will connect to the database and create the required tables.
+ *
+ * @return A JDBCMemoryStore instance configured with the provided database connection.
+ * @deprecated Use {@link #buildAsync()} instead.
+ */
+ @Override
+ @Deprecated
+ public JDBCMemoryStore build() {
+ return this.buildAsync().block();
+ }
+
+ /**
+ * Asynchronously builds and returns a JDBCMemoryStore instance with the specified database
+ * connection.
+ *
+ * @return A Mono with a JDBCMemoryStore instance configured with the provided database
+ * connection.
+ */
+ @Override
+ @CheckReturnValue
+ public Mono buildAsync() {
+ JDBCConnector connector = new JDBCConnector(connection);
+ JDBCMemoryStore memoryStore = new JDBCMemoryStore(connector);
+ return connector.createTableAsync().thenReturn(memoryStore);
+ }
+
+ /**
+ * Sets the database connection to be used by the memory store being built.
+ *
+ * @param connection The Connection object representing the database connection.
+ * @return The updated Builder instance to continue the building process for a
+ * JDBCMemoryStore.
+ */
+ @Override
+ public Builder withConnection(Connection connection) {
+ this.connection = connection;
+ return this;
+ }
+ }
+}
diff --git a/java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/SQLConnector.java b/java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/SQLConnector.java
new file mode 100644
index 0000000000..453ed1fad1
--- /dev/null
+++ b/java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/SQLConnector.java
@@ -0,0 +1,125 @@
+// Copyright (c) Microsoft. All rights reserved.
+package com.microsoft.semantickernel.connectors.memory.jdbc;
+
+import java.time.ZonedDateTime;
+import java.util.Collection;
+import java.util.List;
+import reactor.core.publisher.Mono;
+
+public interface SQLConnector {
+
+ /**
+ * Asynchronously creates a table in the database.
+ *
+ * @return A Mono representing the completion of the table creation operation.
+ */
+ Mono createTableAsync();
+
+ /**
+ * Asynchronously creates a collection within the database.
+ *
+ * @param collectionName The name of the collection to be created.
+ * @return A Mono representing the completion of the collection creation operation.
+ */
+ Mono createCollectionAsync(String collectionName);
+
+ /**
+ * Asynchronously inserts or updates an entry.
+ *
+ * @param collection The name of the collection to insert the entry into.
+ * @param key The key identifying the entry.
+ * @param metadata The metadata associated with the entry.
+ * @param embedding The embedding data associated with the entry.
+ * @param timestamp The timestamp of the insertion.
+ * @return A Mono representing the completion of the insertion operation.
+ */
+ Mono upsertAsync(
+ String collection,
+ String key,
+ String metadata,
+ String embedding,
+ ZonedDateTime timestamp);
+
+ /**
+ * Asynchronously inserts or updates a batch of entries.
+ *
+ * @param records
+ * @return A Mono emitting a list of keys identifying the entries.
+ */
+ Mono> upsertBatchAsync(String collection, Collection records);
+
+ /**
+ * Asynchronously checks if a collection exists.
+ *
+ * @param collectionName The name of the collection to check.
+ * @return A Mono emitting a boolean value indicating whether the collection exists.
+ */
+ Mono doesCollectionExistsAsync(String collectionName);
+
+ /**
+ * Asynchronously retrieves a list of collection names.
+ *
+ * @return A Mono emitting a list of collection names.
+ */
+ Mono> getCollectionsAsync();
+
+ /**
+ * Asynchronously reads all entries within a collection.
+ *
+ * @param collectionName The name of the collection to read from.
+ * @return A Mono emitting a list of DatabaseEntry objects representing the entries.
+ */
+ Mono> readAllAsync(String collectionName);
+
+ /**
+ * Asynchronously reads a specific entry within a collection.
+ *
+ * @param collectionName The name of the collection to read from.
+ * @param key The key identifying the entry to read.
+ * @return A Mono emitting a DatabaseEntry object representing the entry.
+ */
+ Mono readAsync(String collectionName, String key);
+
+ /**
+ * Asynchronously reads a batch of entries within a collection.
+ *
+ * @param collectionName The name of the collection to read from.
+ * @param keys The keys identifying the entries to read.
+ * @return A Mono emitting a list of DatabaseEntry objects representing the entries.
+ */
+ Mono> readBatchAsync(String collectionName, Collection keys);
+
+ /**
+ * Asynchronously deletes a collection.
+ *
+ * @param collectionName The name of the collection to delete.
+ * @return A Mono representing the completion of the collection deletion operation.
+ */
+ Mono deleteCollectionAsync(String collectionName);
+
+ /**
+ * Asynchronously deletes an entry within a collection.
+ *
+ * @param collectionName The name of the collection to delete the entry from.
+ * @param key The key identifying the entry to delete.
+ * @return A Mono representing the completion of the entry deletion operation.
+ */
+ Mono deleteAsync(String collectionName, String key);
+
+ /**
+ * Asynchronously deletes a batch of entries within a collection.
+ *
+ * @param collectionName The name of the collection to delete the entries from.
+ * @param keys The keys identifying the entries to delete.
+ * @return A Mono representing the completion of the entry deletion operation.
+ */
+ Mono deleteBatchAsync(String collectionName, Collection keys);
+
+ /**
+ * Asynchronously deletes all empty entries within a collection.
+ *
+ * @param collectionName The name of the collection to delete empty entries from.
+ * @return A Mono representing the completion of the empty entry deletion operation.
+ */
+ Mono deleteEmptyAsync(String collectionName);
+}
diff --git a/java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/SQLConnectorException.java b/java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/SQLConnectorException.java
new file mode 100644
index 0000000000..45574e1756
--- /dev/null
+++ b/java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/SQLConnectorException.java
@@ -0,0 +1,61 @@
+// Copyright (c) Microsoft. All rights reserved.
+package com.microsoft.semantickernel.connectors.memory.jdbc;
+
+import com.microsoft.semantickernel.SKException;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/** Exception thrown by the SQL connector. */
+public class SQLConnectorException extends SKException {
+
+ private final ErrorCodes errorCode;
+
+ /**
+ * Create an exception with a message
+ *
+ * @param errorCode The error code
+ */
+ public SQLConnectorException(@Nonnull ErrorCodes errorCode) {
+ this(errorCode, null, null);
+ }
+
+ /**
+ * Create an exception with a message
+ *
+ * @param message a description of the cause of the exception
+ */
+ public SQLConnectorException(@Nonnull ErrorCodes errorCode, @Nullable String message) {
+ this(errorCode, message, null);
+ }
+
+ /**
+ * Create an exception with a message and a cause
+ *
+ * @param errorCode the error code
+ * @param message a description of the cause of the exception
+ * @param cause the cause of the exception
+ */
+ public SQLConnectorException(
+ @Nonnull ErrorCodes errorCode, @Nullable String message, @Nullable Throwable cause) {
+ super(message, cause);
+ this.errorCode = errorCode;
+ }
+
+ public ErrorCodes getErrorCode() {
+ return errorCode;
+ }
+
+ public enum ErrorCodes {
+ SQL_ERROR("SQL error");
+
+ final String message;
+
+ ErrorCodes(String message) {
+ this.message = message;
+ }
+
+ public String getMessage() {
+ return message;
+ }
+ }
+}
diff --git a/java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/SQLMemoryStore.java b/java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/SQLMemoryStore.java
new file mode 100644
index 0000000000..ab14f849cb
--- /dev/null
+++ b/java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/SQLMemoryStore.java
@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft. All rights reserved.
+package com.microsoft.semantickernel.connectors.memory.jdbc;
+
+import com.microsoft.semantickernel.memory.MemoryStore;
+import java.sql.Connection;
+import reactor.core.publisher.Mono;
+
+public interface SQLMemoryStore extends MemoryStore {
+ interface Builder extends MemoryStore.Builder {
+ /**
+ * Asynchronously builds the SQLMemoryStore.
+ *
+ * @return A Mono with the SQLMemoryStore.
+ */
+ Mono buildAsync();
+
+ /**
+ * Sets the database connection to be used by the SQL memory store.
+ *
+ * @param connection The Connection object representing the database connection.
+ * @return The updated SQLMemoryStoreBuilder instance.
+ */
+ Builder withConnection(Connection connection);
+ }
+}
diff --git a/java/connectors/semantickernel-connectors-memory-jdbc/src/test/java/com/microsoft/semantickernel/connectors/memory/jdbc/JDBCConnectorTest.java b/java/connectors/semantickernel-connectors-memory-jdbc/src/test/java/com/microsoft/semantickernel/connectors/memory/jdbc/JDBCConnectorTest.java
new file mode 100644
index 0000000000..f1a20dce7a
--- /dev/null
+++ b/java/connectors/semantickernel-connectors-memory-jdbc/src/test/java/com/microsoft/semantickernel/connectors/memory/jdbc/JDBCConnectorTest.java
@@ -0,0 +1,35 @@
+// Copyright (c) Microsoft. All rights reserved.
+package com.microsoft.semantickernel.connectors.memory.jdbc;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import org.junit.jupiter.api.Test;
+
+public class JDBCConnectorTest {
+ @Test
+ void testBatchQuery() {
+ Collection keys = new ArrayList<>();
+ keys.add("key1");
+ keys.add("key2");
+ keys.add("key3");
+
+ JDBCConnector connector = new JDBCConnector(null);
+ String read = connector.batchQuery(JDBCConnector.BatchOperation.SELECT, keys);
+ String delete = connector.batchQuery(JDBCConnector.BatchOperation.DELETE, keys);
+
+ assertNotNull(read);
+ assertNotNull(delete);
+ assertEquals(
+ "SELECT * FROM "
+ + JDBCConnector.TABLE_NAME
+ + " WHERE collection = ? AND key IN (?,?,?)",
+ read);
+ assertEquals(
+ "DELETE FROM "
+ + JDBCConnector.TABLE_NAME
+ + " WHERE collection = ? AND key IN (?,?,?)",
+ delete);
+ }
+}
diff --git a/java/connectors/semantickernel-connectors-memory-mysql/pom.xml b/java/connectors/semantickernel-connectors-memory-mysql/pom.xml
new file mode 100644
index 0000000000..5e77701e59
--- /dev/null
+++ b/java/connectors/semantickernel-connectors-memory-mysql/pom.xml
@@ -0,0 +1,36 @@
+
+
+ 4.0.0
+
+ com.microsoft.semantic-kernel
+ semantickernel-parent
+ 0.2.10-alpha-SNAPSHOT
+ ../../pom.xml
+
+
+ semantickernel-connectors-memory-mysql
+ Semantic Kernel MySQL Memory Connector
+
+
+
+ com.microsoft.semantic-kernel
+ semantickernel-connectors-memory-jdbc
+
+
+ io.projectreactor
+ reactor-core
+
+
+ org.junit.jupiter
+ junit-jupiter-api
+
+
+
+ com.mysql
+ mysql-connector-j
+ 8.0.33
+
+
+
\ No newline at end of file
diff --git a/java/connectors/semantickernel-connectors-memory-mysql/src/main/java/com/microsoft/semantickernel/connectors/memory/mysql/MySQLConnector.java b/java/connectors/semantickernel-connectors-memory-mysql/src/main/java/com/microsoft/semantickernel/connectors/memory/mysql/MySQLConnector.java
new file mode 100644
index 0000000000..c382ee7ac4
--- /dev/null
+++ b/java/connectors/semantickernel-connectors-memory-mysql/src/main/java/com/microsoft/semantickernel/connectors/memory/mysql/MySQLConnector.java
@@ -0,0 +1,370 @@
+// Copyright (c) Microsoft. All rights reserved.
+package com.microsoft.semantickernel.connectors.memory.mysql;
+
+import com.microsoft.semantickernel.connectors.memory.jdbc.DatabaseEntry;
+import com.microsoft.semantickernel.connectors.memory.jdbc.JDBCConnector;
+import com.microsoft.semantickernel.connectors.memory.jdbc.SQLConnector;
+import com.microsoft.semantickernel.connectors.memory.jdbc.SQLConnectorException;
+import java.sql.*;
+import java.time.ZonedDateTime;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Schedulers;
+
+public class MySQLConnector extends JDBCConnector implements SQLConnector {
+ public MySQLConnector(Connection connection) {
+ super(connection);
+ }
+
+ @Override
+ public Mono createTableAsync() {
+ return Mono.fromRunnable(
+ () -> {
+ try (Statement statement = connection.createStatement()) {
+ String createCollectionKeyTable =
+ "CREATE TABLE IF NOT EXISTS "
+ + COLLECTIONS_TABLE_NAME
+ + " (id VARCHAR(255) PRIMARY KEY);";
+
+ String createSKMemoryTable =
+ "CREATE TABLE IF NOT EXISTS "
+ + TABLE_NAME
+ + " ("
+ + "collection VARCHAR(255) NOT NULL, "
+ + "id VARCHAR(255) NOT NULL, "
+ + "metadata TEXT, "
+ + "embedding TEXT, "
+ + "timestamp TEXT, "
+ + "PRIMARY KEY (collection, id), "
+ + "FOREIGN KEY (collection) REFERENCES "
+ + COLLECTIONS_TABLE_NAME
+ + "(id));";
+
+ statement.addBatch(createCollectionKeyTable);
+ statement.addBatch(createSKMemoryTable);
+
+ boolean indexExists = doesIndexExist();
+
+ if (!indexExists) {
+ String createIndex =
+ "CREATE INDEX "
+ + INDEX_NAME
+ + " ON "
+ + TABLE_NAME
+ + "(collection)";
+ statement.addBatch(createIndex);
+ }
+
+ statement.executeBatch();
+ } catch (SQLException e) {
+ throw new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "\"CREATE TABLE\" failed",
+ e);
+ }
+ })
+ .subscribeOn(Schedulers.boundedElastic())
+ .then();
+ }
+
+ private boolean doesIndexExist() throws SQLException {
+ DatabaseMetaData meta = connection.getMetaData();
+ try (ResultSet indexes = meta.getIndexInfo(null, null, TABLE_NAME, false, false)) {
+ while (indexes.next()) {
+ String existingIndexName = indexes.getString("INDEX_NAME");
+ if (existingIndexName != null && existingIndexName.equalsIgnoreCase(INDEX_NAME)) {
+ return true;
+ }
+ }
+ }
+ return false;
+ }
+
+ @Override
+ public Mono createCollectionAsync(String collectionName) {
+ return Mono.fromRunnable(
+ () -> {
+ String query =
+ "INSERT INTO " + COLLECTIONS_TABLE_NAME + " (id) VALUES (?)";
+ try (PreparedStatement statement =
+ this.connection.prepareStatement(query)) {
+ statement.setString(1, collectionName);
+ statement.executeUpdate();
+ } catch (SQLException e) {
+ throw new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "\"INSERT INTO\" failed",
+ e);
+ }
+ })
+ .subscribeOn(Schedulers.boundedElastic())
+ .then();
+ }
+
+ @Override
+ public Mono upsertAsync(
+ String collection,
+ String key,
+ String metadata,
+ String embedding,
+ ZonedDateTime timestamp) {
+ return Mono.fromRunnable(
+ () -> {
+ String query =
+ "INSERT INTO "
+ + TABLE_NAME
+ + " (collection, id, metadata, embedding, timestamp)"
+ + " VALUES (?, ?, ?, ?, ?) ON DUPLICATE KEY UPDATE"
+ + " metadata = VALUES(metadata), embedding ="
+ + " VALUES(embedding), timestamp = VALUES(timestamp)";
+ try (PreparedStatement statement =
+ this.connection.prepareStatement(query)) {
+ String metadataString = metadata != null ? metadata : "";
+ String embeddingString = embedding != null ? embedding : "";
+ String timestampString = formatDatetime(timestamp);
+ statement.setString(1, collection);
+ statement.setString(2, key);
+ statement.setString(3, metadataString);
+ statement.setString(4, embeddingString);
+ statement.setString(5, timestampString);
+ statement.executeUpdate();
+ } catch (SQLException e) {
+ throw new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "\"INSERT INTO\" failed",
+ e);
+ }
+ })
+ .subscribeOn(Schedulers.boundedElastic())
+ .thenReturn(key);
+ }
+
+ @Override
+ public Mono> upsertBatchAsync(
+ String collection, Collection records) {
+ Collection keys = new ArrayList<>();
+ return Mono.fromRunnable(
+ () -> {
+ String query =
+ "INSERT INTO "
+ + TABLE_NAME
+ + " (collection, id, metadata, embedding, timestamp)"
+ + " VALUES (?, ?, ?, ?, ?) ON DUPLICATE KEY UPDATE"
+ + " metadata = VALUES(metadata), embedding ="
+ + " VALUES(embedding), timestamp = VALUES(timestamp)";
+ try (PreparedStatement statement =
+ this.connection.prepareStatement(query)) {
+ for (DatabaseEntry entry : records) {
+ String metadataString =
+ entry.getMetadata() != null ? entry.getMetadata() : "";
+ String embeddingString =
+ entry.getEmbedding() != null
+ ? entry.getEmbedding()
+ : "";
+ String timestampString = formatDatetime(entry.getTimestamp());
+ statement.setString(1, collection);
+ statement.setString(2, entry.getKey());
+ statement.setString(3, metadataString);
+ statement.setString(4, embeddingString);
+ statement.setString(5, timestampString);
+ statement.addBatch();
+ keys.add(entry.getKey());
+ }
+
+ statement.executeBatch();
+ } catch (SQLException e) {
+ throw new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "\"INSERT INTO\" failed",
+ e);
+ }
+ })
+ .subscribeOn(Schedulers.boundedElastic())
+ .thenReturn(keys);
+ }
+
+ @Override
+ public Mono> readAllAsync(String collectionName) {
+ return Mono.defer(
+ () -> {
+ List entries = new ArrayList<>();
+ String query = "SELECT * FROM " + TABLE_NAME + " WHERE collection = ?";
+ try (PreparedStatement statement =
+ this.connection.prepareStatement(query)) {
+ statement.setString(1, collectionName);
+ ResultSet resultSet = statement.executeQuery();
+ while (resultSet.next()) {
+ String key = resultSet.getString("id");
+ String metadata = resultSet.getString("metadata");
+ String embedding = resultSet.getString("embedding");
+ String timestamp = resultSet.getString("timestamp");
+ ZonedDateTime zonedDateTime = parseDatetime(timestamp);
+ entries.add(
+ new DatabaseEntry(
+ key, metadata, embedding, zonedDateTime));
+ }
+ } catch (SQLException e) {
+ return Mono.error(
+ new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "\"SELECT * FROM\" failed",
+ e));
+ }
+ return Mono.just(entries);
+ })
+ .subscribeOn(Schedulers.boundedElastic());
+ }
+
+ @Override
+ public Mono readAsync(String collectionName, String key) {
+ return Mono.defer(
+ () -> {
+ String query =
+ "SELECT * FROM "
+ + TABLE_NAME
+ + " WHERE collection = ?"
+ + " AND id = ?";
+ try (PreparedStatement statement =
+ this.connection.prepareStatement(query)) {
+ statement.setString(1, collectionName);
+ statement.setString(2, key != null && !key.isEmpty() ? key : null);
+ ResultSet resultSet = statement.executeQuery();
+ if (resultSet.next()) {
+ String metadata = resultSet.getString("metadata");
+ String embedding = resultSet.getString("embedding");
+ String timestamp = resultSet.getString("timestamp");
+ ZonedDateTime zonedDateTime = parseDatetime(timestamp);
+ return Mono.just(
+ new DatabaseEntry(
+ key, metadata, embedding, zonedDateTime));
+ }
+ } catch (SQLException e) {
+ return Mono.error(
+ new SQLConnectorException(
+ SQLConnectorException.ErrorCodes.SQL_ERROR,
+ "\"SELECT * FROM\" failed",
+ e));
+ }
+ return Mono.empty();
+ })
+ .subscribeOn(Schedulers.boundedElastic());
+ }
+
+ @Override
+ protected String batchQuery(JDBCConnector.BatchOperation operation, Collection