Java: Add MySQL memory connector (#3099)

### Motivation and Context

<!-- Thank you for your contribution to the semantic-kernel repo!
Please help reviewers and future users, providing the following
information:
  1. Why is this change required?
  2. What problem does it solve?
  3. What scenario does it contribute to?
  4. If it fixes an open issue, please link to the issue here.
-->

Adding memory connector for MySQL with a couple of updates to the JDBC
connector column names to make most methods compatible with MySQL and
other SQL engines.

Also adding one more test to check successful memory store
initialization when existing database objects have already been created.

### Description

<!-- Describe your changes, the overall approach, the underlying design.
These notes will help understanding how your code works. Thanks! -->

### Contribution Checklist

<!-- Before submitting this PR, please make sure: -->

- [ ] The code builds clean without any errors or warnings
- [ ] The PR follows the [SK Contribution
Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md)
and the [pre-submission formatting
script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts)
raises no violations
- [ ] All unit tests pass, and I have added new tests where possible
- [ ] I didn't break anyone 😄

---------

Co-authored-by: Luigi Montoya <[email protected]>
Co-authored-by: joe-braley <[email protected]>
Co-authored-by: Luigi96 <[email protected]>
Co-authored-by: Mark Wallace <[email protected]>
Co-authored-by: John Oliver <[email protected]>
Co-authored-by: David Grieve <[email protected]>
This commit is contained in:
Milder Hernandez
2023-10-19 14:18:23 -04:00
committed by GitHub
co-authored by Luigi Montoya joe-braley Luigi96 Mark Wallace John Oliver David Grieve
parent d16b3ebf51
commit d9520ee716
13 changed files with 1277 additions and 62 deletions
+26 -1
View File
@@ -70,6 +70,13 @@
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>com.microsoft.semantic-kernel</groupId>
<artifactId>semantickernel-connectors-memory-jdbc</artifactId>
<version>${project.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>com.microsoft.semantic-kernel</groupId>
<artifactId>semantickernel-connectors-memory-postgresql</artifactId>
@@ -77,6 +84,20 @@
<scope>test</scope>
</dependency>
<dependency>
<groupId>com.microsoft.semantic-kernel</groupId>
<artifactId>semantickernel-connectors-memory-mysql</artifactId>
<version>${project.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>com.microsoft.semantic-kernel</groupId>
<artifactId>semantickernel-connectors-memory-sqlite</artifactId>
<version>${project.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.xerial</groupId>
<artifactId>sqlite-jdbc</artifactId>
@@ -91,7 +112,11 @@
<artifactId>postgresql</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>mysql</artifactId>
<scope>test</scope>
</dependency>
</dependencies>
<dependencyManagement>
@@ -32,6 +32,7 @@ import reactor.util.function.Tuple2;
public class JDBCMemoryStoreTest {
private static JDBCMemoryStore.Builder builder;
private static MemoryStore db;
private static int collectionNum = 0;
private static final String NULL_ADDITIONAL_METADATA = null;
@@ -40,11 +41,9 @@ public class JDBCMemoryStoreTest {
@BeforeAll
static void setUp() throws SQLException {
db =
new JDBCMemoryStore.Builder()
.withConnection(DriverManager.getConnection("jdbc:sqlite::memory:"))
.buildAsync()
.block();
builder = new JDBCMemoryStore.Builder()
.withConnection(DriverManager.getConnection("jdbc:sqlite::memory:"));
db = builder.buildAsync().block();
}
private Collection<MemoryRecord> createBatchRecords(int numRecords) {
@@ -86,6 +85,15 @@ public class JDBCMemoryStoreTest {
assertNotNull(db);
}
@Test
void itDoesNotFailWithExistingTables() {
// Act
db = builder.buildAsync().block();
// Assert
assertNotNull(db);
}
@Test
void itCanCreateAndGetCollectionAsync() {
// Arrange
@@ -0,0 +1,720 @@
// Copyright (c) Microsoft. All rights reserved.
package com.microsoft.semantickernel.connectors.memory.mysql;
import static org.junit.jupiter.api.Assertions.*;
import com.microsoft.semantickernel.ai.embeddings.Embedding;
import com.microsoft.semantickernel.memory.MemoryException;
import com.microsoft.semantickernel.memory.MemoryRecord;
import com.microsoft.semantickernel.memory.MemoryStore;
import java.io.IOException;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.time.ZoneId;
import java.time.ZonedDateTime;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
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 reactor.core.publisher.Flux;
import reactor.util.function.Tuple2;
@Testcontainers
public class MySQLMemoryStoreTest {
@Container private static final MySQLContainer CONTAINER = new MySQLContainer();
private static final String MYSQL_USER = "test";
private static final String MYSQL_PASSWORD = "test";
private static MySQLMemoryStore.Builder builder;
private static MemoryStore db;
private static int collectionNum = 0;
private static final String NULL_ADDITIONAL_METADATA = null;
private static final String NULL_KEY = null;
private static final ZonedDateTime NULL_TIMESTAMP = null;
@BeforeAll
static void setUp() throws SQLException {
builder = new MySQLMemoryStore.Builder()
.withConnection(
DriverManager.getConnection(
CONTAINER.getJdbcUrl(), MYSQL_USER, MYSQL_PASSWORD));
db = builder.buildAsync().block();
}
private Collection<MemoryRecord> createBatchRecords(int numRecords) {
assertTrue(numRecords % 2 == 0, "Number of records must be even");
assertTrue(numRecords > 0, "Number of records must be greater than 0");
List<MemoryRecord> records = new ArrayList<>(numRecords);
for (int i = 0; i < numRecords / 2; i++) {
MemoryRecord testRecord =
MemoryRecord.localRecord(
"test" + i,
"text" + i,
"description" + i,
new Embedding(Arrays.asList(1f, 1f, 1f)),
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
records.add(testRecord);
}
for (int i = numRecords / 2; i < numRecords; i++) {
MemoryRecord testRecord =
MemoryRecord.referenceRecord(
"test" + i,
"sourceName" + i,
"description" + i,
new Embedding(Arrays.asList(1f, 2f, 3f)),
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
records.add(testRecord);
}
return records;
}
@Test
void initializeDbConnectionSucceeds() {
assertNotNull(db);
}
@Test
void itDoesNotFailWithExistingTables() {
// Act
db = builder.buildAsync().block();
// Assert
assertNotNull(db);
}
@Test
void itCanCreateAndGetCollectionAsync() throws IOException {
// Arrange
String collection = "test_collection" + collectionNum;
collectionNum++;
// Act
db.createCollectionAsync(collection).block();
Collection<String> collections = db.getCollectionsAsync().block();
// Assert
assertNotNull(collections);
assertFalse(collections.isEmpty());
assertTrue(collections.contains(collection));
}
@Test
void itHandlesExceptionsWhenCreatingCollectionAsync() {
// Arrange
String collection = null;
// Assert
assertThrows(
NullPointerException.class,
() -> db.createCollectionAsync(collection).block(),
"Should not be able to create collection with null name");
}
@Test
void itCannotInsertIntoNonExistentCollectionAsync() {
// Arrange
MemoryRecord testRecord =
MemoryRecord.localRecord(
"test",
"text",
"description",
new Embedding(Arrays.asList(1f, 2f, 3f)),
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
String collection = "test_collection" + collectionNum;
collectionNum++;
// Assert
assertThrows(
MemoryException.class,
() -> db.upsertAsync(collection, testRecord).block(),
"Should not be able to insert into a non-existent collection");
}
@Test
void getAsyncReturnsEmptyEmbeddingUnlessSpecifiedAsync() {
// Arrange
MemoryRecord testRecord =
MemoryRecord.localRecord(
"test",
"text",
"description",
new Embedding(Arrays.asList(1f, 2f, 3f)),
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
String collection = "test_collection" + collectionNum;
collectionNum++;
// Act
db.createCollectionAsync(collection).block();
String key = db.upsertAsync(collection, testRecord).block();
MemoryRecord actualDefault = db.getAsync(collection, key, false).block();
MemoryRecord actualWithEmbedding = db.getAsync(collection, key, true).block();
// Assert
assertNotNull(actualDefault);
assertNotNull(actualDefault.getEmbedding());
assertNotNull(actualDefault.getEmbedding().getVector());
assertTrue(actualDefault.getEmbedding().getVector().isEmpty());
assertNotNull(actualWithEmbedding);
assertNotNull(actualWithEmbedding.getEmbedding());
assertNotNull(actualWithEmbedding.getEmbedding().getVector());
assertFalse(actualWithEmbedding.getEmbedding().getVector().isEmpty());
assertNotEquals(testRecord, actualDefault);
assertEquals(testRecord.getMetadata(), actualWithEmbedding.getMetadata());
assertEquals(
testRecord.getEmbedding().getVector(),
actualWithEmbedding.getEmbedding().getVector());
}
@Test
void itCanUpsertAndRetrieveARecordWithNoTimestampAsync() {
// Arrange
MemoryRecord testRecord =
MemoryRecord.localRecord(
"test",
"text",
"description",
new Embedding(Arrays.asList(1f, 2f, 3f)),
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
String collection = "test_collection" + collectionNum;
collectionNum++;
// Act
db.createCollectionAsync(collection).block();
String key = db.upsertAsync(collection, testRecord).block();
MemoryRecord actual = db.getAsync(collection, key, true).block();
// Assert
assertNotNull(actual);
assertEquals(testRecord.getMetadata(), actual.getMetadata());
assertEquals(testRecord.getEmbedding().getVector(), actual.getEmbedding().getVector());
}
@Test
void itCanUpsertAndRetrieveARecordWithTimestampAsync() {
// Arrange
MemoryRecord testRecord =
MemoryRecord.localRecord(
"test",
"text",
"description",
new Embedding(Arrays.asList(1f, 2f, 3f)),
NULL_ADDITIONAL_METADATA,
NULL_KEY,
ZonedDateTime.now(ZoneId.of("UTC")));
String collection = "test_collection" + collectionNum;
collectionNum++;
// Act
db.createCollectionAsync(collection).block();
String key = db.upsertAsync(collection, testRecord).block();
MemoryRecord actual = db.getAsync(collection, key, true).block();
// Assert
assertNotNull(actual);
assertEquals(testRecord.getMetadata(), actual.getMetadata());
assertEquals(testRecord.getEmbedding().getVector(), actual.getEmbedding().getVector());
}
@Test
void upsertReplacesExistingRecordWithSameIdAsync() {
// Arrange
String commonId = "test";
MemoryRecord testRecord =
MemoryRecord.localRecord(
commonId,
"text",
"description",
new Embedding(Arrays.asList(1f, 2f, 3f)),
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
MemoryRecord testRecord2 =
MemoryRecord.localRecord(
commonId,
"text2",
"description2",
new Embedding(Arrays.asList(1f, 2f, 4f)),
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
String collection = "test_collection" + collectionNum;
collectionNum++;
// Act
db.createCollectionAsync(collection).block();
String key = db.upsertAsync(collection, testRecord).block();
String key2 = db.upsertAsync(collection, testRecord2).block();
MemoryRecord actual = db.getAsync(collection, key, true).block();
// Assert
assertNotNull(actual);
assertNotEquals(testRecord2, actual);
assertEquals(key, key2);
assertEquals(testRecord2.getMetadata(), actual.getMetadata());
assertEquals(testRecord2.getEmbedding().getVector(), actual.getEmbedding().getVector());
}
@Test
void existingRecordCanBeRemovedAsync() {
// Arrange
MemoryRecord testRecord =
MemoryRecord.localRecord(
"test",
"text",
"description",
new Embedding(Arrays.asList(1f, 2f, 3f)),
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
String collection = "test_collection" + collectionNum;
collectionNum++;
// Act
db.createCollectionAsync(collection).block();
String key = db.upsertAsync(collection, testRecord).block();
assertNotNull(key);
db.removeAsync(collection, key).block();
MemoryRecord actual = db.getAsync(collection, key, false).block();
// Assert
assertNull(actual);
}
@Test
void removingNonExistingRecordDoesNothingAsync() {
// Arrange
String collection = "test_collection" + collectionNum;
collectionNum++;
// Act
db.removeAsync(collection, "key").block();
MemoryRecord actual = db.getAsync(collection, "key", false).block();
// Assert
assertNull(actual);
}
@Test
void itCanListAllDatabaseCollectionsAsync() {
// Arrange
int numCollections = 3;
String[] testCollections =
IntStream.range(collectionNum, collectionNum += numCollections)
.mapToObj(i -> "test_collection" + i)
.toArray(String[]::new);
Collection<String> collections = db.getCollectionsAsync().block();
assertNotNull(collections);
int initialSize = collections.size();
Flux.fromArray(testCollections)
.concatMap(collection -> db.createCollectionAsync(collection))
.blockLast();
// Act
collections = db.getCollectionsAsync().block();
// Assert
assertNotNull(collections);
assertEquals(initialSize + numCollections, collections.size());
for (String collection : testCollections) {
assertTrue(
collections.contains(collection),
"Collections does not contain the newly-created collection " + collection);
}
}
private String setUpNearestMatches() {
String collection = "test_collection" + collectionNum;
collectionNum++;
db.createCollectionAsync(collection).block();
int i = 0;
MemoryRecord testRecord =
MemoryRecord.localRecord(
"test" + i,
"text" + i,
"description" + i,
new Embedding(Arrays.asList(1f, 1f, 1f)),
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
db.upsertAsync(collection, testRecord).block();
i++;
testRecord =
MemoryRecord.localRecord(
"test" + i,
"text" + i,
"description" + i,
new Embedding(Arrays.asList(-1f, -1f, -1f)),
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
db.upsertAsync(collection, testRecord).block();
i++;
testRecord =
MemoryRecord.localRecord(
"test" + i,
"text" + i,
"description" + i,
new Embedding(Arrays.asList(1f, 2f, 3f)),
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
db.upsertAsync(collection, testRecord).block();
i++;
testRecord =
MemoryRecord.localRecord(
"test" + i,
"text" + i,
"description" + i,
new Embedding(Arrays.asList(-1f, -2f, -3f)),
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
db.upsertAsync(collection, testRecord).block();
i++;
testRecord =
MemoryRecord.localRecord(
"test" + i,
"text" + i,
"description" + i,
new Embedding(Arrays.asList(1f, -1f, -2f)),
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
db.upsertAsync(collection, testRecord).block();
return collection;
}
@Test
void getNearestMatchesReturnsAllResultsWithNoMinScoreAsync() {
// Arrange
Embedding compareEmbedding = new Embedding(Arrays.asList(1f, 1f, 1f));
int topN = 4;
String collection = setUpNearestMatches();
// Act
float threshold = -1f;
Collection<Tuple2<MemoryRecord, Float>> topNResults =
db.getNearestMatchesAsync(collection, compareEmbedding, topN, threshold, false)
.block();
// Assert
assertNotNull(topNResults);
assertEquals(topN, topNResults.size());
Tuple2<MemoryRecord, Float>[] topNResultsArray = topNResults.toArray(new Tuple2[0]);
for (int j = 0; j < topN - 1; j++) {
int compare =
Double.compare(topNResultsArray[j].getT2(), topNResultsArray[j + 1].getT2());
assertTrue(compare >= 0);
}
}
@Test
void getNearestMatchesReturnsLimit() {
// Arrange
Embedding compareEmbedding = new Embedding(Arrays.asList(1f, 1f, 1f));
String collection = setUpNearestMatches();
// Act
float threshold = -1f;
Collection<Tuple2<MemoryRecord, Float>> topNResults =
db.getNearestMatchesAsync(collection, compareEmbedding, 0, threshold, false)
.block();
// Assert
assertNotNull(topNResults);
assertTrue(topNResults.isEmpty());
}
@Test
void getNearestMatchesReturnsEmptyIfLimitZero() {
// Arrange
Embedding compareEmbedding = new Embedding(Arrays.asList(1f, 1f, 1f));
String collection = setUpNearestMatches();
// Act
float threshold = -1f;
Collection<Tuple2<MemoryRecord, Float>> topNResults =
db.getNearestMatchesAsync(collection, compareEmbedding, 0, threshold, false)
.block();
// Assert
assertNotNull(topNResults);
assertTrue(topNResults.isEmpty());
}
@Test
void getNearestMatchesReturnsEmptyIfCollectionEmpty() {
// Arrange
Embedding compareEmbedding = new Embedding(Arrays.asList(1f, 1f, 1f));
String collection = "test_collection" + collectionNum;
collectionNum++;
db.createCollectionAsync(collection).block();
// Act
float threshold = -1f;
Collection<Tuple2<MemoryRecord, Float>> topNResults =
db.getNearestMatchesAsync(
collection, compareEmbedding, Integer.MAX_VALUE, threshold, false)
.block();
// Assert
assertNotNull(topNResults);
assertTrue(topNResults.isEmpty());
}
@Test
void getNearestMatchAsyncReturnsEmptyEmbeddingUnlessSpecifiedAsync() {
// Arrange
Embedding compareEmbedding = new Embedding(Arrays.asList(1f, 1f, 1f));
int topN = 4;
String collection = setUpNearestMatches();
// Act
float threshold = 0.75f;
Tuple2<MemoryRecord, Float> topNResultDefault =
db.getNearestMatchAsync(collection, compareEmbedding, threshold, false).block();
Tuple2<MemoryRecord, Float> topNResultWithEmbedding =
db.getNearestMatchAsync(collection, compareEmbedding, threshold, true).block();
// Assert
assertNotNull(topNResultDefault);
assertNotNull(topNResultWithEmbedding);
assertNotNull(topNResultDefault.getT1().getEmbedding());
assertNotNull(topNResultDefault.getT1().getEmbedding().getVector());
assertTrue(topNResultDefault.getT1().getEmbedding().getVector().isEmpty());
assertNotNull(topNResultWithEmbedding.getT1().getEmbedding());
assertNotNull(topNResultWithEmbedding.getT1().getEmbedding().getVector());
assertFalse(topNResultWithEmbedding.getT1().getEmbedding().getVector().isEmpty());
}
@Test
void getNearestMatchAsyncReturnsExpectedAsync() {
// Arrange
Embedding compareEmbedding = new Embedding(Arrays.asList(1f, 1f, 1f));
int topN = 4;
String collection = setUpNearestMatches();
// Act
float threshold = 0.75f;
Tuple2<MemoryRecord, Float> topNResult =
db.getNearestMatchAsync(collection, compareEmbedding, threshold, false).block();
// Assert
assertNotNull(topNResult);
assertEquals("test0", topNResult.getT1().getMetadata().getId());
assertTrue(topNResult.getT2() >= threshold);
}
@Test
void getNearestMatcheReturnsEmptyIfCollectionEmpty() {
// Arrange
Embedding compareEmbedding = new Embedding(Arrays.asList(1f, 1f, 1f));
int topN = 4;
String collection = "test_collection" + collectionNum;
collectionNum++;
db.createCollectionAsync(collection).block();
// Act
float threshold = -1f;
Tuple2<MemoryRecord, Float> topNResults =
db.getNearestMatchAsync(collection, compareEmbedding, threshold, false).block();
// Assert
assertNull(topNResults);
}
@Test
void getNearestMatchesDifferentiatesIdenticalVectorsByKeyAsync() {
// Arrange
Embedding compareEmbedding = new Embedding(Arrays.asList(1f, 1f, 1f));
int topN = 4;
String collection = "test_collection" + collectionNum;
collectionNum++;
db.createCollectionAsync(collection).block();
for (int i = 0; i < 10; i++) {
MemoryRecord testRecord =
MemoryRecord.localRecord(
"test" + i,
"text" + i,
"description" + i,
new Embedding(Arrays.asList(1f, 1f, 1f)),
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
db.upsertAsync(collection, testRecord).block();
}
// Act
Collection<Tuple2<MemoryRecord, Float>> topNResults =
db.getNearestMatchesAsync(collection, compareEmbedding, topN, 0.75f, true).block();
Collection<String> topNKeys =
topNResults.stream()
.map(tuple -> tuple.getT1().getKey())
.sorted()
.collect(Collectors.toList());
// Assert
assertEquals(topN, topNResults.size());
assertEquals(topN, topNKeys.size());
for (Iterator<Tuple2<MemoryRecord, Float>> iterator = topNResults.iterator();
iterator.hasNext(); ) {
Tuple2<MemoryRecord, Float> tuple = iterator.next();
int compare = Float.compare(tuple.getT2(), 0.75f);
assertTrue(topNKeys.contains(tuple.getT1().getKey()));
assertTrue(compare >= 0);
}
}
@Test
void itCanBatchUpsertRecordsAsync() {
// Arrange
int numRecords = 10;
String collection = "test_collection" + collectionNum;
collectionNum++;
db.createCollectionAsync(collection).block();
Collection<MemoryRecord> records = this.createBatchRecords(numRecords);
// Act
Collection<String> keys = db.upsertBatchAsync(collection, records).block();
Collection<MemoryRecord> resultRecords = db.getBatchAsync(collection, keys, false).block();
// Assert
assertNotNull(keys);
assertEquals(numRecords, keys.size());
assertEquals(numRecords, resultRecords.size());
}
@Test
void itCanBatchGetRecordsAsync() {
// Arrange
int numRecords = 10;
String collection = "test_collection" + collectionNum;
collectionNum++;
db.createCollectionAsync(collection).block();
Collection<MemoryRecord> records = this.createBatchRecords(numRecords);
Collection<String> keys = db.upsertBatchAsync(collection, records).block();
// Act
Collection<MemoryRecord> results = db.getBatchAsync(collection, keys, false).block();
// Assert
assertNotNull(keys);
assertNotNull(results);
assertEquals(numRecords, results.size());
}
@Test
void itCanBatchRemoveRecordsAsync() {
// Arrange
int numRecords = 10;
String collection = "test_collection" + collectionNum;
collectionNum++;
db.createCollectionAsync(collection).block();
Collection<MemoryRecord> records = this.createBatchRecords(numRecords);
List<String> keys = new ArrayList<>();
for (String key : db.upsertBatchAsync(collection, records).block()) {
keys.add(key);
}
// Act
db.removeBatchAsync(collection, keys).block();
// Assert
for (MemoryRecord result : db.getBatchAsync(collection, keys, true).block()) {
assertNull(result);
}
}
@Test
void collectionsCanBeDeletedAsync() {
// Arrange
int numCollections = 3;
String[] testCollections =
IntStream.range(collectionNum, collectionNum += numCollections)
.mapToObj(i -> "test_collection" + i)
.toArray(String[]::new);
Collection<String> collections = db.getCollectionsAsync().block();
assertNotNull(collections);
int initialSize = collections.size();
Flux.fromArray(testCollections)
.concatMap(collection -> db.createCollectionAsync(collection))
.blockLast();
collectionNum += numCollections;
// Act
collections = db.getCollectionsAsync().block();
assertNotNull(collections);
assertEquals(initialSize + numCollections, collections.size());
// Act
for (String collection : collections) {
db.deleteCollectionAsync(collection).block();
}
// Assert
collections = db.getCollectionsAsync().block();
assertNotNull(collections);
assertEquals(0, collections.size());
}
@Test
void itThrowsWhenDeletingNonExistentCollectionAsync() {
// Arrange
String collection = "test_collection" + collectionNum;
collectionNum++;
// Act
assertThrows(MemoryException.class, () -> db.deleteCollectionAsync(collection).block());
}
@Test
void doesCollectionExistAsyncReturnTrueForExistingCollection() {
// Arrange
String collection = "test_collection" + collectionNum;
collectionNum++;
db.createCollectionAsync(collection).block();
// Act
assertTrue(db.doesCollectionExistAsync(collection).block());
}
@Test
void doesCollectionExistAsyncReturnFalseForNonExistentCollection() {
// Arrange
String collection = "test_collection" + collectionNum;
collectionNum++;
// Act
assertFalse(db.doesCollectionExistAsync(collection).block());
}
}
@@ -29,6 +29,7 @@ public class PostgreSQLMemoryStoreTest {
private static final String POSTGRES_USER = "test";
private static final String POSTGRES_PASSWORD = "test";
private static PostgreSQLMemoryStore.Builder builder;
private static MemoryStore db;
private static int collectionNum = 0;
private static final String NULL_ADDITIONAL_METADATA = null;
@@ -37,13 +38,11 @@ public class PostgreSQLMemoryStoreTest {
@BeforeAll
static void setUp() throws SQLException {
db =
new PostgreSQLMemoryStore.Builder()
builder = new PostgreSQLMemoryStore.Builder()
.withConnection(
DriverManager.getConnection(
CONTAINER.getJdbcUrl(), POSTGRES_USER, POSTGRES_PASSWORD))
.buildAsync()
.block();
CONTAINER.getJdbcUrl(), POSTGRES_USER, POSTGRES_PASSWORD));
db = builder.buildAsync().block();
}
private Collection<MemoryRecord> createBatchRecords(int numRecords) {
@@ -85,6 +84,15 @@ public class PostgreSQLMemoryStoreTest {
assertNotNull(db);
}
@Test
void itDoesNotFailWithExistingTables() {
// Act
db = builder.buildAsync().block();
// Assert
assertNotNull(db);
}
@Test
void itCanCreateAndGetCollectionAsync() throws IOException {
// Arrange
@@ -1,35 +1,29 @@
// Copyright (c) Microsoft. All rights reserved.
package com.microsoft.semantickernel.connectors.memory.sqlite;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import com.microsoft.semantickernel.ai.embeddings.Embedding;
import com.microsoft.semantickernel.connectors.memory.jdbc.JDBCMemoryStore;
import com.microsoft.semantickernel.connectors.memory.jdbc.JDBCMemoryStoreTest;
import com.microsoft.semantickernel.memory.MemoryException;
import com.microsoft.semantickernel.memory.MemoryRecord;
import com.microsoft.semantickernel.memory.MemoryStore;
import java.sql.SQLException;
import java.time.ZoneId;
import java.time.ZonedDateTime;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import reactor.core.publisher.Flux;
import reactor.util.function.Tuple2;
public class SQLiteMemoryStoreTest {
import java.sql.SQLException;
import java.time.ZoneId;
import java.time.ZonedDateTime;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import static org.junit.jupiter.api.Assertions.*;
public class SQLiteMemoryStoreTest extends JDBCMemoryStoreTest {
private static SQLiteMemoryStore.Builder builder;
private static MemoryStore db;
private static int collectionNum = 0;
private static final String NULL_ADDITIONAL_METADATA = null;
@@ -38,7 +32,8 @@ public class SQLiteMemoryStoreTest {
@BeforeAll
static void setUp() throws SQLException {
db = new SQLiteMemoryStore.Builder().withFilename(":memory:").buildAsync().block();
builder = new SQLiteMemoryStore.Builder().withFilename(":memory:");
db = builder.buildAsync().block();
}
private Collection<MemoryRecord> createBatchRecords(int numRecords) {
@@ -80,6 +75,15 @@ public class SQLiteMemoryStoreTest {
assertNotNull(db);
}
@Test
void itDoesNotFailWithExistingTables() {
// Act
db = builder.buildAsync().block();
// Assert
assertNotNull(db);
}
@Test
void itCanCreateAndGetCollectionAsync() {
// Arrange
@@ -312,7 +312,7 @@ public class JDBCConnector implements SQLConnector, Closeable {
.subscribeOn(Schedulers.boundedElastic());
}
enum BatchOperation {
protected enum BatchOperation {
SELECT,
DELETE
}
@@ -0,0 +1,36 @@
<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
<modelVersion>4.0.0</modelVersion>
<parent>
<groupId>com.microsoft.semantic-kernel</groupId>
<artifactId>semantickernel-parent</artifactId>
<version>0.2.10-alpha-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>semantickernel-connectors-memory-mysql</artifactId>
<name>Semantic Kernel MySQL Memory Connector</name>
<dependencies>
<dependency>
<groupId>com.microsoft.semantic-kernel</groupId>
<artifactId>semantickernel-connectors-memory-jdbc</artifactId>
</dependency>
<dependency>
<groupId>io.projectreactor</groupId>
<artifactId>reactor-core</artifactId>
</dependency>
<dependency>
<groupId>org.junit.jupiter</groupId>
<artifactId>junit-jupiter-api</artifactId>
</dependency>
<dependency>
<groupId>com.mysql</groupId>
<artifactId>mysql-connector-j</artifactId>
<version>8.0.33</version>
</dependency>
</dependencies>
</project>
@@ -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<Void> 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<Void> 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<String> 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<Collection<String>> upsertBatchAsync(
String collection, Collection<DatabaseEntry> records) {
Collection<String> 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<List<DatabaseEntry>> readAllAsync(String collectionName) {
return Mono.defer(
() -> {
List<DatabaseEntry> 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<DatabaseEntry> 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<String> 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 id 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<Collection<DatabaseEntry>> readBatchAsync(
String collectionName, Collection<String> keys) {
return Mono.defer(
() -> {
Collection<DatabaseEntry> entries = new ArrayList<>();
String query = batchQuery(JDBCConnector.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("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<Void> deleteAsync(String collectionName, String key) {
return Mono.fromRunnable(
() -> {
String query =
"DELETE 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);
statement.executeUpdate();
} catch (SQLException e) {
throw new SQLConnectorException(
SQLConnectorException.ErrorCodes.SQL_ERROR,
"\"DELETE FROM\" failed",
e);
}
})
.subscribeOn(Schedulers.boundedElastic())
.then();
}
@Override
public Mono<Void> deleteEmptyAsync(String collectionName) {
return Mono.fromRunnable(
() -> {
String query =
"DELETE FROM "
+ TABLE_NAME
+ " WHERE collection = ?"
+ " AND id 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();
}
}
@@ -0,0 +1,61 @@
// Copyright (c) Microsoft. All rights reserved.
package com.microsoft.semantickernel.connectors.memory.mysql;
import com.microsoft.semantickernel.connectors.memory.jdbc.JDBCMemoryStore;
import com.microsoft.semantickernel.connectors.memory.jdbc.SQLConnector;
import com.microsoft.semantickernel.connectors.memory.jdbc.SQLMemoryStore;
import java.sql.Connection;
import javax.annotation.CheckReturnValue;
import reactor.core.publisher.Mono;
public class MySQLMemoryStore extends JDBCMemoryStore {
private MySQLMemoryStore(SQLConnector connector) {
super(connector);
}
/** Builds a MySQLMemoryStore. */
public static class Builder implements SQLMemoryStore.Builder<MySQLMemoryStore> {
private Connection connection;
/**
* Builds and returns a MySQLMemoryStore instance with the specified database connection.
* The build process will connect to the database and create the required tables.
*
* @return A MySQLMemoryStore instance configured with the provided database connection.
* @deprecated Use {@link #buildAsync()} instead.
*/
@Override
@Deprecated
public MySQLMemoryStore build() {
return this.buildAsync().block();
}
/**
* Asynchronously builds and returns a MySQLMemoryStore instance with the specified database
* connection.
*
* @return A Mono with a MySQLMemoryStore instance configured with the provided database
* connection.
*/
@Override
@CheckReturnValue
public Mono<MySQLMemoryStore> buildAsync() {
MySQLConnector connector = new MySQLConnector(connection);
MySQLMemoryStore memoryStore = new MySQLMemoryStore(connector);
return connector.createTableAsync().thenReturn(memoryStore);
}
/**
* Sets the MySQL database connection to be used by the MySQL memory store being built.
*
* @param connection The Connection object representing the MySQL database connection.
* @return The updated Builder instance to continue the building process for a
* MySQLMemoryStore.
*/
@Override
public Builder withConnection(Connection connection) {
this.connection = connection;
return this;
}
}
}
@@ -83,8 +83,9 @@ public class PostgreSQLConnector extends JDBCConnector implements SQLConnector {
+ TABLE_NAME
+ " (collection, key, metadata, embedding, timestamp)"
+ " VALUES (?, ?, ?, ?, ?) ON CONFLICT (collection,"
+ " key) DO UPDATE SET metadata = ?, embedding = ?,"
+ " timestamp = ?";
+ " key) DO UPDATE SET metadata = EXCLUDED.metadata,"
+ " embedding = EXCLUDED.embedding, timestamp ="
+ " EXCLUDED.timestamp";
try (PreparedStatement statement =
this.connection.prepareStatement(query)) {
String metadataString = metadata != null ? metadata : "";
@@ -95,9 +96,6 @@ public class PostgreSQLConnector extends JDBCConnector implements SQLConnector {
statement.setString(3, metadataString);
statement.setString(4, embeddingString);
statement.setString(5, timestampString);
statement.setString(6, metadataString);
statement.setString(7, embeddingString);
statement.setString(8, timestampString);
statement.executeUpdate();
} catch (SQLException e) {
throw new SQLConnectorException(
@@ -121,8 +119,9 @@ public class PostgreSQLConnector extends JDBCConnector implements SQLConnector {
+ TABLE_NAME
+ " (collection, key, metadata, embedding, timestamp)"
+ " VALUES (?, ?, ?, ?, ?) ON CONFLICT (collection,"
+ " key) DO UPDATE SET metadata = ?, embedding = ?,"
+ " timestamp = ?";
+ " key) DO UPDATE SET metadata = EXCLUDED.metadata,"
+ " embedding = EXCLUDED.embedding, timestamp ="
+ " EXCLUDED.timestamp";
try (PreparedStatement statement =
this.connection.prepareStatement(query)) {
for (DatabaseEntry entry : records) {
@@ -138,9 +137,6 @@ public class PostgreSQLConnector extends JDBCConnector implements SQLConnector {
statement.setString(3, metadataString);
statement.setString(4, embeddingString);
statement.setString(5, timestampString);
statement.setString(6, metadataString);
statement.setString(7, embeddingString);
statement.setString(8, timestampString);
statement.addBatch();
keys.add(entry.getKey());
}
@@ -1,20 +0,0 @@
<?xml version="1.0" encoding="UTF-8"?>
<Configuration status="info">
<Appenders>
<Console name="Console" target="SYSTEM_OUT">
<PatternLayout pattern="%d{HH:mm:ss} %-5level %c{1.} - %msg%n"/>
</Console>
</Appenders>
<Loggers>
<Logger name="reactor.netty" level="warn">
</Logger>
<Logger name="io.netty" level="warn">
</Logger>
<Logger name="io.quarkus" level="warn">
</Logger>
<Root level="debug">
<AppenderRef ref="Console"/>
</Root>
</Loggers>
</Configuration>
+1
View File
@@ -69,6 +69,7 @@
<module>connectors/semantickernel-connectors-memory-azurecognitivesearch</module>
<module>connectors/semantickernel-connectors-memory-jdbc</module>
<module>connectors/semantickernel-connectors-memory-postgresql</module>
<module>connectors/semantickernel-connectors-memory-mysql</module>
<module>extensions/semantickernel-planners</module>
<module>extensions/semantickernel-gpt3-tokenizer</module>
+6
View File
@@ -100,6 +100,12 @@
<Bug pattern="EI_EXPOSE_REP2"/>
</Match>
<Match>
<Class name="com.microsoft.semantickernel.connectors.memory.mysql.MySQLMemoryStore$Builder"/>
<Method name="withConnection"/>
<Bug pattern="EI_EXPOSE_REP2"/>
</Match>
<Match>
<Class name="com.microsoft.semantickernel.connectors.memory.azurecognitivesearch.AzureCognitiveSearchMemoryRecord"/>
<Bug pattern="EI_EXPOSE_REP"/>