Add SQLiteVectorStoreQueryProvider

Author: Milder Hernandez Cagua

api-test/integration-tests/pom.xml

@@ -78,6 +78,11 @@
             <artifactId>postgresql</artifactId>
             <version>42.7.3</version> <!-- Use the latest version -->
         </dependency>
+        <dependency>
+            <groupId>org.xerial</groupId>
+            <artifactId>sqlite-jdbc</artifactId>
+            <version>3.46.1.0</version>
+        </dependency>
 
         <dependency>
             <groupId>org.testcontainers</groupId>

api-test/integration-tests/src/test/java/com/microsoft/semantickernel/tests/connectors/memory/jdbc/JDBCVectorStoreRecordCollectionTest.java

@@ -6,6 +6,7 @@
 import com.microsoft.semantickernel.connectors.data.mysql.MySQLVectorStoreQueryProvider;
 import com.microsoft.semantickernel.connectors.data.postgres.PostgreSQLVectorStoreQueryProvider;
 import com.microsoft.semantickernel.connectors.data.jdbc.filter.SQLEqualToFilterClause;
+import com.microsoft.semantickernel.connectors.data.sqlite.SQLiteVectorStoreQueryProvider;
 import com.microsoft.semantickernel.data.vectorsearch.VectorSearchFilter;
 import com.microsoft.semantickernel.data.vectorsearch.VectorSearchResult;
 import com.microsoft.semantickernel.data.vectorstorage.options.GetRecordOptions;
@@ -17,6 +18,7 @@
 import org.junit.jupiter.params.provider.EnumSource;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.postgresql.ds.PGSimpleDataSource;
+import org.sqlite.SQLiteDataSource;
 import org.testcontainers.containers.MySQLContainer;
 import org.testcontainers.containers.PostgreSQLContainer;
 import org.testcontainers.junit.jupiter.Container;
