Merge pull request #3241 from cloudflare/ggu/vectorize-range-query

cloudflare · Dec 16, 2024 · 1299c61 · 1299c61
2 parents e0d6c79 + 6ce07d8
commit 1299c61
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 17 deletions.
diff --git a/src/cloudflare/internal/test/vectorize/vectorize-api-test.js b/src/cloudflare/internal/test/vectorize/vectorize-api-test.js
@@ -238,6 +238,30 @@ export const test_vector_search_vector_query = {
       };
       assert.deepStrictEqual(results, expected);
     }
+
+    {
+      const results = await IDX.query(new Float32Array(5), {
+        topK: 1,
+        filter: {
+          text: {
+            $gt: 'Peter Piper picked a peck of pickled peppers',
+            $lt: 'You know New York, you need New York, you know you need unique New York',
+          },
+        },
+      });
+      assert.equal(true, results.count > 0);
+      /** @type {VectorizeMatches}  */
+      const expected = {
+        matches: [
+          {
+            id: 'b0daca4a-ffd8-4865-926b-e24800af2a2d',
+            score: 0.71151,
+          },
+        ],
+        count: 1,
+      };
+      assert.deepStrictEqual(results, expected);
+    }
   },
 };
 

diff --git a/src/cloudflare/internal/test/vectorize/vectorize-mock.js b/src/cloudflare/internal/test/vectorize/vectorize-mock.js
@@ -102,24 +102,34 @@ export default {
         let returnSet = structuredClone(exampleVectorMatches);
         if (
           body?.filter?.['text'] &&
-          typeof body?.filter?.['text'] === 'object' &&
-          body?.filter?.['text']?.['$eq'] !== undefined
+          typeof body?.filter?.['text'] === 'object'
         ) {
-          const criteria = body?.filter?.['text']?.['$eq'];
-          returnSet = returnSet.filter(
-            (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?.filter?.['text']?.['$eq'] !== undefined) {
+            const criteria = body?.filter?.['text']?.['$eq'];
+            returnSet = returnSet.filter(
+              (m) => m.metadata?.['text'] === criteria
+            );
+          }
+          if (body?.filter?.['text']?.['$in'] !== undefined) {
+            const criteria = body?.filter?.['text']?.['$in'];
+            returnSet = returnSet.filter((m) =>
+              criteria.includes(m.metadata?.['text'])
+            );
+          }
+          if (body?.filter?.['text']?.['$gt'] !== undefined) {
+            const criteria = body?.filter?.['text']?.['$gt'];
+            returnSet = returnSet.filter(
+              (m) => m.metadata?.['text'] > criteria
+            );
+          }
+          if (body?.filter?.['text']?.['$lt'] !== undefined) {
+            const criteria = body?.filter?.['text']?.['$lt'];
+            returnSet = returnSet.filter(
+              (m) => m.metadata?.['text'] < criteria
+            );
+          }
         }
+
         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
@@ -33,7 +33,13 @@ interface VectorizeError {
  *
  * This list is expected to grow as support for more operations are released.
  */
-type VectorizeVectorMetadataFilterOp = '$eq' | '$ne';
+type VectorizeVectorMetadataFilterOp =
+  | '$eq'
+  | '$ne'
+  | '$lt'
+  | '$lte'
+  | '$gt'
+  | '$gte';
 type VectorizeVectorMetadataFilterCollectionOp = '$in' | '$nin';
 
 /**