Merge pull request #226 from milderhc/clean-interfaces

Add VolatileVectorStoreRecordCollection search and VectorSearch interfaces refactor

Author: Milder Hernandez

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

@@ -100,9 +100,7 @@ public void buildRecordCollection(QueryProvider provider) {
     }
 
     private List<Hotel> getHotels() {
-        ArrayList<Hotel> embeddings = new ArrayList<>();
-
-        return List.of(
+        return Arrays.asList(
                 new Hotel("id_1", "Hotel 1", 1, "Hotel 1 description", Arrays.asList(0.5f, 3.2f, 7.1f, -4.0f, 2.8f, 10.0f, -1.3f, 5.5f),null, null, null, 4.0),
                 new Hotel("id_2", "Hotel 2", 2, "Hotel 2 description", Arrays.asList(-2.0f, 8.1f, 0.9f, 5.4f, -3.3f, 2.2f, 9.9f, -4.5f),null, null, null, 4.0),
                 new Hotel("id_3", "Hotel 3", 3, "Hotel 3 description", Arrays.asList(4.5f, -6.2f, 3.1f, 7.7f, -0.8f, 1.1f, -2.2f, 8.3f),null, null, null, 5.0),

api-test/integration-tests/src/test/java/com/microsoft/semantickernel/tests/connectors/memory/redis/RedisHashSetVectorStoreRecordCollectionTest.java

@@ -132,7 +132,7 @@ private static RedisHashSetVectorStoreRecordCollection<Hotel> createCollection(@
     }
 
     private static List<Hotel> getHotels() {
-        return List.of(
+        return Arrays.asList(
                 new Hotel("id_1", "Hotel 1", 1, "Hotel 1 description", Arrays.asList(0.5f, 3.2f, 7.1f, -4.0f, 2.8f, 10.0f, -1.3f, 5.5f),null, null, 4.0),
                 new Hotel("id_2", "Hotel 2", 2, "Hotel 2 description", Arrays.asList(-2.0f, 8.1f, 0.9f, 5.4f, -3.3f, 2.2f, 9.9f, -4.5f),null, null, 4.0),
                 new Hotel("id_3", "Hotel 3", 3, "Hotel 3 description", Arrays.asList(4.5f, -6.2f, 3.1f, 7.7f, -0.8f, 1.1f, -2.2f, 8.3f),null, null, 5.0),

api-test/integration-tests/src/test/java/com/microsoft/semantickernel/tests/connectors/memory/redis/RedisJsonVectorStoreRecordCollectionTest.java

@@ -132,7 +132,7 @@ private static RedisJsonVectorStoreRecordCollection<Hotel> createCollection(@Non
     }
 
     private static List<Hotel> getHotels() {
-        return List.of(
+        return Arrays.asList(
                 new Hotel("id_1", "Hotel 1", 1, "Hotel 1 description", Arrays.asList(0.5f, 3.2f, 7.1f, -4.0f, 2.8f, 10.0f, -1.3f, 5.5f),null, null, 4.0),
                 new Hotel("id_2", "Hotel 2", 2, "Hotel 2 description", Arrays.asList(-2.0f, 8.1f, 0.9f, 5.4f, -3.3f, 2.2f, 9.9f, -4.5f),null, null, 4.0),
                 new Hotel("id_3", "Hotel 3", 3, "Hotel 3 description", Arrays.asList(4.5f, -6.2f, 3.1f, 7.7f, -0.8f, 1.1f, -2.2f, 8.3f),null, null, 5.0),

semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/azureaisearch/AzureAISearchVectorStoreRecordCollection.java

@@ -14,7 +14,7 @@
 import com.azure.search.documents.models.VectorQuery;
 import com.azure.search.documents.models.VectorizableTextQuery;
 import com.azure.search.documents.models.VectorizedQuery;
-import com.microsoft.semantickernel.data.vectorsearch.VectorizableSearch;
+import com.microsoft.semantickernel.data.vectorsearch.VectorizableTextSearch;
 import com.microsoft.semantickernel.data.vectorsearch.VectorSearchResult;
 import com.microsoft.semantickernel.data.vectorsearch.VectorizedSearch;
 import com.microsoft.semantickernel.data.vectorstorage.VectorStoreRecordCollection;
@@ -28,9 +28,6 @@
 import com.microsoft.semantickernel.data.vectorstorage.options.GetRecordOptions;
 import com.microsoft.semantickernel.data.vectorstorage.options.UpsertRecordOptions;
 import com.microsoft.semantickernel.data.vectorstorage.options.VectorSearchOptions;
-import com.microsoft.semantickernel.data.vectorsearch.queries.VectorSearchQuery;
-import com.microsoft.semantickernel.data.vectorsearch.queries.VectorizableTextSearchQuery;
-import com.microsoft.semantickernel.data.vectorsearch.queries.VectorizedSearchQuery;
 import com.microsoft.semantickernel.exceptions.SKException;
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.time.OffsetDateTime;
@@ -49,7 +46,7 @@
 public class AzureAISearchVectorStoreRecordCollection<Record> implements
     VectorStoreRecordCollection<String, Record>,
     VectorizedSearch<Record>,
-    VectorizableSearch<Record> {
+    VectorizableTextSearch<Record> {
 
     private static final HashSet<Class<?>> supportedKeyTypes = new HashSet<>(
         Collections.singletonList(
@@ -277,8 +274,25 @@ public Mono<Void> deleteBatchAsync(List<String> keys, DeleteRecordOptions option
         }).collect(Collectors.toList())).then();
     }
 
-    private Mono<List<VectorSearchResult<Record>>> searchAndMapAsync(SearchOptions searchOptions,
+    private Mono<List<VectorSearchResult<Record>>> searchAndMapAsync(
+        List<VectorQuery> vectorQueries, VectorSearchOptions options,
         GetRecordOptions getRecordOptions) {
+
+        String filter = AzureAISearchVectorStoreCollectionSearchMapping
+            .buildFilterString(options.getVectorSearchFilter(), recordDefinition);
+
+        SearchOptions searchOptions = new SearchOptions()
+            .setFilter(filter)
+            .setTop(options.getLimit())
+            .setSkip(options.getOffset())
+            .setScoringParameters()
+            .setVectorSearchOptions(new com.azure.search.documents.models.VectorSearchOptions()
+                .setQueries(vectorQueries));
+
+        if (!options.isIncludeVectors()) {
+            searchOptions.setSelect(nonVectorFields.toArray(new String[0]));
+        }
+
         VectorStoreRecordMapper<Record, SearchDocument> mapper = this.options
             .getVectorStoreRecordMapper();
 
@@ -300,70 +314,32 @@ record = response.getDocument(this.options.getRecordClass());
     }
 
     /**
-     * Search the vector store for records that match the given embedding and filter.
+     * Vectorizable text search. This method searches for records that are similar to the given text.
      *
-     * @param query The vector search query.
+     * @param searchText The text to search with.
+     * @param options    The options to use for the search.
      * @return A list of search results.
      */
     @Override
-    public Mono<List<VectorSearchResult<Record>>> searchAsync(VectorSearchQuery query) {
+    public Mono<List<VectorSearchResult<Record>>> searchAsync(String searchText,
+        VectorSearchOptions options) {
         if (firstVectorFieldName == null) {
             throw new SKException("No vector fields defined. Cannot perform vector search");
         }
 
-        VectorSearchOptions options = query.getSearchOptions();
         if (options == null) {
             options = VectorSearchOptions.createDefault(firstVectorFieldName);
         }
 
         List<VectorQuery> vectorQueries = new ArrayList<>();
-
-        if (query instanceof VectorizedSearchQuery) {
-            vectorQueries.add(new VectorizedQuery(((VectorizedSearchQuery) query).getVector())
-                .setFields(recordDefinition.getField(options.getVectorFieldName() != null
-                    ? options.getVectorFieldName()
-                    : firstVectorFieldName).getEffectiveStorageName())
-                .setKNearestNeighborsCount(options.getLimit()));
-        } else if (query instanceof VectorizableTextSearchQuery) {
-            vectorQueries
-                .add(new VectorizableTextQuery(((VectorizableTextSearchQuery) query).getQueryText())
-                    .setFields(recordDefinition.getField(options.getVectorFieldName() != null
-                        ? options.getVectorFieldName()
-                        : firstVectorFieldName).getEffectiveStorageName())
-                    .setKNearestNeighborsCount(options.getLimit()));
-        } else {
-            throw new SKException("Unsupported query type: " + query.getQueryType());
-        }
-
-        String filter = AzureAISearchVectorStoreCollectionSearchMapping
-            .buildFilterString(options.getVectorSearchFilter(), recordDefinition);
-
-        SearchOptions searchOptions = new SearchOptions()
-            .setFilter(filter)
-            .setTop(options.getLimit())
-            .setSkip(options.getOffset())
-            .setScoringParameters()
-            .setVectorSearchOptions(new com.azure.search.documents.models.VectorSearchOptions()
-                .setQueries(vectorQueries));
-
-        if (!options.isIncludeVectors()) {
-            searchOptions.setSelect(nonVectorFields.toArray(new String[0]));
-        }
-
-        return searchAndMapAsync(searchOptions, new GetRecordOptions(options.isIncludeVectors()));
-    }
-
-    /**
-     * Vectorizable text search. This method searches for records that are similar to the given text.
-     *
-     * @param searchText The text to search with.
-     * @param options    The options to use for the search.
-     * @return A list of search results.
-     */
-    @Override
-    public Mono<List<VectorSearchResult<Record>>> searchAsync(String searchText,
-        VectorSearchOptions options) {
-        return searchAsync(VectorSearchQuery.createQuery(searchText, options));
+        vectorQueries.add(new VectorizableTextQuery(searchText)
+            .setFields(recordDefinition.getField(options.getVectorFieldName() != null
+                ? options.getVectorFieldName()
+                : firstVectorFieldName).getEffectiveStorageName())
+            .setKNearestNeighborsCount(options.getLimit()));
+
+        return searchAndMapAsync(vectorQueries, options,
+            new GetRecordOptions(options.isIncludeVectors()));
     }
 
     /**
@@ -376,6 +352,22 @@ public Mono<List<VectorSearchResult<Record>>> searchAsync(String searchText,
     @Override
     public Mono<List<VectorSearchResult<Record>>> searchAsync(List<Float> vector,
         VectorSearchOptions options) {
-        return searchAsync(VectorSearchQuery.createQuery(vector, options));
+        if (firstVectorFieldName == null) {
+            throw new SKException("No vector fields defined. Cannot perform vector search");
+        }
+
+        if (options == null) {
+            options = VectorSearchOptions.createDefault(firstVectorFieldName);
+        }
+
+        List<VectorQuery> vectorQueries = new ArrayList<>();
+        vectorQueries.add(new VectorizedQuery(vector)
+            .setFields(recordDefinition.getField(options.getVectorFieldName() != null
+                ? options.getVectorFieldName()
+                : firstVectorFieldName).getEffectiveStorageName())
+            .setKNearestNeighborsCount(options.getLimit()));
+
+        return searchAndMapAsync(vectorQueries, options,
+            new GetRecordOptions(options.isIncludeVectors()));
     }
 }

semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/JDBCVectorStoreQueryProvider.java

@@ -6,8 +6,6 @@
 import com.fasterxml.jackson.databind.node.ArrayNode;
 import com.microsoft.semantickernel.data.vectorsearch.VectorOperations;
 import com.microsoft.semantickernel.data.vectorsearch.VectorSearchResult;
-import com.microsoft.semantickernel.data.vectorsearch.queries.VectorSearchQuery;
-import com.microsoft.semantickernel.data.vectorsearch.queries.VectorizedSearchQuery;
 import com.microsoft.semantickernel.data.vectorstorage.VectorStoreRecordMapper;
 import com.microsoft.semantickernel.data.vectorstorage.definition.DistanceFunction;
 import com.microsoft.semantickernel.data.vectorstorage.definition.VectorStoreRecordDefinition;
@@ -471,101 +469,46 @@ protected <Record> List<Record> getRecordsWithFilter(String collectionName,
      *
      * @param <Record> the record type
      * @param collectionName the collection name
-     * @param query the vectorized search query, containing the vector and search options
+     * @param vector the vector to search with
+     * @param options the search options
      * @param recordDefinition the record definition
      * @param mapper the mapper, responsible for mapping the result set to the record type.
      * @return the search results
      */
     @Override
     public <Record> List<VectorSearchResult<Record>> search(String collectionName,
-        VectorSearchQuery query, VectorStoreRecordDefinition recordDefinition,
+        List<Float> vector, VectorSearchOptions options,
+        VectorStoreRecordDefinition recordDefinition,
         VectorStoreRecordMapper<Record, ResultSet> mapper) {
         if (recordDefinition.getVectorFields().isEmpty()) {
             throw new SKException("No vector fields defined. Cannot perform vector search");
         }
 
-        if (query instanceof VectorizedSearchQuery) {
-            VectorizedSearchQuery vectorizedSearchQuery = (VectorizedSearchQuery) query;
-            VectorSearchOptions options = query.getSearchOptions();
+        VectorStoreRecordVectorField firstVectorField = recordDefinition.getVectorFields()
+            .get(0);
+        if (options == null) {
+            options = VectorSearchOptions.createDefault(firstVectorField.getName());
+        }
 
-            VectorStoreRecordVectorField firstVectorField = recordDefinition.getVectorFields()
-                .get(0);
-            if (options == null) {
-                options = VectorSearchOptions.createDefault(firstVectorField.getName());
-            }
+        VectorStoreRecordVectorField vectorField = options.getVectorFieldName() == null
+            ? firstVectorField
+            : (VectorStoreRecordVectorField) recordDefinition
+                .getField(options.getVectorFieldName());
 
-            VectorStoreRecordVectorField vectorField = options.getVectorFieldName() == null
-                ? firstVectorField
-                : (VectorStoreRecordVectorField) recordDefinition
-                    .getField(options.getVectorFieldName());
-
-            String filter = SQLVectorStoreRecordCollectionSearchMapping
-                .buildFilter(options.getVectorSearchFilter(), recordDefinition);
-            List<Object> parameters = SQLVectorStoreRecordCollectionSearchMapping
-                .getFilterParameters(options.getVectorSearchFilter());
-
-            List<Record> records = getRecordsWithFilter(collectionName, recordDefinition, mapper,
-                new GetRecordOptions(true), filter, parameters);
-            List<VectorSearchResult<Record>> results = new ArrayList<>();
-
-            DistanceFunction distanceFunction = vectorField.getDistanceFunction() == null
-                ? DistanceFunction.EUCLIDEAN_DISTANCE
-                : vectorField.getDistanceFunction();
-
-            for (Record record : records) {
-                List<Float> vector;
-                try {
-                    String json = new ObjectMapper().writeValueAsString(record);
-                    ArrayNode arrayNode = (ArrayNode) new ObjectMapper().readTree(json)
-                        .get(vectorField.getEffectiveStorageName());
-
-                    vector = Stream.iterate(0, i -> i + 1)
-                        .limit(arrayNode.size())
-                        .map(i -> arrayNode.get(i).floatValue())
-                        .collect(Collectors.toList());
-                } catch (JsonProcessingException e) {
-                    throw new RuntimeException(e);
-                }
+        String filter = SQLVectorStoreRecordCollectionSearchMapping
+            .buildFilter(options.getVectorSearchFilter(), recordDefinition);
+        List<Object> parameters = SQLVectorStoreRecordCollectionSearchMapping
+            .getFilterParameters(options.getVectorSearchFilter());
 
-                double score;
-                switch (distanceFunction) {
-                    case COSINE_SIMILARITY:
-                        score = VectorOperations.cosineSimilarity(vectorizedSearchQuery.getVector(),
-                            vector);
-                        break;
-                    case COSINE_DISTANCE:
-                        score = VectorOperations.cosineDistance(vectorizedSearchQuery.getVector(),
-                            vector);
-                        break;
-                    case EUCLIDEAN_DISTANCE:
-                        score = VectorOperations
-                            .euclideanDistance(vectorizedSearchQuery.getVector(), vector);
-                        break;
-                    case DOT_PRODUCT:
-                        score = VectorOperations.dot(vectorizedSearchQuery.getVector(), vector);
-                        break;
-                    default:
-                        throw new SKException("Unsupported distance function");
-                }
+        List<Record> records = getRecordsWithFilter(collectionName, recordDefinition, mapper,
+            new GetRecordOptions(true), filter, parameters);
 
-                results.add(new VectorSearchResult<>(record, score));
-            }
-
-            Comparator<VectorSearchResult<Record>> comparator = Comparator
-                .comparingDouble(VectorSearchResult::getScore);
-            // Higher scores are better
-            if (distanceFunction == DistanceFunction.COSINE_SIMILARITY
-                || distanceFunction == DistanceFunction.DOT_PRODUCT) {
-                comparator = comparator.reversed();
-            }
-            return results.stream()
-                .sorted(comparator)
-                .skip(options.getOffset())
-                .limit(options.getLimit())
-                .collect(Collectors.toList());
-        }
+        DistanceFunction distanceFunction = vectorField.getDistanceFunction() == null
+            ? DistanceFunction.EUCLIDEAN_DISTANCE
+            : vectorField.getDistanceFunction();
 
-        throw new SKException("Unsupported query type");
+        return VectorOperations.exactSimilaritySearch(records, vector, vectorField,
+            distanceFunction, options);
     }
 
     /**

semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/JDBCVectorStoreRecordCollection.java

@@ -7,7 +7,6 @@
 import com.microsoft.semantickernel.connectors.data.postgres.PostgreSQLVectorStoreRecordMapper;
 import com.microsoft.semantickernel.data.vectorsearch.VectorSearchResult;
 import com.microsoft.semantickernel.data.vectorsearch.VectorizedSearch;
-import com.microsoft.semantickernel.data.vectorsearch.queries.VectorSearchQuery;
 import com.microsoft.semantickernel.data.vectorstorage.VectorStoreRecordMapper;
 import com.microsoft.semantickernel.data.vectorstorage.VectorStoreRecordCollection;
 import com.microsoft.semantickernel.data.vectorstorage.definition.VectorStoreRecordDefinition;
@@ -298,14 +297,6 @@ public Mono<Void> prepareAsync() {
             .subscribeOn(Schedulers.boundedElastic()).then();
     }
 
-    @Override
-    public Mono<List<VectorSearchResult<Record>>> searchAsync(VectorSearchQuery query) {
-        return Mono.fromCallable(
-            () -> queryProvider.search(this.collectionName, query, recordDefinition,
-                vectorStoreRecordMapper))
-            .subscribeOn(Schedulers.boundedElastic());
-    }
-
     /**
      * Vectorized search. This method searches for records that are similar to the given vector.
      *
@@ -316,7 +307,11 @@ public Mono<List<VectorSearchResult<Record>>> searchAsync(VectorSearchQuery quer
     @Override
     public Mono<List<VectorSearchResult<Record>>> searchAsync(List<Float> vector,
         VectorSearchOptions vectorSearchOptions) {
-        return this.searchAsync(VectorSearchQuery.createQuery(vector, vectorSearchOptions));
+        return Mono.fromCallable(
+            () -> queryProvider.search(this.collectionName, vector, vectorSearchOptions,
+                recordDefinition,
+                vectorStoreRecordMapper))
+            .subscribeOn(Schedulers.boundedElastic());
     }
 
     public static class Builder<Record>

semantickernel-experimental/src/main/java/com/microsoft/semantickernel/connectors/data/jdbc/SQLVectorStoreQueryProvider.java

@@ -3,12 +3,12 @@
 
 import com.microsoft.semantickernel.builders.SemanticKernelBuilder;
 import com.microsoft.semantickernel.data.vectorsearch.VectorSearchResult;
-import com.microsoft.semantickernel.data.vectorsearch.queries.VectorSearchQuery;
 import com.microsoft.semantickernel.data.vectorstorage.VectorStoreRecordMapper;
 import com.microsoft.semantickernel.data.vectorstorage.definition.VectorStoreRecordDefinition;
 import com.microsoft.semantickernel.data.vectorstorage.options.DeleteRecordOptions;
 import com.microsoft.semantickernel.data.vectorstorage.options.GetRecordOptions;
 import com.microsoft.semantickernel.data.vectorstorage.options.UpsertRecordOptions;
+import com.microsoft.semantickernel.data.vectorstorage.options.VectorSearchOptions;
 
 import java.sql.ResultSet;
 import java.util.List;
@@ -138,13 +138,15 @@ void deleteRecords(String collectionName, List<String> keys,
      *
      * @param <Record> the record type
      * @param collectionName the collection name
-     * @param query the vectorized search query, containing the vector and search options
+     * @param vector the vector to search with
+     * @param options the vector search options
      * @param recordDefinition the record definition
      * @param mapper the mapper, responsible for mapping the result set to the record type.
      * @return the search results
      */
     <Record> List<VectorSearchResult<Record>> search(String collectionName,
-        VectorSearchQuery query,
+        List<Float> vector,
+        VectorSearchOptions options,
         VectorStoreRecordDefinition recordDefinition,
         VectorStoreRecordMapper<Record, ResultSet> mapper);