Merge pull request #3074 from cloudflare/ggu/vectorize-in-nin

cloudflare · Nov 8, 2024 · 8434636 · 8434636
2 parents df6d46c + fdd5c8f
commit 8434636
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 0 deletions.
diff --git a/src/cloudflare/internal/test/vectorize/vectorize-api-test.js b/src/cloudflare/internal/test/vectorize/vectorize-api-test.js
@@ -207,6 +207,37 @@ export const test_vector_search_vector_query = {
       };
       assert.deepStrictEqual(results, expected);
     }
+
+    {
+      // with returnValues = unset (false), returnMetadata = unset (none), filter = "Peter Piper picked a peck of pickled peppers"
+      const results = await IDX.query(new Float32Array(new Array(5).fill(0)), {
+        topK: 1,
+        filter: {
+          text: {
+            $in: [
+              'Peter Piper picked a peck of pickled peppers',
+              'She sells seashells by the seashore',
+            ],
+          },
+        },
+      });
+      assert.equal(true, results.count > 0);
+      /** @type {VectorizeMatches}  */
+      const expected = {
+        matches: [
+          {
+            id: 'b0daca4a-ffd8-4865-926b-e24800af2a2d',
+            score: 0.71151,
+          },
+          {
+            id: 'a44706aa-a366-48bc-8cc1-3feffd87d548',
+            score: 0.68913,
+          },
+        ],
+        count: 2,
+      };
+      assert.deepStrictEqual(results, expected);
+    }
   },
 };
 

diff --git a/src/cloudflare/internal/test/vectorize/vectorize-mock.js b/src/cloudflare/internal/test/vectorize/vectorize-mock.js
@@ -110,6 +110,16 @@ export default {
             (m) => m.metadata?.['text'] === criteria
           );
         }
+        if (
+          body?.filter?.['text'] &&
+          typeof body?.filter?.['text'] === 'object' &&
+          body?.filter?.['text']?.['$in'] !== undefined
+        ) {
+          const criteria = body?.filter?.['text']?.['$in'];
+          returnSet = returnSet.filter((m) =>
+            criteria.includes(m.metadata?.['text'])
+          );
+        }
         if (!body?.returnValues)
           returnSet.forEach((v) => {
             delete v.values;

diff --git a/src/cloudflare/internal/vectorize.d.ts b/src/cloudflare/internal/vectorize.d.ts
@@ -34,6 +34,7 @@ interface VectorizeError {
  * This list is expected to grow as support for more operations are released.
  */
 type VectorizeVectorMetadataFilterOp = '$eq' | '$ne';
+type VectorizeVectorMetadataFilterCollectionOp = '$in' | '$nin';
 
 /**
  * Filter criteria for vector metadata used to limit the retrieved query result set.
@@ -47,6 +48,12 @@ type VectorizeVectorMetadataFilter = {
           VectorizeVectorMetadataValue,
           string[]
         > | null;
+      }
+    | {
+        [Op in VectorizeVectorMetadataFilterCollectionOp]?: Exclude<
+          VectorizeVectorMetadataValue,
+          string[]
+        >[];
       };
 };