refactor: deprecated ExactMatchFilter and use MetadataFilter with man…

…y operator support
run-llama · Jul 11, 2024 · 12d80c6 · 12d80c6
1 parent 4205967
commit 12d80c6
Show file tree

Hide file tree

Showing 6 changed files with 66 additions and 20 deletions.
diff --git a/examples/metadata-filter/simple.ts b/examples/metadata-filter/simple.ts
@@ -74,12 +74,12 @@ async function main() {
         {
           key: "private",
           value: "false",
-          filterType: "ExactMatch",
+          operator: "==",
         },
         {
           key: "dogId",
           value: "3",
-          filterType: "ExactMatch",
+          operator: "==",
         },
       ],
     },

diff --git a/packages/llamaindex/src/storage/vectorStore/PGVectorStore.ts b/packages/llamaindex/src/storage/vectorStore/PGVectorStore.ts
@@ -272,7 +272,10 @@ export class PGVectorStore
     query.filters?.filters.forEach((filter, index) => {
       const paramIndex = params.length + 1;
       whereClauses.push(`metadata->>'${filter.key}' = $${paramIndex}`);
-      params.push(filter.value);
+      // TODO: support filter with other operators
+      if (!Array.isArray(filter.value)) {
+        params.push(filter.value);
+      }
     });
 
     const where =

diff --git a/packages/llamaindex/src/storage/vectorStore/PineconeVectorStore.ts b/packages/llamaindex/src/storage/vectorStore/PineconeVectorStore.ts
@@ -1,7 +1,7 @@
 import {
   VectorStoreBase,
-  type ExactMatchFilter,
   type IEmbedModel,
+  type MetadataFilter,
   type MetadataFilters,
   type VectorStoreNoEmbedModel,
   type VectorStoreQuery,
@@ -199,8 +199,12 @@ export class PineconeVectorStore
   }
 
   toPineconeFilter(stdFilters?: MetadataFilters) {
-    return stdFilters?.filters?.reduce((carry: any, item: ExactMatchFilter) => {
-      carry[item.key] = item.value;
+    return stdFilters?.filters?.reduce((carry: any, item: MetadataFilter) => {
+      // Use MetadataFilter with EQ operator to replace ExactMatchFilter
+      // TODO: support filter with other operators
+      if (item.operator === "==") {
+        carry[item.key] = item.value;
+      }
       return carry;
     }, {});
   }

diff --git a/packages/llamaindex/src/storage/vectorStore/SimpleVectorStore.ts b/packages/llamaindex/src/storage/vectorStore/SimpleVectorStore.ts
@@ -9,9 +9,11 @@ import {
 import { exists } from "../FileSystem.js";
 import { DEFAULT_PERSIST_DIR } from "../constants.js";
 import {
+  FilterOperator,
   VectorStoreBase,
   VectorStoreQueryMode,
   type IEmbedModel,
+  type MetadataFilter,
   type VectorStoreNoEmbedModel,
   type VectorStoreQuery,
   type VectorStoreQueryResult,
@@ -26,10 +28,12 @@ const LEARNER_MODES = new Set<VectorStoreQueryMode>([
 
 const MMR_MODE = VectorStoreQueryMode.MMR;
 
+type MetadataValue = Record<string, any>;
+
 class SimpleVectorStoreData {
   embeddingDict: Record<string, number[]> = {};
   textIdToRefDocId: Record<string, string> = {};
-  metadataDict: Record<string, Record<string, any>> = {};
+  metadataDict: Record<string, MetadataValue> = {};
 }
 
 export class SimpleVectorStore
@@ -105,27 +109,30 @@ export class SimpleVectorStore
   }> {
     const items = Object.entries(this.data.embeddingDict);
 
-    const metadataLookup = {
-      ExactMatch: (
-        metadata: Record<string, any>,
-        key: string,
-        value: string | number,
-      ) => {
-        return String(metadata[key]) === value.toString(); // compare as string
+    const operatorToFilterFn: Record<
+      FilterOperator,
+      (
+        input: Omit<MetadataFilter, "operator"> & {
+          metadata: MetadataValue;
+        },
+      ) => boolean
+    > = {
+      "==": (input) => {
+        return String(input.metadata[input.key]) === input.value.toString(); // compare as string
       },
     };
 
     const queryFilterFn = (nodeId: string) => {
       if (!query.filters) return true;
       const filters = query.filters.filters;
       for (const filter of filters) {
-        const { key, value, filterType } = filter;
-        const metadataLookupFn = metadataLookup[filterType];
+        const { key, value, operator } = filter;
+        const metadataLookupFn = operatorToFilterFn[operator];
         const metadata = this.data.metadataDict[nodeId];
         const isMatch =
           metadataLookupFn &&
           metadata &&
-          metadataLookupFn(metadata, key, value);
+          metadataLookupFn({ metadata, key, value });
         if (!isMatch) return false; // TODO: handle condition OR AND
       }
       return true;

diff --git a/packages/llamaindex/src/storage/vectorStore/types.ts b/packages/llamaindex/src/storage/vectorStore/types.ts
@@ -20,19 +20,51 @@ export enum VectorStoreQueryMode {
   MMR = "mmr",
 }
 
+/**
+ * @deprecated Use MetadataFilter with operator EQ instead
+ */
 export interface ExactMatchFilter {
   filterType: "ExactMatch";
   key: string;
   value: string | number;
 }
 
+export enum FilterOperator {
+  EQ = "==", // default operator (string, number)
+  // GT = ">", // greater than (number)
+  // LT = "<", // less than (number)
+  // NE = "!=", // not equal to (string, number)
+  // GTE = ">=", // greater than or equal to (number)
+  // LTE = "<=", // less than or equal to (number)
+  // IN = "in", // In array (string or number)
+  // NIN = "nin", // Not in array (string or number)
+  // ANY = "any", // Contains any (array of strings)
+  // ALL = "all", // Contains all (array of strings)
+  // TEXT_MATCH = "text_match", // full text match (allows you to search for a specific substring, token or phrase within the text field)
+  // CONTAINS = "contains", // metadata array contains value (string or number)
+}
+
+export enum FilterCondition {
+  AND = "and",
+  OR = "or",
+}
+
+export type MetadataFilterValue = string | number | string[] | number[];
+
+export interface MetadataFilter {
+  key: string;
+  value: MetadataFilterValue;
+  operator: `${FilterOperator}`; // ==, any, all,...
+}
+
 export interface MetadataFilters {
-  filters: ExactMatchFilter[];
+  filters: Array<MetadataFilter>;
+  condition?: `${FilterCondition}`; // and, or
 }
 
 export interface VectorStoreQuerySpec {
   query: string;
-  filters: ExactMatchFilter[];
+  filters: MetadataFilter[];
   topK?: number;
 }
 

diff --git a/packages/llamaindex/tests/vectorStores/SimpleVectorStore.test.ts b/packages/llamaindex/tests/vectorStores/SimpleVectorStore.test.ts
@@ -74,7 +74,7 @@ describe("SimpleVectorStore", () => {
             {
               key: "private",
               value: "false",
-              filterType: "ExactMatch",
+              operator: "==",
             },
           ],
         },