mirror of
https://github.com/microsoft/semantic-kernel.git
synced 2026-10-12 02:19:06 +00:00
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:
co-authored by
Luigi Montoya
joe-braley
Luigi96
Mark Wallace
John Oliver
David Grieve
parent
d16b3ebf51
commit
d9520ee716
@@ -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>
|
||||
|
||||
+13
-5
@@ -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
|
||||
|
||||
+720
@@ -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());
|
||||
}
|
||||
}
|
||||
+13
-5
@@ -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
|
||||
|
||||
+24
-20
@@ -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
|
||||
+1
-1
@@ -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>
|
||||
+370
@@ -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();
|
||||
}
|
||||
}
|
||||
+61
@@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
+6
-10
@@ -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>
|
||||
@@ -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>
|
||||
|
||||
@@ -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"/>
|
||||
|
||||
Reference in New Issue
Block a user