feat(vector-stores): Support filtering in VertexAI Matching Engine (#136

)

https://cloud.google.com/vertex-ai/docs/matching-engine/filtering

packages/langchain/lib/src/documents/vector_stores/models/models.dart

@@ -1,34 +1,54 @@
 /// {@template vector_store_search_type}
 /// Vector store search type.
+///
+/// In general, we provide support two types of search:
+/// - Similarity search.
+/// - Maximal Marginal Relevance (MMR) search.
+///
+/// But it depends on the actual implementation of the vector store whether
+/// these are supported or not. Vector stores may also provide their own
+/// subclasses of this class to support additional configuration options.
+/// For example, [VertexAIMatchingEngine] provides
+/// [VertexAIMatchingEngineSimilaritySearch] which is a subclass of
+/// [VectorStoreSimilaritySearch]. Check the documentation of the vector store
+/// you are using for more information.
 /// {@endtemplate}
 sealed class VectorStoreSearchType {
   /// {@macro vector_store_search_type}
   const VectorStoreSearchType({
     required this.k,
+    this.filter,
   });
 
   /// The number of documents to return.
   final int k;
 
+  /// The filter to apply to the search.
+  final Map<String, dynamic>? filter;
+
   /// Similarity search.
   factory VectorStoreSearchType.similarity({
     final int k = 4,
+    final Map<String, dynamic>? filter,
     final double? scoreThreshold,
   }) {
     return VectorStoreSimilaritySearch(
       k: k,
+      filter: filter,
       scoreThreshold: scoreThreshold,
     );
   }
 
   /// Maximal Marginal Relevance (MMR) search.
   factory VectorStoreSearchType.mmr({
     final int k = 4,
+    final Map<String, dynamic>? filter,
     final int fetchK = 20,
     final double lambdaMult = 0.5,
   }) {
     return VectorStoreMMRSearch(
       k: k,
+      filter: filter,
       fetchK: fetchK,
       lambdaMult: lambdaMult,
     );
@@ -43,6 +63,7 @@ class VectorStoreSimilaritySearch extends VectorStoreSearchType {
   /// {@macro vector_store_similarity_search}
   const VectorStoreSimilaritySearch({
     super.k = 4,
+    super.filter,
     this.scoreThreshold,
   });
 
@@ -61,6 +82,7 @@ class VectorStoreMMRSearch extends VectorStoreSearchType {
   /// {@macro vector_store_mmr_search}
   const VectorStoreMMRSearch({
     super.k = 4,
+    super.filter,
     this.fetchK = 20,
     this.lambdaMult = 0.5,
   });

packages/langchain_google/lib/src/vector_stores/matching_engine.dart

@@ -6,6 +6,8 @@ import 'package:langchain/langchain.dart';
 import 'package:uuid/uuid.dart';
 import 'package:vertex_ai/vertex_ai.dart';
 
+import 'models/models.dart';
+
 /// A vector store that uses Vertex AI Matching Engine.
 ///
 /// Vertex AI Matching Engine provides a high-scale low latency vector database.
@@ -95,7 +97,7 @@ import 'package:vertex_ai/vertex_ai.dart';
 ///
 /// You can use the constant `VertexAIMatchingEngine.cloudPlatformScopes`.
 ///
-/// ### Vector attributes
+/// ### Vector attributes filtering
 ///
 /// Vertex AI Matching Engine allows you to add attributes to the vectors that
 /// you can later use to restrict vector matching searches to a subset of the
@@ -125,6 +127,26 @@ import 'package:vertex_ai/vertex_ai.dart';
 ///
 /// Check out the documentation for more details:
 /// https://cloud.google.com/vertex-ai/docs/matching-engine/filtering
+///
+/// After adding the attributes to the documents, you can use the use them to
+/// restrict the similarity search results. Example:
+///
+/// ```dart
+/// final vectorStore = VertexAIMatchingEngine(...);
+/// final res = await vectorStore.similaritySearch(
+///   query: 'What should I feed my cat?',
+///   config: VertexAIMatchingEngineSimilaritySearch(
+///     k: 5,
+///     scoreThreshold: 0.8,
+///     filters: [
+///       const VertexAIMatchingEngineFilter(
+///         namespace: 'class',
+///         allowList: ['cat'],
+///       ),
+///     ],
+///   ),
+/// );
+/// ```
 class VertexAIMatchingEngine extends VectorStore {
   /// Creates a new Vertex AI Matching Engine vector store.
   ///
@@ -321,6 +343,8 @@ class VertexAIMatchingEngine extends VectorStore {
           datapoint: VertexAIIndexDatapoint(
             datapointId: _uuid.v4(),
             featureVector: embedding,
+            restricts: config
+                .filter?[VertexAIMatchingEngineSimilaritySearch.filterKey],
           ),
           neighborCount: config.k,
         ),

packages/langchain_google/lib/src/vector_stores/models/mappers.dart

@@ -0,0 +1,15 @@
+import 'package:vertex_ai/vertex_ai.dart';
+
+import 'models.dart';
+
+abstract class VertexAIMatchingEngineFilterMapper {
+  static VertexAIIndexDatapointRestriction toDto(
+    final VertexAIMatchingEngineFilter filter,
+  ) {
+    return VertexAIIndexDatapointRestriction(
+      namespace: filter.namespace,
+      allowList: filter.allowList,
+      denyList: filter.denyList,
+    );
+  }
+}

packages/langchain_google/lib/src/vector_stores/models/models.dart

@@ -0,0 +1,74 @@
+import 'package:langchain/langchain.dart';
+import 'package:meta/meta.dart';
+
+import 'mappers.dart';
+
+/// {@template vector_store_similarity_search}
+/// Vertex AI Matching Engine similarity search config.
+///
+/// Example:
+/// ```dart
+/// VertexAIMatchingEngineSimilaritySearch(
+///   k: 5,
+///   filters: [
+///     const VertexAIMatchingEngineFilter(
+///       namespace: 'class',
+///       allowList: ['pet'],
+///     ),
+///     const VertexAIMatchingEngineFilter(
+///       namespace: 'category',
+///       denyList: ['canine'],
+///     ),
+///   ]
+/// ),
+/// ```
+/// {@endtemplate}
+class VertexAIMatchingEngineSimilaritySearch
+    extends VectorStoreSimilaritySearch {
+  VertexAIMatchingEngineSimilaritySearch({
+    super.k = 4,
+    final List<VertexAIMatchingEngineFilter>? filters,
+    super.scoreThreshold,
+  }) : super(
+          filter: filters != null
+              ? {
+                  filterKey: filters
+                      .map(
+                        VertexAIMatchingEngineFilterMapper.toDto,
+                      )
+                      .toList(growable: false),
+                }
+              : null,
+        );
+
+  static const filterKey = 'restricts';
+}
+
+/// {@template vertex_ai_matching_engine_filter}
+/// Filter for the Vertex AI Matching Engine.
+/// See: https://cloud.google.com/vertex-ai/docs/matching-engine/filtering
+/// {@endtemplate}
+@immutable
+class VertexAIMatchingEngineFilter {
+  /// {@macro vertex_ai_matching_engine_filter}
+  const VertexAIMatchingEngineFilter({
+    required this.namespace,
+    this.allowList = const [],
+    this.denyList = const [],
+  });
+
+  /// The namespace of this restriction.
+  ///
+  /// eg: color.
+  final String namespace;
+
+  /// The attributes to allow in this namespace.
+  ///
+  /// eg: 'red'
+  final List<String> allowList;
+
+  /// The attributes to deny in this namespace.
+  ///
+  /// eg: 'blue'
+  final List<String> denyList;
+}

packages/langchain_google/lib/src/vector_stores/vector_stores.dart

@@ -1 +1,2 @@
 export 'matching_engine.dart';
+export 'models/models.dart';

packages/langchain_google/test/vector_stores/matching_engine.dart

@@ -59,11 +59,29 @@ void main() async {
     test('Test VertexAIMatchingEngine query with scoreThreshold', () async {
       final res = await vectorStore.similaritySearchWithScores(
         query: 'Can I pay by credit card?',
-        config: const VectorStoreSimilaritySearch(scoreThreshold: 0.6),
+        config: VertexAIMatchingEngineSimilaritySearch(scoreThreshold: 0.6),
       );
       for (final (_, score) in res) {
         expect(score, greaterThan(0.6));
       }
     });
+
+    test('Test VertexAIMatchingEngine query with filters', () async {
+      final res = await vectorStore.similaritySearch(
+        query: 'Can I pay by credit card?',
+        config: VertexAIMatchingEngineSimilaritySearch(
+          k: 10,
+          filters: [
+            const VertexAIMatchingEngineFilter(
+              namespace: 'type',
+              allowList: ['faq'],
+            ),
+          ],
+        ),
+      );
+      for (final doc in res) {
+        expect(doc.metadata['type'], 'faq');
+      }
+    });
   });
 }