feat: Allow boolean returnMetadata for Vectorize V2 queries (#2596)

Co-authored-by: James M Snell <[email protected]>
cloudflare · Aug 30, 2024 · 1739783 · 1739783
1 parent 2f13452
commit 1739783
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 12 deletions.
diff --git a/src/cloudflare/internal/test/vectorize/vectorize-api-test.js b/src/cloudflare/internal/test/vectorize/vectorize-api-test.js
@@ -57,6 +57,69 @@ export const test_vector_search_vector_query = {
       assert.deepStrictEqual(results, expected);
     }
 
+    {
+      // with returnValues = unset (false), returnMetadata = false ("none")
+      const results = await IDX.query(new Float32Array(new Array(5).fill(0)), {
+        topK: 3,
+        returnMetadata: false,
+      });
+      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,
+          },
+          {
+            id: '43cfcb31-07e2-411f-8bf9-f82a95ba8b96',
+            score: 0.94812,
+          },
+        ],
+        count: 3,
+      };
+      assert.deepStrictEqual(results, expected);
+    }
+
+    {
+      // with returnValues = unset (false), returnMetadata = true ("all")
+      const results = await IDX.query(new Float32Array(new Array(5).fill(0)), {
+        topK: 3,
+        returnMetadata: true,
+      });
+      assert.equal(true, results.count > 0);
+      /** @type {VectorizeMatches}  */
+      const expected = {
+        matches: [
+          {
+            id: 'b0daca4a-ffd8-4865-926b-e24800af2a2d',
+            metadata: { text: 'She sells seashells by the seashore' },
+            score: 0.71151,
+          },
+          {
+            id: 'a44706aa-a366-48bc-8cc1-3feffd87d548',
+            metadata: {
+              text: 'Peter Piper picked a peck of pickled peppers',
+            },
+            score: 0.68913,
+          },
+          {
+            id: '43cfcb31-07e2-411f-8bf9-f82a95ba8b96',
+            metadata: {
+              text: 'You know New York, you need New York, you know you need unique New York',
+            },
+            score: 0.94812,
+          },
+        ],
+        count: 3,
+      };
+      assert.deepStrictEqual(results, expected);
+    }
+
     {
       // with returnValues = unset (false), returnMetadata = unset (none)
       const results = await IDX.query(new Float32Array(new Array(5).fill(0)), {

diff --git a/src/cloudflare/internal/vectorize-api.ts b/src/cloudflare/internal/vectorize-api.ts
@@ -45,14 +45,20 @@ class VectorizeIndexImpl implements Vectorize {
     options?: VectorizeQueryOptions
   ): Promise<VectorizeMatches> {
     if (this.indexVersion === 'v2') {
-      if (
-        options &&
-        options.returnMetadata &&
-        !isVectorizeMetadataRetrievalLevel(options.returnMetadata)
-      ) {
-        throw new Error(
-          `Invalid returnMetadata option. Expected: "none", "indexed" or "all"; got: ${options.returnMetadata}`
-        );
+      if (options?.returnMetadata) {
+        if (
+          typeof options.returnMetadata !== 'boolean' &&
+          !isVectorizeMetadataRetrievalLevel(options.returnMetadata)
+        ) {
+          throw new Error(
+            `Invalid returnMetadata option. Expected: true, false, "none", "indexed" or "all"; got: ${options.returnMetadata}`
+          );
+        }
+
+        if (typeof options.returnMetadata === 'boolean') {
+          // Allow boolean returnMetadata for backward compatibility. true converts to 'all' and false converts to 'none'
+          options.returnMetadata = options.returnMetadata ? 'all' : 'none';
+        }
       }
       const res = await this._send(Operation.VECTOR_QUERY, `query`, {
         method: 'POST',
@@ -238,9 +244,7 @@ class VectorizeIndexImpl implements Vectorize {
   }
 }
 
-function isVectorizeMetadataRetrievalLevel(
-  value: unknown
-): value is VectorizeMetadataRetrievalLevel {
+function isVectorizeMetadataRetrievalLevel(value: unknown): boolean {
   return (
     typeof value === 'string' &&
     (value === 'all' || value === 'indexed' || value === 'none')

diff --git a/types/defines/vectorize.d.ts b/types/defines/vectorize.d.ts
@@ -59,7 +59,7 @@ type VectorizeDistanceMetric = "euclidean" | "cosine" | "dot-product";
  */
 type VectorizeMetadataRetrievalLevel = "all" | "indexed" | "none";
 
-interface VectorizeQueryOptions{
+interface VectorizeQueryOptions {
   topK?: number;
   namespace?: string;
   returnValues?: boolean;