@@ -47,7 +49,8 @@ public class JDBCVectorStoreRecordCollectionTest {
 
     public enum QueryProvider {
         MySQL,
-        PostgreSQL
+        PostgreSQL,
+        SQLite
     }
 
     private JDBCVectorStoreRecordCollection<Hotel> buildRecordCollection(QueryProvider provider, @Nonnull String collectionName) {
@@ -75,6 +78,14 @@ private JDBCVectorStoreRecordCollection<Hotel> buildRecordCollection(QueryProvid
                         .withDataSource(dataSource)
                         .build();
                 break;
+            case SQLite:
+                SQLiteDataSource sqliteDataSource = new SQLiteDataSource();
+                sqliteDataSource.setUrl("jdbc:sqlite:file:testdb");
+                dataSource = sqliteDataSource;
+                queryProvider = SQLiteVectorStoreQueryProvider.builder()
+                        .withDataSource(sqliteDataSource)
+                        .build();
+                break;
             default:
                 throw new IllegalArgumentException("Unknown query provider: " + provider);
         }
@@ -131,7 +142,7 @@ public void upsertAndGetRecordAsync(QueryProvider provider) {
 
         // Upsert the first time
         for (Hotel hotel : hotels) {
-            Hotel retrievedHotel =  recordCollection.getAsync(hotel.getId(), null).block();
+            Hotel retrievedHotel = recordCollection.getAsync(hotel.getId(), null).block();
             assertNotNull(retrievedHotel);
             assertEquals(hotel.getId(), retrievedHotel.getId());
             assertEquals(hotel.getRating(), retrievedHotel.getRating());
@@ -146,7 +157,7 @@ public void upsertAndGetRecordAsync(QueryProvider provider) {
         }
 
         for (Hotel hotel : hotels) {
-            Hotel retrievedHotel =  recordCollection.getAsync(hotel.getId(), null).block();
+            Hotel retrievedHotel = recordCollection.getAsync(hotel.getId(), null).block();
             assertNotNull(retrievedHotel);
             assertEquals(hotel.getId(), retrievedHotel.getId());
             assertEquals(1.0, retrievedHotel.getRating());
@@ -157,7 +168,7 @@ public void upsertAndGetRecordAsync(QueryProvider provider) {
     @EnumSource(QueryProvider.class)
     public void getBatchAsync(QueryProvider provider) {
         String collectionName = "getBatchAsync";
-        JDBCVectorStoreRecordCollection<Hotel>  recordCollection = buildRecordCollection(provider, collectionName);
+        JDBCVectorStoreRecordCollection<Hotel> recordCollection = buildRecordCollection(provider, collectionName);
 
         List<Hotel> hotels = getHotels();
         for (Hotel hotel : hotels) {
@@ -169,7 +180,7 @@ public void getBatchAsync(QueryProvider provider) {
             keys.add(hotel.getId());
         }
 
-        List<Hotel> retrievedHotels =  recordCollection.getBatchAsync(keys, null).block();
+        List<Hotel> retrievedHotels = recordCollection.getBatchAsync(keys, null).block();
         assertNotNull(retrievedHotels);
         assertEquals(hotels.size(), retrievedHotels.size());
     }
@@ -178,7 +189,7 @@ public void getBatchAsync(QueryProvider provider) {
     @EnumSource(QueryProvider.class)
     public void upsertBatchAndGetBatchAsync(QueryProvider provider) {
         String collectionName = "upsertBatchAndGetBatchAsync";
-        JDBCVectorStoreRecordCollection<Hotel>  recordCollection = buildRecordCollection(provider, collectionName);
+        JDBCVectorStoreRecordCollection<Hotel> recordCollection = buildRecordCollection(provider, collectionName);
 
         List<Hotel> hotels = getHotels();
          recordCollection.upsertBatchAsync(hotels, null).block();
@@ -188,7 +199,7 @@ public void upsertBatchAndGetBatchAsync(QueryProvider provider) {
             keys.add(hotel.getId());
         }
 
-        List<Hotel> retrievedHotels =  recordCollection.getBatchAsync(keys, null).block();
+        List<Hotel> retrievedHotels = recordCollection.getBatchAsync(keys, null).block();
         assertNotNull(retrievedHotels);
         assertEquals(hotels.size(), retrievedHotels.size());
     }
@@ -209,7 +220,7 @@ public void insertAndReplaceAsync(QueryProvider provider) {
             keys.add(hotel.getId());
         }
 
-        List<Hotel> retrievedHotels =  recordCollection.getBatchAsync(keys, null).block();
+        List<Hotel> retrievedHotels = recordCollection.getBatchAsync(keys, null).block();
         assertNotNull(retrievedHotels);
         assertEquals(hotels.size(), retrievedHotels.size());
     }
@@ -225,7 +236,7 @@ public void deleteRecordAsync(QueryProvider provider) {
 
         for (Hotel hotel : hotels) {
              recordCollection.deleteAsync(hotel.getId(), null).block();
-            Hotel retrievedHotel =  recordCollection.getAsync(hotel.getId(), null).block();
+            Hotel retrievedHotel = recordCollection.getAsync(hotel.getId(), null).block();
             assertNull(retrievedHotel);
         }
     }
@@ -247,7 +258,7 @@ public void deleteBatchAsync(QueryProvider provider) {
          recordCollection.deleteBatchAsync(keys, null).block();
 
         for (String key : keys) {
-            Hotel retrievedHotel =  recordCollection.getAsync(key, null).block();
+            Hotel retrievedHotel = recordCollection.getAsync(key, null).block();
             assertNull(retrievedHotel);
         }
     }
@@ -266,7 +277,7 @@ public void getWithNoVectors(QueryProvider provider) {
             .build();
 
         for (Hotel hotel : hotels) {
-            Hotel retrievedHotel =  recordCollection.getAsync(hotel.getId(), options).block();
+            Hotel retrievedHotel = recordCollection.getAsync(hotel.getId(), options).block();
             assertNotNull(retrievedHotel);
             assertEquals(hotel.getId(), retrievedHotel.getId());
             assertNull(retrievedHotel.getEuclidean());
@@ -277,7 +288,7 @@ public void getWithNoVectors(QueryProvider provider) {
             .build();
 
         for (Hotel hotel : hotels) {
-            Hotel retrievedHotel =  recordCollection.getAsync(hotel.getId(), options).block();
+            Hotel retrievedHotel = recordCollection.getAsync(hotel.getId(), options).block();
             assertNotNull(retrievedHotel);
             assertEquals(hotel.getId(), retrievedHotel.getId());
             assertNotNull(retrievedHotel.getEuclidean());
@@ -302,7 +313,7 @@ public void getBatchWithNoVectors(QueryProvider provider) {
             keys.add(hotel.getId());
         }
 
-        List<Hotel> retrievedHotels =  recordCollection.getBatchAsync(keys, options).block();
+        List<Hotel> retrievedHotels = recordCollection.getBatchAsync(keys, options).block();
         assertNotNull(retrievedHotels);
         assertEquals(hotels.size(), retrievedHotels.size());
 
@@ -314,7 +325,7 @@ public void getBatchWithNoVectors(QueryProvider provider) {
             .includeVectors(true)
             .build();
 
-        retrievedHotels =  recordCollection.getBatchAsync(keys, options).block();
+        retrievedHotels = recordCollection.getBatchAsync(keys, options).block();
         assertNotNull(retrievedHotels);
         assertEquals(hotels.size(), retrievedHotels.size());
 

semantickernel-experimental/pom.xml

@@ -112,6 +112,11 @@
             <artifactId>postgresql</artifactId>
             <version>42.7.3</version> <!-- Use the latest version -->
         </dependency>
+        <dependency>
+            <groupId>org.xerial</groupId>
+            <artifactId>sqlite-jdbc</artifactId>
+            <version>3.46.1.0</version>
+        </dependency>
 
     </dependencies>
 

semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/sqlite/SQLiteVectorStoreQueryProvider.java

@@ -0,0 +1,159 @@
+// Copyright (c) Microsoft. All rights reserved.
+package com.microsoft.semantickernel.connectors.data.sqlite;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.microsoft.semantickernel.connectors.data.jdbc.JDBCVectorStoreQueryProvider;
+import com.microsoft.semantickernel.connectors.data.jdbc.SQLVectorStoreQueryProvider;
+import com.microsoft.semantickernel.data.vectorstorage.definition.VectorStoreRecordDefinition;
+import com.microsoft.semantickernel.data.vectorstorage.definition.VectorStoreRecordField;
+import com.microsoft.semantickernel.data.vectorstorage.definition.VectorStoreRecordVectorField;
+import com.microsoft.semantickernel.data.vectorstorage.options.UpsertRecordOptions;
+import com.microsoft.semantickernel.exceptions.SKException;
+import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
+
+import javax.annotation.Nonnull;
+import javax.sql.DataSource;
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.SQLException;
+import java.util.List;
+import java.util.stream.Collectors;
+
+public class SQLiteVectorStoreQueryProvider extends
+    JDBCVectorStoreQueryProvider implements SQLVectorStoreQueryProvider {
+
+    private final DataSource dataSource;
+    private final ObjectMapper objectMapper;
+
+    @SuppressFBWarnings("EI_EXPOSE_REP2")
+    private SQLiteVectorStoreQueryProvider(
+        @Nonnull DataSource dataSource,
+        @Nonnull String collectionsTable,
+        @Nonnull String prefixForCollectionTables,
+        @Nonnull ObjectMapper objectMapper) {
+        super(dataSource, collectionsTable, prefixForCollectionTables);
+        this.dataSource = dataSource;
+        this.objectMapper = objectMapper;
+    }
+
+    /**
+     * Creates a new builder.
+     * @return the builder
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    private void setUpsertStatementValues(PreparedStatement statement, Object record,
+        List<VectorStoreRecordField> fields) {
+        JsonNode jsonNode = objectMapper.valueToTree(record);
+
+        for (int i = 0; i < fields.size(); ++i) {
+            VectorStoreRecordField field = fields.get(i);
+            try {
+                JsonNode valueNode = jsonNode.get(field.getEffectiveStorageName());
+
+                if (field instanceof VectorStoreRecordVectorField) {
+                    // Convert the vector field to a string
+                    if (!field.getFieldType().equals(String.class)) {
+                        statement.setObject(i + 1, objectMapper.writeValueAsString(valueNode));
+                        continue;
+                    }
+                }
+
+                statement.setObject(i + 1,
+                    objectMapper.convertValue(valueNode, field.getFieldType()));
+            } catch (SQLException | JsonProcessingException e) {
+                throw new RuntimeException(e);
+            }
+        }
+    }
+
+    /**
+     * Upserts records into the collection.
+     * @param collectionName the collection name
+     * @param records the records to upsert
+     * @param recordDefinition the record definition
+     * @param options the upsert options
+     * @throws SKException if the upsert fails
+     */
+    @Override
+    @SuppressFBWarnings("SQL_PREPARED_STATEMENT_GENERATED_FROM_NONCONSTANT_STRING") // SQL query is generated dynamically with valid identifiers
+    public void upsertRecords(String collectionName, List<?> records,
+        VectorStoreRecordDefinition recordDefinition, UpsertRecordOptions options) {
+        List<VectorStoreRecordField> fields = recordDefinition.getAllFields();
+
+        String query = formatQuery("INSERT OR REPLACE INTO %s (%s) VALUES (%s)",
+            getCollectionTableName(collectionName),
+            getQueryColumnsFromFields(fields),
+            getWildcardString(fields.size()));
+
+        try (Connection connection = dataSource.getConnection();
+            PreparedStatement statement = connection.prepareStatement(query)) {
+            for (Object record : records) {
+                setUpsertStatementValues(statement, record, recordDefinition.getAllFields());
+                statement.addBatch();
+            }
+
+            statement.executeBatch();
+        } catch (SQLException e) {
+            throw new SKException("Failed to upsert records", e);
+        }
+    }
+
+    public static class Builder
+        extends JDBCVectorStoreQueryProvider.Builder {
+        private DataSource dataSource;
+        private String collectionsTable = DEFAULT_COLLECTIONS_TABLE;
+        private String prefixForCollectionTables = DEFAULT_PREFIX_FOR_COLLECTION_TABLES;
+        private ObjectMapper objectMapper = new ObjectMapper();
+
+        @SuppressFBWarnings("EI_EXPOSE_REP2")
+        public Builder withDataSource(DataSource dataSource) {
+            this.dataSource = dataSource;
+            return this;
+        }
+
+        /**
+         * Sets the collections table name.
+         * @param collectionsTable the collections table name
+         * @return the builder
+         */
+        public Builder withCollectionsTable(String collectionsTable) {
+            this.collectionsTable = validateSQLidentifier(collectionsTable);
+            return this;
+        }
+
+        /**
+         * Sets the prefix for collection tables.
+         * @param prefixForCollectionTables the prefix for collection tables
+         * @return the builder
+         */
+        public Builder withPrefixForCollectionTables(String prefixForCollectionTables) {
+            this.prefixForCollectionTables = validateSQLidentifier(prefixForCollectionTables);
+            return this;
+        }
+
+        /**
+         * Sets the object mapper.
+         * @param objectMapper the object mapper
+         * @return the builder
+         */
+        @SuppressFBWarnings("EI_EXPOSE_REP2")
+        public Builder withObjectMapper(ObjectMapper objectMapper) {
+            this.objectMapper = objectMapper;
+            return this;
+        }
+
+        public SQLiteVectorStoreQueryProvider build() {
+            if (dataSource == null) {
+                throw new SKException("DataSource is required");
+            }
+
+            return new SQLiteVectorStoreQueryProvider(dataSource, collectionsTable,
+                prefixForCollectionTables, objectMapper);
+        }
+    }
+}