VS-257: Fix Vectorize query options

cloudflare · Aug 2, 2024 · 28f21c5 · 28f21c5
1 parent ce15e0b
commit 28f21c5
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 20 deletions.
diff --git a/src/cloudflare/internal/test/vectorize/vectorize-api-test.js b/src/cloudflare/internal/test/vectorize/vectorize-api-test.js
@@ -19,7 +19,7 @@ export const test_vector_search_vector_query = {
   async test(_, env) {
     const IDX = env["vector-search"];
     {
-      // with returnValues = true, returnMetadata = true
+      // with returnValues = true, returnMetadata = "indexed"
       const results = await IDX.query(new Float32Array(new Array(5).fill(0)), {
         topK: 3,
         returnValues: true,
@@ -58,7 +58,7 @@ export const test_vector_search_vector_query = {
     }
 
     {
-      // with returnValues = unset (false), returnMetadata = unset (false)
+      // with returnValues = unset (false), returnMetadata = unset (none)
       const results = await IDX.query(new Float32Array(new Array(5).fill(0)), {
         topK: 3,
       });
@@ -85,7 +85,7 @@ export const test_vector_search_vector_query = {
     }
 
     {
-      // with returnValues = unset (false), returnMetadata = unset (false), filter = "Peter Piper picked a peck of pickled peppers"
+      // 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: {

diff --git a/src/cloudflare/internal/test/vectorize/vectorize-mock.js b/src/cloudflare/internal/test/vectorize/vectorize-mock.js
@@ -97,7 +97,7 @@ export default {
       ) {
         return Response.json({});
       } else if (request.method === "POST" && pathname.endsWith("/query")) {
-        /** @type {VectorizeQueryOptions<VectorizeMetadataRetrievalLevel> & {vector: number[]}} */
+        /** @type {VectorizeQueryOptions & {vector: number[]}} */
         const body = await request.json();
         let returnSet = structuredClone(exampleVectorMatches);
         if (
@@ -114,7 +114,7 @@ export default {
           returnSet.forEach((v) => {
             delete v.values;
           });
-        if (!body?.returnMetadata)
+        if (!body?.returnMetadata || body?.returnMetadata === "none")
           returnSet.forEach((v) => {
             delete v.metadata;
           });

diff --git a/src/cloudflare/internal/vectorize-api.ts b/src/cloudflare/internal/vectorize-api.ts
@@ -42,13 +42,18 @@ class VectorizeIndexImpl implements Vectorize {
 
   public async query(
     vector: VectorFloatArray | number[],
-    options: VectorizeQueryOptions<VectorizeMetadataRetrievalLevel>
+    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}`
+        );
+      }
       const res = await this._send(Operation.VECTOR_QUERY, `query`, {
         method: "POST",
         body: JSON.stringify({
-          ...options,
+          ...(options ?? {}),
           vector: Array.isArray(vector) ? vector : Array.from(vector),
         }),
         headers: {
@@ -59,6 +64,11 @@ class VectorizeIndexImpl implements Vectorize {
 
       return await toJson<VectorizeMatches>(res);
     } else {
+      if (options && options.returnMetadata && typeof options.returnMetadata !== 'boolean') {
+        throw new Error(
+          `Invalid returnMetadata option. Expected boolean; got: ${options.returnMetadata}`
+        );
+      }
       const compat = {
         queryMetadataOptional: flags.vectorizeQueryMetadataOptional,
       };
@@ -68,7 +78,7 @@ class VectorizeIndexImpl implements Vectorize {
         {
           method: "POST",
           body: JSON.stringify({
-            ...options,
+            ...(options ?? {}),
             vector: Array.isArray(vector) ? vector : Array.from(vector),
             compat,
           }),
@@ -220,6 +230,10 @@ class VectorizeIndexImpl implements Vectorize {
   }
 }
 
+function isVectorizeMetadataRetrievalLevel(value: any): value is VectorizeMetadataRetrievalLevel {
+  return value === 'all' || value === 'indexed' || value === 'none';
+}
+
 const maxBodyLogChars = 1_000;
 async function toJson<T = unknown>(response: Response): Promise<T> {
   const body = await response.text();

diff --git a/src/cloudflare/internal/vectorize.d.ts b/src/cloudflare/internal/vectorize.d.ts
@@ -67,13 +67,11 @@ type VectorizeDistanceMetric = "euclidean" | "cosine" | "dot-product";
  */
 type VectorizeMetadataRetrievalLevel = "all" | "indexed" | "none";
 
-interface VectorizeQueryOptions<
-  MetadataReturn extends boolean | VectorizeMetadataRetrievalLevel = boolean,
-> {
+interface VectorizeQueryOptions{
   topK?: number;
   namespace?: string;
   returnValues?: boolean;
-  returnMetadata?: MetadataReturn;
+  returnMetadata?: boolean | VectorizeMetadataRetrievalLevel;
   filter?: VectorizeVectorMetadataFilter;
 }
 
@@ -196,7 +194,7 @@ declare abstract class VectorizeIndex {
    */
   public query(
     vector: VectorFloatArray | number[],
-    options: VectorizeQueryOptions
+    options?: VectorizeQueryOptions
   ): Promise<VectorizeMatches>;
   /**
    * Insert a list of vectors into the index dataset. If a provided id exists, an error will be thrown.
@@ -243,7 +241,7 @@ declare abstract class Vectorize {
    */
   public query(
     vector: VectorFloatArray | number[],
-    options: VectorizeQueryOptions<VectorizeMetadataRetrievalLevel>
+    options?: VectorizeQueryOptions
   ): Promise<VectorizeMatches>;
   /**
    * Insert a list of vectors into the index dataset. If a provided id exists, an error will be thrown.

diff --git a/types/defines/vectorize.d.ts b/types/defines/vectorize.d.ts
@@ -59,13 +59,11 @@ type VectorizeDistanceMetric = "euclidean" | "cosine" | "dot-product";
  */
 type VectorizeMetadataRetrievalLevel = "all" | "indexed" | "none";
 
-interface VectorizeQueryOptions<
-  MetadataReturn extends boolean | VectorizeMetadataRetrievalLevel = boolean,
-> {
+interface VectorizeQueryOptions{
   topK?: number;
   namespace?: string;
   returnValues?: boolean;
-  returnMetadata?: MetadataReturn;
+  returnMetadata?: boolean | VectorizeMetadataRetrievalLevel;
   filter?: VectorizeVectorMetadataFilter;
 }
 
@@ -188,7 +186,7 @@ declare abstract class VectorizeIndex {
    */
   public query(
     vector: VectorFloatArray | number[],
-    options: VectorizeQueryOptions
+    options?: VectorizeQueryOptions
   ): Promise<VectorizeMatches>;
   /**
    * Insert a list of vectors into the index dataset. If a provided id exists, an error will be thrown.
@@ -235,7 +233,7 @@ declare abstract class Vectorize {
    */
   public query(
     vector: VectorFloatArray | number[],
-    options: VectorizeQueryOptions<VectorizeMetadataRetrievalLevel>
+    options?: VectorizeQueryOptions
   ): Promise<VectorizeMatches>;
   /**
    * Insert a list of vectors into the index dataset. If a provided id exists, an error will be thrown.