microsoft/semantic-kernel
1
0
Fork 0
mirror of https://github.com/microsoft/semantic-kernel.git synced 2026-10-12 02:19:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Java: Add JDBC and Postgres connectors (#2489)

Browse Source 
### 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 JDBC connector from SQLite and Postgres connector overriding
what's necessary.

Also adding new integration tests for the JDBC and Postgres memory
connectors under
`api-test/integration-tests/src/test/java/com/microsoft/semantickernel/connectors`.

### 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-09-07 18:37:14 +01:00
committed by GitHub
co-authored by Luigi Montoya joe-braley Luigi96 Mark Wallace John Oliver David Grieve
parent beeed7b7a7
commit 7a9268d29e
25 changed files with 2611 additions and 828 deletions
Download Patch File Download Diff File Expand all files Collapse all files
java/api-test/integration-tests/pom.xml
+35
View File
@@ -69,5 +69,40 @@
<artifactId>sample-code</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>com.microsoft.semantic-kernel</groupId>
<artifactId>semantickernel-connectors-memory-postgresql</artifactId>
<version>${project.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.xerial</groupId>
<artifactId>sqlite-jdbc</artifactId>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>junit-jupiter</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>postgresql</artifactId>
<scope>test</scope>
</dependency>
</dependencies>
<dependencyManagement>
<dependencies>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>testcontainers-bom</artifactId>
<version>1.18.3</version>
<type>pom</type>
<scope>import</scope>
</dependency>
</dependencies>
</dependencyManagement>
</project>
java/api-test/integration-tests/src/test/java/com/microsoft/semantickernel/connectors/memory/jdbc/JDBCMemoryStoreTest.java
+713
View File
@@ -0,0 +1,713 @@
// Copyright (c) Microsoft. All rights reserved.
package com.microsoft.semantickernel.connectors.memory.jdbc;
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.memory.MemoryException;
import com.microsoft.semantickernel.memory.MemoryRecord;
import com.microsoft.semantickernel.memory.MemoryStore;
import java.awt.*;
import java.sql.DriverManager;
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 JDBCMemoryStoreTest {
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 {
db =
new JDBCMemoryStore.Builder()
.withConnection(DriverManager.getConnection("jdbc:sqlite::memory:"))
.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 itCanCreateAndGetCollectionAsync() {
// 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 = this.db.getCollectionsAsync().block();
assertNotNull(collections);
int initialSize = collections.size();
Flux.fromArray(testCollections)
.concatMap(collection -> this.db.createCollectionAsync(collection))
.blockLast();
// Act
collections = this.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
double threshold = -1;
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
double threshold = -1;
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
double threshold = -1;
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
double threshold = -1;
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
double threshold = 0.75;
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
double threshold = 0.75;
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
double threshold = -1;
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.75, 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 = this.db.getCollectionsAsync().block();
assertNotNull(collections);
int initialSize = collections.size();
Flux.fromArray(testCollections)
.concatMap(collection -> this.db.createCollectionAsync(collection))
.blockLast();
collectionNum += numCollections;
// Act
collections = this.db.getCollectionsAsync().block();
assertNotNull(collections);
assertEquals(initialSize + numCollections, collections.size());
// Act
for (String collection : collections) {
this.db.deleteCollectionAsync(collection).block();
}
// Assert
collections = this.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());
}
}
java/api-test/integration-tests/src/test/java/com/microsoft/semantickernel/connectors/memory/postgresql/PostgreSQLMemoryStoreTest.java
+712
View File
@@ -0,0 +1,712 @@
// Copyright (c) Microsoft. All rights reserved.
package com.microsoft.semantickernel.connectors.memory.postgresql;
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.PostgreSQLContainer;
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 PostgreSQLMemoryStoreTest {
@Container private static final PostgreSQLContainer CONTAINER = new PostgreSQLContainer();
private static final String POSTGRES_USER = "test";
private static final String POSTGRES_PASSWORD = "test";
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 {
db =
new PostgreSQLMemoryStore.Builder()
.withConnection(
DriverManager.getConnection(
CONTAINER.getJdbcUrl(), POSTGRES_USER, POSTGRES_PASSWORD))
.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 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
double threshold = -1;
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
double threshold = -1;
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
double threshold = -1;
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
double threshold = -1;
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
double threshold = 0.75;
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
double threshold = 0.75;
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
double threshold = -1;
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.75, 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());
}
}
java/connectors/semantickernel-connectors-memory-azurecognitivesearch/src/main/java/com/microsoft/semantickernel/connectors/memory/azurecognitivesearch/AzureCognitiveSearchMemory.java
+30 -30
View File
@@ -95,12 +95,12 @@ public class AzureCognitiveSearchMemory implements SemanticTextMemory {
/**
* @throws IllegalArgumentException if the collection name does not conform to Azure Cognitive
* Search
* <a href="https://learn.microsoft.com/en-us/rest/api/searchservice/naming-rules#chart-of-naming-rules">Index naming rules</a>.
* The name must be lower-case, start with a letter or number,
* have no slashes or dots, and be 2 to 128 characters in length. After starting the name
* with a letter or number, the rest of the name can include any letter, number and dashes,
* as long as the dashes aren't consecutive.
* Search <a
* href="https://learn.microsoft.com/en-us/rest/api/searchservice/naming-rules#chart-of-naming-rules">Index
* naming rules</a>. The name must be lower-case, start with a letter or number, have no
* slashes or dots, and be 2 to 128 characters in length. After starting the name with a
* letter or number, the rest of the name can include any letter, number and dashes, as long
* as the dashes aren't consecutive.
*/
@Override
public Mono<String> saveInformationAsync(
@@ -119,12 +119,12 @@ public class AzureCognitiveSearchMemory implements SemanticTextMemory {
/**
* @throws IllegalArgumentException if the collection name does not conform to Azure Cognitive
* Search
* <a href="https://learn.microsoft.com/en-us/rest/api/searchservice/naming-rules#chart-of-naming-rules">Index naming rules</a>.
* The name must be lower-case, start with a letter or number,
* have no slashes or dots, and be 2 to 128 characters in length. After starting the name
* with a letter or number, the rest of the name can include any letter, number and dashes,
* as long as the dashes aren't consecutive.
* Search <a
* href="https://learn.microsoft.com/en-us/rest/api/searchservice/naming-rules#chart-of-naming-rules">Index
* naming rules</a>. The name must be lower-case, start with a letter or number, have no
* slashes or dots, and be 2 to 128 characters in length. After starting the name with a
* letter or number, the rest of the name can include any letter, number and dashes, as long
* as the dashes aren't consecutive.
*/
@Override
public Mono<String> saveReferenceAsync(
@@ -149,12 +149,12 @@ public class AzureCognitiveSearchMemory implements SemanticTextMemory {
/**
* @throws IllegalArgumentException if the collection name does not conform to Azure Cognitive
* Search
* <a href="https://learn.microsoft.com/en-us/rest/api/searchservice/naming-rules#chart-of-naming-rules">Index naming rules</a>.
* The name must be lower-case, start with a letter or number,
* have no slashes or dots, and be 2 to 128 characters in length. After starting the name
* with a letter or number, the rest of the name can include any letter, number and dashes,
* as long as the dashes aren't consecutive.
* Search <a
* href="https://learn.microsoft.com/en-us/rest/api/searchservice/naming-rules#chart-of-naming-rules">Index
* naming rules</a>. The name must be lower-case, start with a letter or number, have no
* slashes or dots, and be 2 to 128 characters in length. After starting the name with a
* letter or number, the rest of the name can include any letter, number and dashes, as long
* as the dashes aren't consecutive.
*/
@Override
public Mono<MemoryQueryResult> getAsync(
@@ -176,12 +176,12 @@ public class AzureCognitiveSearchMemory implements SemanticTextMemory {
/**
* @throws IllegalArgumentException if the collection name does not conform to Azure Cognitive
* Search
* <a href="https://learn.microsoft.com/en-us/rest/api/searchservice/naming-rules#chart-of-naming-rules">Index naming rules</a>.
* The name must be lower-case, start with a letter or number,
* have no slashes or dots, and be 2 to 128 characters in length. After starting the name
* with a letter or number, the rest of the name can include any letter, number and dashes,
* as long as the dashes aren't consecutive.
* Search <a
* href="https://learn.microsoft.com/en-us/rest/api/searchservice/naming-rules#chart-of-naming-rules">Index
* naming rules</a>. The name must be lower-case, start with a letter or number, have no
* slashes or dots, and be 2 to 128 characters in length. After starting the name with a
* letter or number, the rest of the name can include any letter, number and dashes, as long
* as the dashes aren't consecutive.
*/
@Override
public Mono<Void> removeAsync(@Nonnull String collection, @Nonnull String key) {
@@ -221,12 +221,12 @@ public class AzureCognitiveSearchMemory implements SemanticTextMemory {
/**
* @throws IllegalArgumentException if the collection name does not conform to Azure Cognitive
* Search
* <a href="https://learn.microsoft.com/en-us/rest/api/searchservice/naming-rules#chart-of-naming-rules">Index naming rules</a>.
* The name must be lower-case, start with a letter or number,
* have no slashes or dots, and be 2 to 128 characters in length. After starting the name
* with a letter or number, the rest of the name can include any letter, number and dashes,
* as long as the dashes aren't consecutive.
* Search <a
* href="https://learn.microsoft.com/en-us/rest/api/searchservice/naming-rules#chart-of-naming-rules">Index
* naming rules</a>. The name must be lower-case, start with a letter or number, have no
* slashes or dots, and be 2 to 128 characters in length. After starting the name with a
* letter or number, the rest of the name can include any letter, number and dashes, as long
* as the dashes aren't consecutive.
*/
@Override
public Mono<List<MemoryQueryResult>> searchAsync(
java/connectors/semantickernel-connectors-memory-jdbc/pom.xml
+57
View File
@@ -0,0 +1,57 @@
<?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.9-alpha-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>semantickernel-connectors-memory-jdbc</artifactId>
<dependencies>
<dependency>
<groupId>com.fasterxml.jackson.core</groupId>
<artifactId>jackson-core</artifactId>
<scope>compile</scope>
</dependency>
<dependency>
<groupId>com.fasterxml.jackson.core</groupId>
<artifactId>jackson-databind</artifactId>
<scope>compile</scope>
</dependency>
<dependency>
<groupId>com.microsoft.semantic-kernel</groupId>
<artifactId>semantickernel-api</artifactId>
</dependency>
<dependency>
<groupId>org.mockito</groupId>
<artifactId>mockito-core</artifactId>
</dependency>
<dependency>
<groupId>org.junit.jupiter</groupId>
<artifactId>junit-jupiter</artifactId>
</dependency>
<dependency>
<groupId>org.apache.logging.log4j</groupId>
<artifactId>log4j-api</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.apache.logging.log4j</groupId>
<artifactId>log4j-core</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>com.microsoft.semantic-kernel</groupId>
<artifactId>semantickernel-api</artifactId>
</dependency>
</dependencies>
</project>
java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/DatabaseEntry.java
+43
View File
@@ -0,0 +1,43 @@
// Copyright (c) Microsoft. All rights reserved.
package com.microsoft.semantickernel.connectors.memory.jdbc;
import com.microsoft.semantickernel.memory.DataEntryBase;
import java.time.ZonedDateTime;
/** Represents an entry in the Semantic Kernel Memory Table. */
public class DatabaseEntry extends DataEntryBase {
private final String embedding; // JSON with the embedding information
private final String metadata; // JSON with the metadata associated
/**
* Creates a new instance of DatabaseEntry.
*
* @param key The key identifying the entry.
* @param metadata The metadata associated.
* @param embedding The embedding information associated.
* @param timestamp The timestamp indicating when the entry was created or modified.
*/
public DatabaseEntry(String key, String metadata, String embedding, ZonedDateTime timestamp) {
super(key, timestamp);
this.metadata = metadata;
this.embedding = embedding;
}
/**
* Gets the metadata associated with the entry.
*
* @return The metadata associated with the entry.
*/
public String getMetadata() {
return metadata;
}
/**
* Gets the embedding information associated with the entry.
*
* @return The embedding information associated with the entry.
*/
public String getEmbedding() {
return embedding;
}
}
java/connectors/semantickernel-connectors-memory-sqlite/src/main/java/com/microsoft/semantickernel/connectors/memory/sqlite/Database.java → java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/JDBCConnector.java
+97 -104
View File
@@ -1,9 +1,9 @@
// Copyright (c) Microsoft. All rights reserved.
package com.microsoft.semantickernel.connectors.memory.sqlite;
package com.microsoft.semantickernel.connectors.memory.jdbc;
import com.microsoft.semantickernel.memory.DataEntryBase;
import com.microsoft.semantickernel.memory.MemoryException;
import com.microsoft.semantickernel.memory.MemoryException.ErrorCodes;
import java.io.Closeable;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
@@ -13,70 +13,79 @@ import java.time.ZonedDateTime;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.List;
import org.sqlite.SQLiteErrorCode;
import org.sqlite.SQLiteException;
import javax.annotation.Nullable;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;
public class Database {
public class JDBCConnector implements SQLConnector, Closeable {
protected final Connection connection;
static class DatabaseEntry extends DataEntryBase {
private final String metadata;
private final String embedding;
public DatabaseEntry(
String key, String metadata, String embedding, ZonedDateTime timestamp) {
super(key, timestamp);
this.metadata = metadata;
this.embedding = embedding;
}
public String getMetadata() {
return metadata;
}
public String getEmbedding() {
return embedding;
}
public JDBCConnector(Connection connection) {
this.connection = connection;
}
// Convenience method to format a ZonedDateTime in a format acceptable to SQLite
private static String formatDatetime(ZonedDateTime datetime) {
// Convenience method to format a ZonedDateTime in a format acceptable to SQL
protected static String formatDatetime(@Nullable ZonedDateTime datetime) {
if (datetime == null) return "";
return datetime.format(DateTimeFormatter.ISO_OFFSET_DATE_TIME);
}
// Convenience method to parse a SQLite datetime string into a ZonedDateTime
private static ZonedDateTime parseDatetime(String datetime) {
// Convenience method to parse a SQL datetime string into a ZonedDateTime
protected static ZonedDateTime parseDatetime(@Nullable String datetime) {
if (datetime == null || datetime.isEmpty()) return null;
return ZonedDateTime.parse(datetime, DateTimeFormatter.ISO_OFFSET_DATE_TIME);
}
private static final String COLLECTIONS_TABLE_NAME = "SKCollectionTable";
private static final String TABLE_NAME = "SKMemoryTable";
private static final String INDEX_NAME = "SKMemoryIndex";
protected static final String COLLECTIONS_TABLE_NAME = "SKCollectionTable";
protected static final String TABLE_NAME = "SKMemoryTable";
protected static final String INDEX_NAME = "SKMemoryIndex";
public Mono<Void> createTableAsync(Connection connection) {
/**
* Returns the name of the table that stores the collection names.
*
* @return The name of the table that stores the collection names.
*/
protected static String DEFAULT_COLLECTIONS_TABLE_NAME() {
return COLLECTIONS_TABLE_NAME;
}
/**
* Returns the name of the Semantic Kernel table.
*
* @return The name of the Semantic Kernel table.
*/
protected static String DEFAULT_TABLE_NAME() {
return TABLE_NAME;
}
/**
* Returns the name of the index on the Semantic Kernel table.
*
* @return The name of the index on the Semantic Kernel table.
*/
protected static String DEFAULT_INDEX_NAME() {
return INDEX_NAME;
}
public Mono<Void> createTableAsync() {
return Mono.fromRunnable(
() -> {
String createCollectionKeyTable =
"CREATE TABLE IF NOT EXISTS "
+ COLLECTIONS_TABLE_NAME
+ " ("
+ "id INTEGER PRIMARY KEY, "
+ "name TEXT NOT NULL UNIQUE"
+ "id TEXT PRIMARY KEY"
+ " )";
String createSKMemoryTable =
"CREATE TABLE IF NOT EXISTS "
+ TABLE_NAME
+ " ("
+ "collection INTEGER, "
+ "collection TEXT, "
+ "key TEXT, "
+ "metadata TEXT, "
+ "embedding TEXT, "
+ "timestamp TEXT, "
+ "FOREIGN KEY(collection) REFERENCES "
+ "PRIMARY KEY (collection, key), "
+ "FOREIGN KEY (collection) REFERENCES "
+ COLLECTIONS_TABLE_NAME
+ "(id)"
+ " )";
@@ -88,7 +97,7 @@ public class Database {
+ TABLE_NAME
+ "(collection)";
try (Statement statement = connection.createStatement()) {
try (Statement statement = this.connection.createStatement()) {
statement.addBatch(createCollectionKeyTable);
statement.addBatch(createSKMemoryTable);
statement.addBatch(createIndex);
@@ -104,21 +113,15 @@ public class Database {
.then();
}
public Mono<Void> createCollectionAsync(Connection connection, String collectionName) {
public Mono<Void> createCollectionAsync(String collectionName) {
return Mono.fromRunnable(
() -> {
String query =
"INSERT INTO " + COLLECTIONS_TABLE_NAME + " (name) VALUES (?)";
try (PreparedStatement statement = connection.prepareStatement(query)) {
"INSERT INTO " + COLLECTIONS_TABLE_NAME + " (id) VALUES (?)";
try (PreparedStatement statement =
this.connection.prepareStatement(query)) {
statement.setString(1, collectionName);
statement.executeUpdate();
} catch (SQLiteException e) {
if (e.getResultCode() != SQLiteErrorCode.SQLITE_CONSTRAINT_UNIQUE) {
throw new SQLConnectorException(
SQLConnectorException.ErrorCodes.SQL_ERROR,
"\"INSERT INTO\" failed",
e);
}
} catch (SQLException e) {
throw new SQLConnectorException(
SQLConnectorException.ErrorCodes.SQL_ERROR,
@@ -131,7 +134,6 @@ public class Database {
}
public Mono<Void> updateAsync(
Connection connection,
String collection,
String key,
String metadata,
@@ -144,11 +146,10 @@ public class Database {
"UPDATE "
+ TABLE_NAME
+ " SET metadata = ?, embedding = ?, timestamp = ?"
+ " WHERE collection = (SELECT id FROM "
+ COLLECTIONS_TABLE_NAME
+ " WHERE name = ?)"
+ " WHERE collection = ?"
+ " AND key = ?";
try (PreparedStatement statement = connection.prepareStatement(query)) {
try (PreparedStatement statement =
this.connection.prepareStatement(query)) {
statement.setString(1, metadata != null ? metadata : "");
statement.setString(2, embedding != null ? embedding : "");
statement.setString(3, formatDatetime(timestamp));
@@ -167,7 +168,6 @@ public class Database {
}
public Mono<Void> insertOrIgnoreAsync(
Connection connection,
String collection,
String key,
String metadata,
@@ -179,10 +179,9 @@ public class Database {
"INSERT OR IGNORE INTO "
+ TABLE_NAME
+ " (collection, key, metadata, embedding, timestamp)"
+ " VALUES ((SELECT id FROM "
+ COLLECTIONS_TABLE_NAME
+ " WHERE name = ?), ?, ?, ?, ?)";
try (PreparedStatement statement = connection.prepareStatement(query)) {
+ " VALUES (?, ?, ?, ?, ?)";
try (PreparedStatement statement =
this.connection.prepareStatement(query)) {
statement.setString(1, collection);
statement.setString(2, key);
statement.setString(3, metadata != null ? metadata : "");
@@ -200,30 +199,30 @@ public class Database {
.then();
}
public Mono<Boolean> doesCollectionExistsAsync(Connection connection, String collectionName) {
return Mono.fromCallable(() -> doesCollectionExists(connection, collectionName))
public Mono<Boolean> doesCollectionExistsAsync(String collectionName) {
return Mono.fromCallable(() -> doesCollectionExists(collectionName))
.subscribeOn(Schedulers.boundedElastic());
}
private boolean doesCollectionExists(Connection connection, String collectionName)
throws SQLException {
String query = "SELECT name FROM " + COLLECTIONS_TABLE_NAME + " WHERE name = ?";
try (PreparedStatement statement = connection.prepareStatement(query)) {
private boolean doesCollectionExists(String collectionName) throws SQLException {
String query = "SELECT id FROM " + COLLECTIONS_TABLE_NAME + " WHERE id = ?";
try (PreparedStatement statement = this.connection.prepareStatement(query)) {
statement.setString(1, collectionName);
ResultSet resultSet = statement.executeQuery();
return resultSet.next();
}
}
public Mono<List<String>> getCollectionsAsync(Connection connection) {
public Mono<List<String>> getCollectionsAsync() {
return Mono.defer(
() -> {
List<String> collections = new ArrayList<>();
String query = "SELECT name FROM " + COLLECTIONS_TABLE_NAME;
try (PreparedStatement statement = connection.prepareStatement(query)) {
String query = "SELECT id FROM " + COLLECTIONS_TABLE_NAME;
try (PreparedStatement statement =
this.connection.prepareStatement(query)) {
ResultSet resultSet = statement.executeQuery();
while (resultSet.next()) {
String collection = resultSet.getString("name");
String collection = resultSet.getString("id");
collections.add(collection);
}
} catch (SQLException e) {
@@ -238,18 +237,13 @@ public class Database {
.subscribeOn(Schedulers.boundedElastic());
}
public Mono<List<DatabaseEntry>> readAllAsync(Connection connection, String collectionName) {
public Mono<List<DatabaseEntry>> readAllAsync(String collectionName) {
return Mono.defer(
() -> {
List<DatabaseEntry> entries = new ArrayList<>();
String query =
"SELECT * FROM "
+ TABLE_NAME
+ " WHERE collection ="
+ " (SELECT id FROM "
+ COLLECTIONS_TABLE_NAME
+ " WHERE name = ?)";
try (PreparedStatement statement = connection.prepareStatement(query)) {
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()) {
@@ -274,18 +268,16 @@ public class Database {
.subscribeOn(Schedulers.boundedElastic());
}
public Mono<DatabaseEntry> readAsync(Connection connection, String collectionName, String key) {
public Mono<DatabaseEntry> readAsync(String collectionName, String key) {
return Mono.defer(
() -> {
String query =
"SELECT * FROM "
+ TABLE_NAME
+ " WHERE collection ="
+ " (SELECT id FROM "
+ COLLECTIONS_TABLE_NAME
+ " WHERE name = ?)"
+ " WHERE collection = ?"
+ " AND key = ?";
try (PreparedStatement statement = connection.prepareStatement(query)) {
try (PreparedStatement statement =
this.connection.prepareStatement(query)) {
statement.setString(1, collectionName);
statement.setString(2, key != null && !key.isEmpty() ? key : null);
ResultSet resultSet = statement.executeQuery();
@@ -310,21 +302,16 @@ public class Database {
.subscribeOn(Schedulers.boundedElastic());
}
public Mono<Void> deleteCollectionAsync(Connection connection, String collectionName) {
public Mono<Void> deleteCollectionAsync(String collectionName) {
return Mono.fromRunnable(
() -> {
String query1 =
"DELETE FROM "
+ TABLE_NAME
+ " WHERE collection ="
+ " (SELECT id FROM "
+ COLLECTIONS_TABLE_NAME
+ " WHERE name = ?)";
String query1 = "DELETE FROM " + TABLE_NAME + " WHERE collection = ?";
String query2 =
"DELETE FROM " + COLLECTIONS_TABLE_NAME + " WHERE name = ?";
try (PreparedStatement statement = connection.prepareStatement(query1);
"DELETE FROM " + COLLECTIONS_TABLE_NAME + " WHERE id = ?";
try (PreparedStatement statement =
this.connection.prepareStatement(query1);
PreparedStatement statement2 =
connection.prepareStatement(query2)) {
this.connection.prepareStatement(query2)) {
statement.setString(1, collectionName);
statement.executeUpdate();
statement2.setString(1, collectionName);
@@ -343,18 +330,16 @@ public class Database {
.then();
}
public Mono<Void> deleteAsync(Connection connection, String collectionName, String key) {
public Mono<Void> deleteAsync(String collectionName, String key) {
return Mono.fromRunnable(
() -> {
String query =
"DELETE FROM "
+ TABLE_NAME
+ " WHERE collection ="
+ " (SELECT id FROM "
+ COLLECTIONS_TABLE_NAME
+ " WHERE name = ?)"
+ " WHERE collection = ?"
+ " AND key = ?";
try (PreparedStatement statement = connection.prepareStatement(query)) {
try (PreparedStatement statement =
this.connection.prepareStatement(query)) {
statement.setString(1, collectionName);
statement.setString(2, key != null && !key.isEmpty() ? key : null);
statement.executeUpdate();
@@ -369,18 +354,16 @@ public class Database {
.then();
}
public Mono<Void> deleteEmptyAsync(Connection connection, String collectionName) {
public Mono<Void> deleteEmptyAsync(String collectionName) {
return Mono.fromRunnable(
() -> {
String query =
"DELETE FROM "
+ TABLE_NAME
+ " WHERE collection ="
+ " (SELECT id FROM "
+ COLLECTIONS_TABLE_NAME
+ " WHERE name = ?)"
+ " WHERE collection = ?"
+ " AND key is NULL";
try (PreparedStatement statement = connection.prepareStatement(query)) {
try (PreparedStatement statement =
this.connection.prepareStatement(query)) {
statement.setString(1, collectionName);
statement.executeUpdate();
} catch (SQLException e) {
@@ -393,4 +376,14 @@ public class Database {
.subscribeOn(Schedulers.boundedElastic())
.then();
}
@Override
public void close() {
try {
this.connection.close();
} catch (SQLException e) {
throw new SQLConnectorException(
SQLConnectorException.ErrorCodes.SQL_ERROR, "Database access error while closing connection", e);
}
}
}
java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/JDBCMemoryStore.java
+309
View File
@@ -0,0 +1,309 @@
// Copyright (c) Microsoft. All rights reserved.
package com.microsoft.semantickernel.connectors.memory.jdbc;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.microsoft.semantickernel.ai.embeddings.Embedding;
import com.microsoft.semantickernel.memory.MemoryException;
import com.microsoft.semantickernel.memory.MemoryException.ErrorCodes;
import com.microsoft.semantickernel.memory.MemoryRecord;
import java.sql.Connection;
import java.util.*;
import java.util.stream.Collectors;
import javax.annotation.CheckReturnValue;
import javax.annotation.Nonnull;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.util.function.Tuple2;
import reactor.util.function.Tuples;
public class JDBCMemoryStore implements SQLMemoryStore {
protected final SQLConnector dbConnector;
protected JDBCMemoryStore(SQLConnector connector) {
this.dbConnector = connector;
}
@Override
public Mono<Void> createCollectionAsync(@Nonnull String collectionName) {
Objects.requireNonNull(collectionName);
return this.dbConnector.createCollectionAsync(collectionName);
}
@Override
public Mono<List<String>> getCollectionsAsync() {
return this.dbConnector.getCollectionsAsync();
}
@Override
public Mono<Boolean> doesCollectionExistAsync(@Nonnull String collectionName) {
Objects.requireNonNull(collectionName);
return this.dbConnector.doesCollectionExistsAsync(collectionName);
}
@Override
public Mono<Void> deleteCollectionAsync(@Nonnull String collectionName) {
Objects.requireNonNull(collectionName);
return this.dbConnector.deleteCollectionAsync(collectionName);
}
@Override
public Mono<String> upsertAsync(@Nonnull String collectionName, @Nonnull MemoryRecord record) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(record);
return doesCollectionExistAsync(collectionName)
.handle(
(exists, sink) -> {
if (!exists) {
sink.error(
new MemoryException(
ErrorCodes
.ATTEMPTED_TO_ACCESS_NONEXISTENT_COLLECTION,
collectionName));
return;
}
sink.next(exists);
})
.then(internalUpsertAsync(collectionName, record));
}
private Mono<String> internalUpsertAsync(String collectionName, MemoryRecord record) {
try {
Mono<Void> update =
this.dbConnector.updateAsync(
collectionName,
record.getMetadata().getId(),
record.getSerializedMetadata(),
record.getSerializedEmbedding(),
record.getTimestamp());
Mono<Void> insert =
this.dbConnector.insertOrIgnoreAsync(
collectionName,
record.getMetadata().getId(),
record.getSerializedMetadata(),
record.getSerializedEmbedding(),
record.getTimestamp());
return update.then(insert).then(Mono.just(record.getMetadata().getId()));
} catch (JsonProcessingException e) {
throw new SQLConnectorException(
SQLConnectorException.ErrorCodes.SQL_ERROR,
"Error serializing MemoryRecord",
e);
}
}
@Override
public Mono<Collection<String>> upsertBatchAsync(
@Nonnull String collectionName, @Nonnull Collection<MemoryRecord> records) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(records);
return doesCollectionExistAsync(collectionName)
.handle(
(exists, sink) -> {
if (!exists) {
sink.error(
new MemoryException(
ErrorCodes
.ATTEMPTED_TO_ACCESS_NONEXISTENT_COLLECTION,
collectionName));
return;
}
sink.next(exists);
})
.then(
Flux.fromIterable(records)
.concatMap(record -> internalUpsertAsync(collectionName, record))
.collect(Collectors.toCollection(ArrayList::new)));
}
@Override
public Mono<MemoryRecord> getAsync(
@Nonnull String collectionName, @Nonnull String key, boolean withEmbedding) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(key);
return this.internalGetAsync(collectionName, key, withEmbedding);
}
private MemoryRecord memoryRecordFromDatabaseEntry(DatabaseEntry entry, boolean withEmbedding) {
try {
Embedding embedding =
withEmbedding
? new ObjectMapper().readValue(entry.getEmbedding(), Embedding.class)
: Embedding.empty();
return MemoryRecord.fromJsonMetadata(
entry.getMetadata(), embedding, entry.getKey(), entry.getTimestamp());
} catch (JsonProcessingException e) {
throw new SQLConnectorException(
SQLConnectorException.ErrorCodes.SQL_ERROR,
"Error deserializing database entry",
e);
}
}
private Mono<MemoryRecord> internalGetAsync(
String collectionName, String key, boolean withEmbedding) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(key);
Mono<DatabaseEntry> entry = this.dbConnector.readAsync(collectionName, key);
return entry.hasElement()
.flatMap(
hasElement -> {
if (!hasElement) {
return Mono.empty();
}
return entry.map(
databaseEntry ->
memoryRecordFromDatabaseEntry(
databaseEntry, withEmbedding));
});
}
@Override
public Mono<Collection<MemoryRecord>> getBatchAsync(
@Nonnull String collectionName,
@Nonnull Collection<String> keys,
boolean withEmbeddings) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(keys);
return Flux.fromIterable(keys)
.concatMap(key -> internalGetAsync(collectionName, key, withEmbeddings))
.collect(Collectors.toCollection(ArrayList::new));
}
@Override
public Mono<Void> removeAsync(@Nonnull String collectionName, @Nonnull String key) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(key);
return this.dbConnector.deleteAsync(collectionName, key);
}
@Override
public Mono<Void> removeBatchAsync(
@Nonnull String collectionName, @Nonnull Collection<String> keys) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(keys);
return Flux.fromIterable(keys)
.concatMap(key -> this.dbConnector.deleteAsync(collectionName, key))
.then();
}
@Override
public Mono<Collection<Tuple2<MemoryRecord, Float>>> getNearestMatchesAsync(
@Nonnull String collectionName,
@Nonnull Embedding embedding,
int limit,
double minRelevanceScore,
boolean withEmbeddings) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(embedding);
Mono<List<DatabaseEntry>> entries = this.dbConnector.readAllAsync(collectionName);
return entries.flatMap(
databaseEntries -> {
List<Tuple2<MemoryRecord, Float>> nearestMatches = new ArrayList<>();
for (DatabaseEntry entry : databaseEntries) {
if (entry.getEmbedding() == null || entry.getEmbedding().isEmpty()) {
continue;
}
try {
Embedding recordEmbedding =
new ObjectMapper()
.readValue(entry.getEmbedding(), Embedding.class);
float similarity = embedding.cosineSimilarity(recordEmbedding);
if (similarity >= (float) minRelevanceScore) {
MemoryRecord record =
MemoryRecord.fromJsonMetadata(
entry.getMetadata(),
withEmbeddings ? recordEmbedding : null,
entry.getKey(),
entry.getTimestamp());
nearestMatches.add(Tuples.of(record, similarity));
}
} catch (JsonProcessingException e) {
throw new SQLConnectorException(
SQLConnectorException.ErrorCodes.SQL_ERROR,
"Error deserializing database entry",
e);
}
}
List<Tuple2<MemoryRecord, Float>> results =
nearestMatches.stream()
.sorted(
Comparator.comparing(
Tuple2::getT2, (a, b) -> Float.compare(b, a)))
.limit(limit)
.collect(Collectors.toList());
return Mono.just(results);
});
}
@Override
public Mono<Tuple2<MemoryRecord, Float>> getNearestMatchAsync(
@Nonnull String collectionName,
@Nonnull Embedding embedding,
double minRelevanceScore,
boolean withEmbedding) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(embedding);
return getNearestMatchesAsync(
collectionName, embedding, 1, minRelevanceScore, withEmbedding)
.flatMap(
nearestMatches -> {
if (nearestMatches.isEmpty()) {
return Mono.empty();
}
return Mono.just(nearestMatches.iterator().next());
});
}
/** Builds a JDBCMemoryStore. */
public static class Builder implements SQLMemoryStore.Builder<JDBCMemoryStore> {
private Connection connection;
/**
* Builds and returns a JDBCMemoryStore instance with the specified database connection. The
* build process will connect to the database and create the required tables.
*
* @return A JDBCMemoryStore instance configured with the provided database connection.
* @deprecated Use {@link #buildAsync()} instead.
*/
@Override
@Deprecated
public JDBCMemoryStore build() {
return this.buildAsync().block();
}
/**
* Asynchronously builds and returns a JDBCMemoryStore instance with the specified database
* connection.
*
* @return A Mono with a JDBCMemoryStore instance configured with the provided database
* connection.
*/
@Override
@CheckReturnValue
public Mono<JDBCMemoryStore> buildAsync() {
JDBCConnector connector = new JDBCConnector(connection);
JDBCMemoryStore memoryStore = new JDBCMemoryStore(connector);
return connector.createTableAsync().thenReturn(memoryStore);
}
/**
* Sets the database connection to be used by the memory store being built.
*
* @param connection The Connection object representing the database connection.
* @return The updated Builder instance to continue the building process for a
* JDBCMemoryStore.
*/
@Override
public Builder withConnection(Connection connection) {
this.connection = connection;
return this;
}
}
}
java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/SQLConnector.java
+115
View File
@@ -0,0 +1,115 @@
// Copyright (c) Microsoft. All rights reserved.
package com.microsoft.semantickernel.connectors.memory.jdbc;
import java.time.ZonedDateTime;
import java.util.List;
import reactor.core.publisher.Mono;
public interface SQLConnector {
/**
* Asynchronously creates a table in the database.
*
* @return A Mono representing the completion of the table creation operation.
*/
Mono<Void> createTableAsync();
/**
* Asynchronously creates a collection within the database.
*
* @param collectionName The name of the collection to be created.
* @return A Mono representing the completion of the collection creation operation.
*/
Mono<Void> createCollectionAsync(String collectionName);
/**
* Asynchronously updates an entry.
*
* @param collection The name of the collection containing the entry.
* @param key The key identifying the entry to update.
* @param metadata The metadata associated with the entry.
* @param embedding The embedding data associated with the entry.
* @param timestamp The timestamp of the update.
* @return A Mono representing the completion of the update operation.
*/
Mono<Void> updateAsync(
String collection,
String key,
String metadata,
String embedding,
ZonedDateTime timestamp);
/**
* Asynchronously inserts or ignores an entry.
*
* @param collection The name of the collection to insert the entry into.
* @param key The key identifying the entry.
* @param metadata The metadata associated with the entry.
* @param embedding The embedding data associated with the entry.
* @param timestamp The timestamp of the insertion.
* @return A Mono representing the completion of the insertion operation.
*/
Mono<Void> insertOrIgnoreAsync(
String collection,
String key,
String metadata,
String embedding,
ZonedDateTime timestamp);
/**
* Asynchronously checks if a collection exists.
*
* @param collectionName The name of the collection to check.
* @return A Mono emitting a boolean value indicating whether the collection exists.
*/
Mono<Boolean> doesCollectionExistsAsync(String collectionName);
/**
* Asynchronously retrieves a list of collection names.
*
* @return A Mono emitting a list of collection names.
*/
Mono<List<String>> getCollectionsAsync();
/**
* Asynchronously reads all entries within a collection.
*
* @param collectionName The name of the collection to read from.
* @return A Mono emitting a list of DatabaseEntry objects representing the entries.
*/
Mono<List<DatabaseEntry>> readAllAsync(String collectionName);
/**
* Asynchronously reads a specific entry within a collection.
*
* @param collectionName The name of the collection to read from.
* @param key The key identifying the entry to read.
* @return A Mono emitting a DatabaseEntry object representing the entry.
*/
Mono<DatabaseEntry> readAsync(String collectionName, String key);
/**
* Asynchronously deletes a collection.
*
* @param collectionName The name of the collection to delete.
* @return A Mono representing the completion of the collection deletion operation.
*/
Mono<Void> deleteCollectionAsync(String collectionName);
/**
* Asynchronously deletes an entry within a collection.
*
* @param collectionName The name of the collection to delete the entry from.
* @param key The key identifying the entry to delete.
* @return A Mono representing the completion of the entry deletion operation.
*/
Mono<Void> deleteAsync(String collectionName, String key);
/**
* Asynchronously deletes all empty entries within a collection.
*
* @param collectionName The name of the collection to delete empty entries from.
* @return A Mono representing the completion of the empty entry deletion operation.
*/
Mono<Void> deleteEmptyAsync(String collectionName);
}
java/connectors/semantickernel-connectors-memory-sqlite/src/main/java/com/microsoft/semantickernel/connectors/memory/sqlite/SQLConnectorException.java → java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/SQLConnectorException.java
+2 -2
View File
@@ -1,11 +1,11 @@
// Copyright (c) Microsoft. All rights reserved.
package com.microsoft.semantickernel.connectors.memory.sqlite;
package com.microsoft.semantickernel.connectors.memory.jdbc;
import com.microsoft.semantickernel.SKException;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
/** Exception thrown by the SQLite connector. */
/** Exception thrown by the SQL connector. */
public class SQLConnectorException extends SKException {
private final ErrorCodes errorCode;
java/connectors/semantickernel-connectors-memory-jdbc/src/main/java/com/microsoft/semantickernel/connectors/memory/jdbc/SQLMemoryStore.java
+25
View File
@@ -0,0 +1,25 @@
// Copyright (c) Microsoft. All rights reserved.
package com.microsoft.semantickernel.connectors.memory.jdbc;
import com.microsoft.semantickernel.memory.MemoryStore;
import java.sql.Connection;
import reactor.core.publisher.Mono;
public interface SQLMemoryStore extends MemoryStore {
interface Builder<T extends SQLMemoryStore> extends MemoryStore.Builder<T> {
/**
* Asynchronously builds the SQLMemoryStore.
*
* @return A Mono with the SQLMemoryStore.
*/
Mono<T> buildAsync();
/**
* Sets the database connection to be used by the SQL memory store.
*
* @param connection The Connection object representing the database connection.
* @return The updated SQLMemoryStoreBuilder instance.
*/
Builder<T> withConnection(Connection connection);
}
}
java/connectors/semantickernel-connectors-memory-postgresql/pom.xml
+53
View File
@@ -0,0 +1,53 @@
<?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.9-alpha-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>semantickernel-connectors-memory-postgresql</artifactId>
<name>Semantic Kernel PostgreSQL Memory Connector</name>
<dependencies>
<dependency>
<groupId>com.fasterxml.jackson.core</groupId>
<artifactId>jackson-core</artifactId>
<scope>compile</scope>
</dependency>
<dependency>
<groupId>com.fasterxml.jackson.core</groupId>
<artifactId>jackson-databind</artifactId>
<scope>compile</scope>
</dependency>
<dependency>
<groupId>com.microsoft.semantic-kernel</groupId>
<artifactId>semantickernel-api</artifactId>
</dependency>
<dependency>
<groupId>org.mockito</groupId>
<artifactId>mockito-core</artifactId>
</dependency>
<dependency>
<groupId>org.junit.jupiter</groupId>
<artifactId>junit-jupiter</artifactId>
</dependency>
<dependency>
<groupId>org.postgresql</groupId>
<artifactId>postgresql</artifactId>
<version>LATEST</version>
</dependency>
<dependency>
<groupId>com.microsoft.semantic-kernel</groupId>
<artifactId>semantickernel-connectors-memory-jdbc</artifactId>
<version>${project.version}</version>
<scope>compile</scope>
</dependency>
</dependencies>
</project>
java/connectors/semantickernel-connectors-memory-postgresql/src/main/java/com/microsoft/semantickernel/connectors/memory/postgresql/PostgreSQLConnector.java
+102
View File
@@ -0,0 +1,102 @@
// Copyright (c) Microsoft. All rights reserved.
package com.microsoft.semantickernel.connectors.memory.postgresql;
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.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Statement;
import java.time.ZonedDateTime;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;
public class PostgreSQLConnector extends JDBCConnector implements SQLConnector {
public PostgreSQLConnector(Connection connection) {
super(connection);
}
@Override
public Mono<Void> createTableAsync() {
return Mono.fromRunnable(
() -> {
String createCollectionKeyTable =
"CREATE TABLE IF NOT EXISTS "
+ COLLECTIONS_TABLE_NAME
+ " ("
+ "id TEXT PRIMARY KEY"
+ " )";
String createSKMemoryTable =
"CREATE TABLE IF NOT EXISTS "
+ TABLE_NAME
+ " ("
+ "collection TEXT, "
+ "key TEXT, "
+ "metadata TEXT, "
+ "embedding TEXT, "
+ "timestamp TEXT, "
+ "PRIMARY KEY (collection, key), "
+ "FOREIGN KEY (collection) REFERENCES "
+ COLLECTIONS_TABLE_NAME
+ "(id)"
+ " )";
String createIndex =
"CREATE INDEX IF NOT EXISTS "
+ INDEX_NAME
+ " ON "
+ TABLE_NAME
+ "(collection)";
try (Statement statement = this.connection.createStatement()) {
statement.addBatch(createCollectionKeyTable);
statement.addBatch(createSKMemoryTable);
statement.addBatch(createIndex);
statement.executeBatch();
} catch (SQLException e) {
throw new SQLConnectorException(
SQLConnectorException.ErrorCodes.SQL_ERROR,
"\"CREATE TABLE\" failed",
e);
}
})
.subscribeOn(Schedulers.boundedElastic())
.then();
}
@Override
public Mono<Void> insertOrIgnoreAsync(
String collection,
String key,
String metadata,
String embedding,
ZonedDateTime timestamp) {
return Mono.fromRunnable(
() -> {
String query =
"INSERT INTO "
+ TABLE_NAME
+ " (collection, key, metadata, embedding, timestamp)"
+ " VALUES (?, ?, ?, ?, ?)"
+ " ON CONFLICT (collection, key) DO NOTHING";
try (PreparedStatement statement =
this.connection.prepareStatement(query)) {
statement.setString(1, collection);
statement.setString(2, key);
statement.setString(3, metadata != null ? metadata : "");
statement.setString(4, embedding != null ? embedding : "");
statement.setString(5, formatDatetime(timestamp));
statement.executeUpdate();
} catch (SQLException e) {
throw new SQLConnectorException(
SQLConnectorException.ErrorCodes.SQL_ERROR,
"\"INSERT INTO\" failed",
e);
}
})
.subscribeOn(Schedulers.boundedElastic())
.then();
}
}
java/connectors/semantickernel-connectors-memory-postgresql/src/main/java/com/microsoft/semantickernel/connectors/memory/postgresql/PostgreSQLMemoryStore.java
+64
View File
@@ -0,0 +1,64 @@
// Copyright (c) Microsoft. All rights reserved.
package com.microsoft.semantickernel.connectors.memory.postgresql;
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 PostgreSQLMemoryStore extends JDBCMemoryStore {
private PostgreSQLMemoryStore(SQLConnector connector) {
super(connector);
}
/** Builds a PostgreSQLMemoryStore. */
public static class Builder implements SQLMemoryStore.Builder<PostgreSQLMemoryStore> {
private Connection connection;
/**
* Builds and returns a PostgreSQLMemoryStore instance with the specified database
* connection. The build process will connect to the database and create the required
* tables.
*
* @return A PostgreSQLMemoryStore instance configured with the provided database
* connection.
* @deprecated Use {@link #buildAsync()} instead.
*/
@Override
@Deprecated
public PostgreSQLMemoryStore build() {
return this.buildAsync().block();
}
/**
* Asynchronously builds and returns a PostgreSQLMemoryStore instance with the specified
* database connection.
*
* @return A Mono with a PostgreSQLMemoryStore instance configured with the provided
* database connection.
*/
@Override
@CheckReturnValue
public Mono<PostgreSQLMemoryStore> buildAsync() {
PostgreSQLConnector connector = new PostgreSQLConnector(connection);
PostgreSQLMemoryStore memoryStore = new PostgreSQLMemoryStore(connector);
return connector.createTableAsync().thenReturn(memoryStore);
}
/**
* Sets the PostgreSQL database connection to be used by the PostgreSQL memory store being
* built.
*
* @param connection The Connection object representing the PostgreSQL database connection.
* @return The updated Builder instance to continue the building process for a
* PostgreSQLMemoryStore.
*/
@Override
public Builder withConnection(Connection connection) {
this.connection = connection;
return this;
}
}
}
java/connectors/semantickernel-connectors-memory-sqlite/pom.xml
+4 -14
View File
@@ -44,21 +44,11 @@
<artifactId>sqlite-jdbc</artifactId>
</dependency>
<dependency>
<groupId>org.apache.logging.log4j</groupId>
<artifactId>log4j-api</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.apache.logging.log4j</groupId>
<artifactId>log4j-core</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.apache.logging.log4j</groupId>
<artifactId>log4j-slf4j2-impl</artifactId>
<scope>test</scope>
<groupId>com.microsoft.semantic-kernel</groupId>
<artifactId>semantickernel-connectors-memory-jdbc</artifactId>
<version>${project.version}</version>
<scope>compile</scope>
</dependency>
</dependencies>
java/connectors/semantickernel-connectors-memory-sqlite/src/main/java/com/microsoft/semantickernel/connectors/memory/sqlite/SQLiteMemoryStore.java
+61 -288
View File
@@ -1,303 +1,76 @@
// Copyright (c) Microsoft. All rights reserved.
package com.microsoft.semantickernel.connectors.memory.sqlite;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.microsoft.semantickernel.ai.embeddings.Embedding;
import com.microsoft.semantickernel.memory.MemoryException;
import com.microsoft.semantickernel.memory.MemoryException.ErrorCodes;
import com.microsoft.semantickernel.memory.MemoryRecord;
import com.microsoft.semantickernel.memory.MemoryStore;
import com.microsoft.semantickernel.connectors.memory.jdbc.JDBCConnector;
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 java.sql.DriverManager;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Comparator;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;
import javax.annotation.Nonnull;
import reactor.core.publisher.Flux;
import javax.annotation.CheckReturnValue;
import reactor.core.publisher.Mono;
import reactor.util.function.Tuple2;
import reactor.util.function.Tuples;
public class SQLiteMemoryStore implements MemoryStore {
private final Database dbConnector;
private Connection dbConnection;
public SQLiteMemoryStore() {
this.dbConnector = new Database();
public class SQLiteMemoryStore extends JDBCMemoryStore {
private SQLiteMemoryStore(SQLConnector connector) {
super(connector);
}
public Mono<Void> connectAsync(@Nonnull String filename) throws SQLException {
Objects.requireNonNull(filename);
this.dbConnection = DriverManager.getConnection("jdbc:sqlite:" + filename);
return this.dbConnector.createTableAsync(this.dbConnection);
}
/** Builds a SQLiteMemoryStore. */
public static class Builder implements SQLMemoryStore.Builder<SQLiteMemoryStore> {
private Connection connection;
@Override
public Mono<Void> createCollectionAsync(@Nonnull String collectionName) {
Objects.requireNonNull(collectionName);
return this.dbConnector.createCollectionAsync(this.dbConnection, collectionName);
}
@Override
public Mono<List<String>> getCollectionsAsync() {
return this.dbConnector.getCollectionsAsync(this.dbConnection);
}
@Override
public Mono<Boolean> doesCollectionExistAsync(@Nonnull String collectionName) {
Objects.requireNonNull(collectionName);
return this.dbConnector.doesCollectionExistsAsync(this.dbConnection, collectionName);
}
@Override
public Mono<Void> deleteCollectionAsync(@Nonnull String collectionName) {
Objects.requireNonNull(collectionName);
return this.dbConnector.deleteCollectionAsync(this.dbConnection, collectionName);
}
@Override
public Mono<String> upsertAsync(@Nonnull String collectionName, @Nonnull MemoryRecord record) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(record);
return doesCollectionExistAsync(collectionName)
.handle(
(exists, sink) -> {
if (!exists) {
sink.error(
new MemoryException(
ErrorCodes
.ATTEMPTED_TO_ACCESS_NONEXISTENT_COLLECTION,
collectionName));
return;
}
sink.next(exists);
})
.then(internalUpsertAsync(collectionName, record));
}
private Mono<String> internalUpsertAsync(
@Nonnull String collectionName, @Nonnull MemoryRecord record) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(record);
try {
Mono<Void> update =
this.dbConnector.updateAsync(
this.dbConnection,
collectionName,
record.getMetadata().getId(),
record.getSerializedMetadata(),
record.getSerializedEmbedding(),
record.getTimestamp());
Mono<Void> insert =
this.dbConnector.insertOrIgnoreAsync(
this.dbConnection,
collectionName,
record.getMetadata().getId(),
record.getSerializedMetadata(),
record.getSerializedEmbedding(),
record.getTimestamp());
return update.then(insert).then(Mono.just(record.getMetadata().getId()));
} catch (JsonProcessingException e) {
throw new MemoryException(
ErrorCodes.UNABLE_TO_SERIALIZE_MEMORY,
String.format(
"collection=%s, key=%s", collectionName, record.getMetadata().getId()),
e);
}
}
@Override
public Mono<Collection<String>> upsertBatchAsync(
@Nonnull String collectionName, @Nonnull Collection<MemoryRecord> records) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(records);
return doesCollectionExistAsync(collectionName)
.handle(
(exists, sink) -> {
if (!exists) {
sink.error(
new MemoryException(
ErrorCodes
.ATTEMPTED_TO_ACCESS_NONEXISTENT_COLLECTION,
collectionName));
return;
}
sink.next(exists);
})
.then(
Flux.fromIterable(records)
.concatMap(record -> internalUpsertAsync(collectionName, record))
.collect(Collectors.toCollection(ArrayList::new)));
}
@Override
public Mono<MemoryRecord> getAsync(
@Nonnull String collectionName, @Nonnull String key, boolean withEmbedding) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(key);
return this.internalGetAsync(collectionName, key, withEmbedding);
}
private Mono<MemoryRecord> internalGetAsync(
@Nonnull String collectionName, @Nonnull String key, boolean withEmbedding) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(key);
Mono<Database.DatabaseEntry> entry =
this.dbConnector.readAsync(this.dbConnection, collectionName, key);
return entry.hasElement()
.flatMap(
hasElement -> {
if (!hasElement) {
return Mono.empty();
}
return entry.map(
databaseEntry -> {
try {
if (withEmbedding) {
return MemoryRecord.fromJsonMetadata(
databaseEntry.getMetadata(),
new ObjectMapper()
.readValue(
databaseEntry
.getEmbedding(),
Embedding.class),
databaseEntry.getKey(),
databaseEntry.getTimestamp());
}
return MemoryRecord.fromJsonMetadata(
databaseEntry.getMetadata(),
Embedding.empty(),
databaseEntry.getKey(),
databaseEntry.getTimestamp());
} catch (JsonProcessingException e) {
throw new MemoryException(
ErrorCodes.UNABLE_TO_DESERIALIZE_MEMORY,
String.format(
"collection=%s, key=%s",
collectionName, databaseEntry.getKey()),
e);
}
});
});
}
@Override
public Mono<Collection<MemoryRecord>> getBatchAsync(
@Nonnull String collectionName,
@Nonnull Collection<String> keys,
boolean withEmbeddings) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(keys);
return Flux.fromIterable(keys)
.flatMap(key -> internalGetAsync(collectionName, key, withEmbeddings))
.collect(Collectors.toCollection(ArrayList::new));
}
@Override
public Mono<Void> removeAsync(@Nonnull String collectionName, @Nonnull String key) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(key);
return this.dbConnector.deleteAsync(this.dbConnection, collectionName, key);
}
@Override
public Mono<Void> removeBatchAsync(
@Nonnull String collectionName, @Nonnull Collection<String> keys) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(keys);
return Flux.fromIterable(keys)
.flatMap(
key -> this.dbConnector.deleteAsync(this.dbConnection, collectionName, key))
.then();
}
@Override
public Mono<Collection<Tuple2<MemoryRecord, Float>>> getNearestMatchesAsync(
@Nonnull String collectionName,
@Nonnull Embedding embedding,
int limit,
double minRelevanceScore,
boolean withEmbeddings) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(embedding);
Mono<List<Database.DatabaseEntry>> entries =
this.dbConnector.readAllAsync(this.dbConnection, collectionName);
return entries.flatMap(
databaseEntries -> {
List<Tuple2<MemoryRecord, Float>> nearestMatches = new ArrayList<>();
for (Database.DatabaseEntry entry : databaseEntries) {
if (entry.getEmbedding() == null || entry.getEmbedding().isEmpty()) {
continue;
}
try {
Embedding recordEmbedding =
new ObjectMapper()
.readValue(entry.getEmbedding(), Embedding.class);
float similarity = embedding.cosineSimilarity(recordEmbedding);
if (similarity >= (float) minRelevanceScore) {
MemoryRecord record =
MemoryRecord.fromJsonMetadata(
entry.getMetadata(),
withEmbeddings ? recordEmbedding : null,
entry.getKey(),
entry.getTimestamp());
nearestMatches.add(Tuples.of(record, similarity));
}
} catch (JsonProcessingException e) {
throw new MemoryException(
ErrorCodes.UNABLE_TO_DESERIALIZE_MEMORY,
String.format(
"collection=%s, key=%s",
collectionName, entry.getKey()),
e);
}
}
List<Tuple2<MemoryRecord, Float>> results =
nearestMatches.stream()
.sorted(
Comparator.comparing(
Tuple2::getT2, (a, b) -> Float.compare(b, a)))
.limit(limit)
.collect(Collectors.toList());
return Mono.just(results);
});
}
@Override
public Mono<Tuple2<MemoryRecord, Float>> getNearestMatchAsync(
@Nonnull String collectionName,
@Nonnull Embedding embedding,
double minRelevanceScore,
boolean withEmbedding) {
Objects.requireNonNull(collectionName);
Objects.requireNonNull(embedding);
return getNearestMatchesAsync(
collectionName, embedding, 1, minRelevanceScore, withEmbedding)
.flatMap(
nearestMatches -> {
if (nearestMatches.isEmpty()) {
return Mono.empty();
}
return Mono.just(nearestMatches.iterator().next());
});
}
public static class Builder implements MemoryStore.Builder {
/**
* Builds and returns an SQLiteMemoryStore instance with the specified database connection.
* The build process will connect to the database and create the required tables.
*
* @return An SQLiteMemoryStore instance configured with the provided database connection.
* @deprecated Use {@link #buildAsync()} instead.
*/
@Override
public MemoryStore build() {
return new SQLiteMemoryStore();
@Deprecated
public SQLiteMemoryStore build() {
return this.buildAsync().block();
}
/**
* Asynchronously builds and returns an SQLiteMemoryStore instance with the specified
* database connection.
*
* @return A Mono with an SQLiteMemoryStore instance configured with the provided SQLite
* database connection.
*/
@Override
@CheckReturnValue
public Mono<SQLiteMemoryStore> buildAsync() {
JDBCConnector connector = new JDBCConnector(connection);
SQLiteMemoryStore memoryStore = new SQLiteMemoryStore(connector);
return connector.createTableAsync().thenReturn(memoryStore);
}
/**
* Sets the SQLite database connection to be used by the SQLite memory store being built.
*
* @param connection The Connection object representing the SQLite database connection.
* @return The updated Builder instance to continue the building process for an
* SQLiteMemoryStore.
*/
@Override
public Builder withConnection(Connection connection) {
this.connection = connection;
return this;
}
/**
* Creates and sets an SQLite database connection using the specified filename.
* This will overwrite any previously set connection.
*
* @param filename The filename of the SQLite database.
* @return The updated Builder instance with the SQLite database connection set.
* @throws SQLException If there is an issue with establishing the database connection.
*/
public Builder withFilename(String filename) throws SQLException {
return withConnection(DriverManager.getConnection("jdbc:sqlite:" + filename));
}
}
}
java/connectors/semantickernel-connectors-memory-sqlite/src/test/java/com/microsoft/semantickernel/connectors/memory/sqlite/SQLiteMemoryStoreTest.java
+135 -384
View File
@@ -23,25 +23,22 @@ import java.util.Iterator;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import org.junit.jupiter.api.BeforeEach;
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 {
private MemoryStore _db;
private int _collectionNum = 0;
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;
@BeforeEach
void setUp() throws SQLException {
_db = new SQLiteMemoryStore.Builder().build();
((SQLiteMemoryStore) _db).connectAsync(":memory:").block();
@BeforeAll
static void setUp() throws SQLException {
db = new SQLiteMemoryStore.Builder().withFilename(":memory:").buildAsync().block();
}
private Collection<MemoryRecord> createBatchRecords(int numRecords) {
@@ -80,18 +77,18 @@ public class SQLiteMemoryStoreTest {
@Test
void initializeDbConnectionSucceeds() {
assertNotNull(this._db);
assertNotNull(db);
}
@Test
void itCanCreateAndGetCollectionAsync() {
// Arrange
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
String collection = "test_collection" + collectionNum;
collectionNum++;
// Act
this._db.createCollectionAsync(collection).block();
Collection<String> collections = this._db.getCollectionsAsync().block();
db.createCollectionAsync(collection).block();
Collection<String> collections = db.getCollectionsAsync().block();
// Assert
assertNotNull(collections);
@@ -107,7 +104,7 @@ public class SQLiteMemoryStoreTest {
// Assert
assertThrows(
NullPointerException.class,
() -> this._db.createCollectionAsync(collection).block(),
() -> db.createCollectionAsync(collection).block(),
"Should not be able to create collection with null name");
}
@@ -124,13 +121,13 @@ public class SQLiteMemoryStoreTest {
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
String collection = "test_collection" + collectionNum;
collectionNum++;
// Assert
assertThrows(
MemoryException.class,
() -> this._db.upsertAsync(collection, testRecord).block(),
() -> db.upsertAsync(collection, testRecord).block(),
"Should not be able to insert into a non-existent collection");
}
@@ -146,14 +143,14 @@ public class SQLiteMemoryStoreTest {
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
String collection = "test_collection" + collectionNum;
collectionNum++;
// Act
this._db.createCollectionAsync(collection).block();
String key = this._db.upsertAsync(collection, testRecord).block();
MemoryRecord actualDefault = this._db.getAsync(collection, key, false).block();
MemoryRecord actualWithEmbedding = this._db.getAsync(collection, key, true).block();
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);
@@ -183,13 +180,13 @@ public class SQLiteMemoryStoreTest {
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
String collection = "test_collection" + collectionNum;
collectionNum++;
// Act
this._db.createCollectionAsync(collection).block();
String key = this._db.upsertAsync(collection, testRecord).block();
MemoryRecord actual = this._db.getAsync(collection, key, true).block();
db.createCollectionAsync(collection).block();
String key = db.upsertAsync(collection, testRecord).block();
MemoryRecord actual = db.getAsync(collection, key, true).block();
// Assert
assertNotNull(actual);
@@ -209,13 +206,13 @@ public class SQLiteMemoryStoreTest {
NULL_ADDITIONAL_METADATA,
NULL_KEY,
ZonedDateTime.now(ZoneId.of("UTC")));
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
String collection = "test_collection" + collectionNum;
collectionNum++;
// Act
this._db.createCollectionAsync(collection).block();
String key = this._db.upsertAsync(collection, testRecord).block();
MemoryRecord actual = this._db.getAsync(collection, key, true).block();
db.createCollectionAsync(collection).block();
String key = db.upsertAsync(collection, testRecord).block();
MemoryRecord actual = db.getAsync(collection, key, true).block();
// Assert
assertNotNull(actual);
@@ -245,14 +242,14 @@ public class SQLiteMemoryStoreTest {
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
String collection = "test_collection" + collectionNum;
collectionNum++;
// Act
this._db.createCollectionAsync(collection).block();
String key = this._db.upsertAsync(collection, testRecord).block();
String key2 = this._db.upsertAsync(collection, testRecord2).block();
MemoryRecord actual = this._db.getAsync(collection, key, true).block();
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);
@@ -274,15 +271,15 @@ public class SQLiteMemoryStoreTest {
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
String collection = "test_collection" + collectionNum;
collectionNum++;
// Act
this._db.createCollectionAsync(collection).block();
String key = this._db.upsertAsync(collection, testRecord).block();
db.createCollectionAsync(collection).block();
String key = db.upsertAsync(collection, testRecord).block();
assertNotNull(key);
this._db.removeAsync(collection, key).block();
MemoryRecord actual = this._db.getAsync(collection, key, false).block();
db.removeAsync(collection, key).block();
MemoryRecord actual = db.getAsync(collection, key, false).block();
// Assert
assertNull(actual);
@@ -291,12 +288,12 @@ public class SQLiteMemoryStoreTest {
@Test
void removingNonExistingRecordDoesNothingAsync() {
// Arrange
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
String collection = "test_collection" + collectionNum;
collectionNum++;
// Act
this._db.removeAsync(collection, "key").block();
MemoryRecord actual = this._db.getAsync(collection, "key", false).block();
db.removeAsync(collection, "key").block();
MemoryRecord actual = db.getAsync(collection, "key", false).block();
// Assert
assertNull(actual);
@@ -307,20 +304,24 @@ public class SQLiteMemoryStoreTest {
// Arrange
int numCollections = 3;
String[] testCollections =
IntStream.range(this._collectionNum, this._collectionNum += numCollections)
IntStream.range(collectionNum, collectionNum += numCollections)
.mapToObj(i -> "test_collection" + i)
.toArray(String[]::new);
Collection<String> collections = this.db.getCollectionsAsync().block();
assertNotNull(collections);
int initialSize = collections.size();
Flux.fromArray(testCollections)
.concatMap(collection -> this._db.createCollectionAsync(collection))
.concatMap(collection -> this.db.createCollectionAsync(collection))
.blockLast();
// Act
Collection<String> collections = this._db.getCollectionsAsync().block();
collections = this.db.getCollectionsAsync().block();
// Assert
assertNotNull(collections);
assertEquals(numCollections, collections.size());
assertEquals(initialSize + numCollections, collections.size());
for (String collection : testCollections) {
assertTrue(
collections.contains(collection),
@@ -328,14 +329,10 @@ public class SQLiteMemoryStoreTest {
}
}
@Test
void getNearestMatchesReturnsAllResultsWithNoMinScoreAsync() {
// Arrange
Embedding compareEmbedding = new Embedding(Arrays.asList(1f, 1f, 1f));
int topN = 4;
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
this._db.createCollectionAsync(collection).block();
private String setUpNearestMatches() {
String collection = "test_collection" + collectionNum;
collectionNum++;
db.createCollectionAsync(collection).block();
int i = 0;
MemoryRecord testRecord =
MemoryRecord.localRecord(
@@ -346,7 +343,7 @@ public class SQLiteMemoryStoreTest {
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
this._db.upsertAsync(collection, testRecord).block();
db.upsertAsync(collection, testRecord).block();
i++;
testRecord =
@@ -358,7 +355,7 @@ public class SQLiteMemoryStoreTest {
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
this._db.upsertAsync(collection, testRecord).block();
db.upsertAsync(collection, testRecord).block();
i++;
testRecord =
@@ -370,7 +367,7 @@ public class SQLiteMemoryStoreTest {
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
this._db.upsertAsync(collection, testRecord).block();
db.upsertAsync(collection, testRecord).block();
i++;
testRecord =
@@ -382,7 +379,7 @@ public class SQLiteMemoryStoreTest {
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
this._db.upsertAsync(collection, testRecord).block();
db.upsertAsync(collection, testRecord).block();
i++;
testRecord =
@@ -394,14 +391,22 @@ public class SQLiteMemoryStoreTest {
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
this._db.upsertAsync(collection, testRecord).block();
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
double threshold = -1;
Collection<Tuple2<MemoryRecord, Float>> topNResults =
this._db
.getNearestMatchesAsync(
collection, compareEmbedding, topN, threshold, false)
db.getNearestMatchesAsync(collection, compareEmbedding, topN, threshold, false)
.block();
// Assert
@@ -419,154 +424,29 @@ public class SQLiteMemoryStoreTest {
void getNearestMatchesReturnsLimit() {
// Arrange
Embedding compareEmbedding = new Embedding(Arrays.asList(1f, 1f, 1f));
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
this._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);
this._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);
this._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);
this._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);
this._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);
this._db.upsertAsync(collection, testRecord).block();
String collection = setUpNearestMatches();
// Act
double threshold = -1;
Collection<Tuple2<MemoryRecord, Float>> topNResults =
this._db
.getNearestMatchesAsync(
collection, compareEmbedding, i / 2, threshold, false)
db.getNearestMatchesAsync(collection, compareEmbedding, 0, threshold, false)
.block();
// Assert
assertNotNull(topNResults);
assertEquals(i / 2, topNResults.size());
assertTrue(topNResults.isEmpty());
}
@Test
void getNearestMatchesReturnsEmptyIfLimitZero() {
// Arrange
Embedding compareEmbedding = new Embedding(Arrays.asList(1f, 1f, 1f));
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
this._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);
this._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);
this._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);
this._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);
this._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);
this._db.upsertAsync(collection, testRecord).block();
String collection = setUpNearestMatches();
// Act
double threshold = -1;
Collection<Tuple2<MemoryRecord, Float>> topNResults =
this._db
.getNearestMatchesAsync(collection, compareEmbedding, 0, threshold, false)
db.getNearestMatchesAsync(collection, compareEmbedding, 0, threshold, false)
.block();
// Assert
@@ -578,15 +458,14 @@ public class SQLiteMemoryStoreTest {
void getNearestMatchesReturnsEmptyIfCollectionEmpty() {
// Arrange
Embedding compareEmbedding = new Embedding(Arrays.asList(1f, 1f, 1f));
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
this._db.createCollectionAsync(collection).block();
String collection = "test_collection" + collectionNum;
collectionNum++;
db.createCollectionAsync(collection).block();
// Act
double threshold = -1;
Collection<Tuple2<MemoryRecord, Float>> topNResults =
this._db
.getNearestMatchesAsync(
db.getNearestMatchesAsync(
collection, compareEmbedding, Integer.MAX_VALUE, threshold, false)
.block();
@@ -600,79 +479,15 @@ public class SQLiteMemoryStoreTest {
// Arrange
Embedding compareEmbedding = new Embedding(Arrays.asList(1f, 1f, 1f));
int topN = 4;
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
this._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);
this._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);
this._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);
this._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);
this._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);
this._db.upsertAsync(collection, testRecord).block();
String collection = setUpNearestMatches();
// Act
double threshold = 0.75;
Tuple2<MemoryRecord, Float> topNResultDefault =
this._db
.getNearestMatchAsync(collection, compareEmbedding, threshold, false)
.block();
db.getNearestMatchAsync(collection, compareEmbedding, threshold, false).block();
Tuple2<MemoryRecord, Float> topNResultWithEmbedding =
this._db
.getNearestMatchAsync(collection, compareEmbedding, threshold, true)
.block();
db.getNearestMatchAsync(collection, compareEmbedding, threshold, true).block();
// Assert
assertNotNull(topNResultDefault);
@@ -690,75 +505,12 @@ public class SQLiteMemoryStoreTest {
// Arrange
Embedding compareEmbedding = new Embedding(Arrays.asList(1f, 1f, 1f));
int topN = 4;
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
this._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);
this._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);
this._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);
this._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);
this._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);
this._db.upsertAsync(collection, testRecord).block();
String collection = setUpNearestMatches();
// Act
double threshold = 0.75;
Tuple2<MemoryRecord, Float> topNResult =
this._db
.getNearestMatchAsync(collection, compareEmbedding, threshold, false)
.block();
db.getNearestMatchAsync(collection, compareEmbedding, threshold, false).block();
// Assert
assertNotNull(topNResult);
@@ -771,16 +523,14 @@ public class SQLiteMemoryStoreTest {
// Arrange
Embedding compareEmbedding = new Embedding(Arrays.asList(1f, 1f, 1f));
int topN = 4;
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
this._db.createCollectionAsync(collection).block();
String collection = "test_collection" + collectionNum;
collectionNum++;
db.createCollectionAsync(collection).block();
// Act
double threshold = -1;
Tuple2<MemoryRecord, Float> topNResults =
this._db
.getNearestMatchAsync(collection, compareEmbedding, threshold, false)
.block();
db.getNearestMatchAsync(collection, compareEmbedding, threshold, false).block();
// Assert
assertNull(topNResults);
@@ -791,9 +541,9 @@ public class SQLiteMemoryStoreTest {
// Arrange
Embedding compareEmbedding = new Embedding(Arrays.asList(1f, 1f, 1f));
int topN = 4;
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
this._db.createCollectionAsync(collection).block();
String collection = "test_collection" + collectionNum;
collectionNum++;
db.createCollectionAsync(collection).block();
for (int i = 0; i < 10; i++) {
MemoryRecord testRecord =
@@ -805,14 +555,12 @@ public class SQLiteMemoryStoreTest {
NULL_ADDITIONAL_METADATA,
NULL_KEY,
NULL_TIMESTAMP);
this._db.upsertAsync(collection, testRecord).block();
db.upsertAsync(collection, testRecord).block();
}
// Act
Collection<Tuple2<MemoryRecord, Float>> topNResults =
this._db
.getNearestMatchesAsync(collection, compareEmbedding, topN, 0.75, true)
.block();
db.getNearestMatchesAsync(collection, compareEmbedding, topN, 0.75, true).block();
Collection<String> topNKeys =
topNResults.stream()
.map(tuple -> tuple.getT1().getKey())
@@ -835,15 +583,14 @@ public class SQLiteMemoryStoreTest {
void itCanBatchUpsertRecordsAsync() {
// Arrange
int numRecords = 10;
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
this._db.createCollectionAsync(collection).block();
String collection = "test_collection" + collectionNum;
collectionNum++;
db.createCollectionAsync(collection).block();
Collection<MemoryRecord> records = this.createBatchRecords(numRecords);
// Act
Collection<String> keys = this._db.upsertBatchAsync(collection, records).block();
Collection<MemoryRecord> resultRecords =
this._db.getBatchAsync(collection, keys, false).block();
Collection<String> keys = db.upsertBatchAsync(collection, records).block();
Collection<MemoryRecord> resultRecords = db.getBatchAsync(collection, keys, false).block();
// Assert
assertNotNull(keys);
@@ -855,14 +602,14 @@ public class SQLiteMemoryStoreTest {
void itCanBatchGetRecordsAsync() {
// Arrange
int numRecords = 10;
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
this._db.createCollectionAsync(collection).block();
String collection = "test_collection" + collectionNum;
collectionNum++;
db.createCollectionAsync(collection).block();
Collection<MemoryRecord> records = this.createBatchRecords(numRecords);
Collection<String> keys = this._db.upsertBatchAsync(collection, records).block();
Collection<String> keys = db.upsertBatchAsync(collection, records).block();
// Act
Collection<MemoryRecord> results = this._db.getBatchAsync(collection, keys, false).block();
Collection<MemoryRecord> results = db.getBatchAsync(collection, keys, false).block();
// Assert
assertNotNull(keys);
@@ -874,21 +621,21 @@ public class SQLiteMemoryStoreTest {
void itCanBatchRemoveRecordsAsync() {
// Arrange
int numRecords = 10;
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
this._db.createCollectionAsync(collection).block();
String collection = "test_collection" + collectionNum;
collectionNum++;
db.createCollectionAsync(collection).block();
Collection<MemoryRecord> records = this.createBatchRecords(numRecords);
List<String> keys = new ArrayList<>();
for (String key : this._db.upsertBatchAsync(collection, records).block()) {
for (String key : db.upsertBatchAsync(collection, records).block()) {
keys.add(key);
}
// Act
this._db.removeBatchAsync(collection, keys).block();
db.removeBatchAsync(collection, keys).block();
// Assert
for (MemoryRecord result : this._db.getBatchAsync(collection, keys, true).block()) {
for (MemoryRecord result : db.getBatchAsync(collection, keys, true).block()) {
assertNull(result);
}
}
@@ -898,26 +645,31 @@ public class SQLiteMemoryStoreTest {
// Arrange
int numCollections = 3;
String[] testCollections =
IntStream.range(this._collectionNum, this._collectionNum += numCollections)
IntStream.range(collectionNum, collectionNum += numCollections)
.mapToObj(i -> "test_collection" + i)
.toArray(String[]::new);
Collection<String> collections = this.db.getCollectionsAsync().block();
assertNotNull(collections);
int initialSize = collections.size();
Flux.fromArray(testCollections)
.concatMap(collection -> this._db.createCollectionAsync(collection))
.concatMap(collection -> this.db.createCollectionAsync(collection))
.blockLast();
collectionNum += numCollections;
// Act
Collection<String> collections = this._db.getCollectionsAsync().block();
collections = this.db.getCollectionsAsync().block();
assertNotNull(collections);
assertEquals(numCollections, collections.size());
assertEquals(initialSize + numCollections, collections.size());
// Act
for (String collection : collections) {
this._db.deleteCollectionAsync(collection).block();
this.db.deleteCollectionAsync(collection).block();
}
// Assert
collections = this._db.getCollectionsAsync().block();
collections = this.db.getCollectionsAsync().block();
assertNotNull(collections);
assertEquals(0, collections.size());
}
@@ -925,32 +677,31 @@ public class SQLiteMemoryStoreTest {
@Test
void itThrowsWhenDeletingNonExistentCollectionAsync() {
// Arrange
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
String collection = "test_collection" + collectionNum;
collectionNum++;
// Act
assertThrows(
MemoryException.class, () -> this._db.deleteCollectionAsync(collection).block());
assertThrows(MemoryException.class, () -> db.deleteCollectionAsync(collection).block());
}
@Test
void doesCollectionExistAsyncReturnTrueForExistingCollection() {
// Arrange
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
this._db.createCollectionAsync(collection).block();
String collection = "test_collection" + collectionNum;
collectionNum++;
db.createCollectionAsync(collection).block();
// Act
assertTrue(this._db.doesCollectionExistAsync(collection).block());
assertTrue(db.doesCollectionExistAsync(collection).block());
}
@Test
void doesCollectionExistAsyncReturnFalseForNonExistentCollection() {
// Arrange
String collection = "test_collection" + this._collectionNum;
this._collectionNum++;
String collection = "test_collection" + collectionNum;
collectionNum++;
// Act
assertFalse(this._db.doesCollectionExistAsync(collection).block());
assertFalse(db.doesCollectionExistAsync(collection).block());
}
}
java/pom.xml
+2
View File
@@ -76,6 +76,8 @@
<module>semantickernel-bom</module>
<module>semantickernel-api</module>
<module>semantickernel-core</module>
<module>connectors/semantickernel-connectors-memory-jdbc</module>
<module>connectors/semantickernel-connectors-memory-postgresql</module>
</modules>
<dependencyManagement>
java/samples/sample-code/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/Example14_SemanticMemory.java
+2 -1
View File
@@ -5,6 +5,7 @@ import com.microsoft.semantickernel.Kernel;
import com.microsoft.semantickernel.SKBuilders;
import com.microsoft.semantickernel.SamplesConfig;
import com.microsoft.semantickernel.connectors.memory.azurecognitivesearch.AzureCognitiveSearchMemory;
import com.microsoft.semantickernel.memory.VolatileMemoryStore;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
@@ -61,7 +62,7 @@ public class Example14_SemanticMemory {
.withOpenAIClient(openAIAsyncClient)
.withModelId("text-embedding-ada-002")
.build())
.withMemoryStorage(SKBuilders.memoryStore().build())
.withMemoryStorage(new VolatileMemoryStore.Builder().build())
.build();
runExampleAsync(kernelWithCustomDb).block();
java/samples/sample-code/src/main/java/com/microsoft/semantickernel/samples/syntaxexamples/Example15_MemorySkill.java
+2 -1
View File
@@ -8,6 +8,7 @@ import com.microsoft.semantickernel.SamplesConfig;
import com.microsoft.semantickernel.ai.embeddings.TextEmbeddingGeneration;
import com.microsoft.semantickernel.coreskills.TextMemorySkill;
import com.microsoft.semantickernel.memory.MemoryStore;
import com.microsoft.semantickernel.memory.VolatileMemoryStore;
import com.microsoft.semantickernel.orchestration.SKContext;
import com.microsoft.semantickernel.orchestration.SKFunction;
import com.microsoft.semantickernel.semanticfunctions.PromptTemplateConfig;
@@ -46,7 +47,7 @@ public class Example15_MemorySkill {
.withModelId("text-embedding-ada-002")
.build();
MemoryStore memoryStore = SKBuilders.memoryStore().build();
MemoryStore memoryStore = new VolatileMemoryStore.Builder().build();
Kernel kernel = SKBuilders.kernel()
.withDefaultAIService(textCompletionService)
java/semantickernel-api/src/main/java/com/microsoft/semantickernel/memory/MemoryException.java
+3
View File
@@ -71,6 +71,9 @@ public class MemoryException extends SKException {
/** Unable to serialize a memory . */
UNABLE_TO_SERIALIZE_MEMORY("Unable to serialize memory"),
/** Unable to build memory. */
UNABLE_TO_BUILD_MEMORY("Unable to build memory"),
/** Attempted to access a memory collection that does not exist. */
ATTEMPTED_TO_ACCESS_NONEXISTENT_COLLECTION("Attempted to access non-existent collection"),
java/semantickernel-api/src/main/java/com/microsoft/semantickernel/memory/MemoryStore.java
+2 -2
View File
@@ -148,9 +148,9 @@ public interface MemoryStore extends Buildable {
double minRelevanceScore,
boolean withEmbedding);
static Builder builder() {
static <T extends MemoryStore> Builder<T> builder() {
return BuildersSingleton.INST.getInstance(Builder.class);
}
interface Builder extends SemanticKernelBuilder<MemoryStore> {}
interface Builder<T extends MemoryStore> extends SemanticKernelBuilder<T> {}
}
java/semantickernel-bom/pom.xml
+11
View File
@@ -75,11 +75,21 @@
<artifactId>semantickernel-gpt3-tokenizer</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>com.microsoft.semantic-kernel</groupId>
<artifactId>semantickernel-connectors-memory-jdbc</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>com.microsoft.semantic-kernel</groupId>
<artifactId>semantickernel-connectors-memory-sqlite</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>com.microsoft.semantic-kernel</groupId>
<artifactId>semantickernel-connectors-memory-postgresql</artifactId>
<version>${project.version}</version>
</dependency>
<dependency>
<groupId>com.microsoft.semantic-kernel</groupId>
<artifactId>semantickernel-connectors-memory-azurecognitivesearch</artifactId>
@@ -155,6 +165,7 @@
<distributionManagement>
<repository>
<id>${releaseRepoId}</id>
<!--suppress UnresolvedMavenProperty -->
<url>${releaseRepoUrl}</url>
<name>${releaseRepoName}</name>
</repository>
java/semantickernel-core/src/main/java/com/microsoft/semantickernel/memory/VolatileMemoryStore.java
+2 -2
View File
@@ -279,9 +279,9 @@ public class VolatileMemoryStore implements MemoryStore {
return collection;
}
public static class Builder implements MemoryStore.Builder {
public static class Builder implements MemoryStore.Builder<VolatileMemoryStore> {
@Override
public MemoryStore build() {
public VolatileMemoryStore build() {
return new VolatileMemoryStore();
}
}
java/spotbugs-exclude.xml
+30
View File
@@ -70,4 +70,34 @@
<Bug pattern="EI_EXPOSE_REP"/>
</Match>
<Match>
<Class name="com.microsoft.semantickernel.connectors.memory.jdbc.JDBCConnector"/>
<Method name="&lt;init&gt;"/>
<Bug pattern="EI_EXPOSE_REP2"/>
</Match>
<Match>
<Class name="com.microsoft.semantickernel.connectors.memory.jdbc.JDBCMemoryStore"/>
<Method name="&lt;init&gt;"/>
<Bug pattern="EI_EXPOSE_REP2"/>
</Match>
<Match>
<Class name="com.microsoft.semantickernel.connectors.memory.jdbc.JDBCMemoryStore$Builder"/>
<Method name="withConnection"/>
<Bug pattern="EI_EXPOSE_REP2"/>
</Match>
<Match>
<Class name="com.microsoft.semantickernel.connectors.memory.sqlite.SQLiteMemoryStore$Builder"/>
<Method name="withConnection"/>
<Bug pattern="EI_EXPOSE_REP2"/>
</Match>
<Match>
<Class name="com.microsoft.semantickernel.connectors.memory.postgresql.PostgreSQLMemoryStore$Builder"/>
<Method name="withConnection"/>
<Bug pattern="EI_EXPOSE_REP2"/>
</Match>
</FindBugsFilter>