VS-284: Add Wrangler support for QueryById (#7004)

* VS-284: Enable query by id via Wrangler

* Apply suggestions from code review

Co-authored-by: Carmen Popoviciu <[email protected]>

* Update .changeset/happy-scissors-deliver.md

Co-authored-by: Edmund Hung <[email protected]>

---------

Co-authored-by: Carmen Popoviciu <[email protected]>
Co-authored-by: emily-shen <[email protected]>
Co-authored-by: Edmund Hung <[email protected]>

Author: 4 people

.changeset/happy-scissors-deliver.md

@@ -0,0 +1,5 @@
+---
+"wrangler": minor
+---
+
+feat: Enable Vectorize query by id via Wrangler

packages/wrangler/src/__tests__/vectorize/vectorize.test.ts

@@ -173,7 +173,8 @@ describe("vectorize help", () => {
 			  -v, --version                   Show version number  [boolean]
 
 			OPTIONS
-			      --vector           Vector to query the Vectorize Index  [array] [required]
+			      --vector           Vector to query the Vectorize Index  [array]
+			      --vector-id        Identifier for a vector in the index against which the index should be queried  [string]
 			      --top-k            The number of results (nearest neighbors) to return  [number] [default: 5]
 			      --return-values    Specify if the vector values should be included in the results  [boolean] [default: false]
 			      --return-metadata  Specify if the vector metadata should be included in the results  [string] [choices: \\"all\\", \\"indexed\\", \\"none\\"] [default: \\"none\\"]
@@ -523,87 +524,25 @@ describe("vectorize commands", () => {
 		await runWrangler(
 			"vectorize query test-index --vector 1 2 3 '4' 1.5 '2.6' a 'b' null 7 abc 8 undefined"
 		);
-		expect(std.out).toMatchInlineSnapshot(`
-			"📋 Searching for relevant vectors...
-{
-  \\"count\\": 2,
-  \\"matches\\": [
-    {
-      \\"id\\": \\"a\\",
-      \\"score\\": 0.5,
-      \\"values\\": [
-        1,
-        2,
-        3,
-        4
-      ],
-      \\"namespace\\": \\"abcd\\",
-      \\"metadata\\": {
-        \\"a\\": true,
-        \\"b\\": 123
-      }
-    },
-    {
-      \\"id\\": \\"b\\",
-      \\"score\\": 0.75,
-      \\"values\\": [
-        5,
-        6,
-        7,
-        8
-      ],
-      \\"metadata\\": {
-        \\"c\\": false,
-        \\"b\\": \\"123\\"
-      }
-    }
-  ]
-}"
-		`);
+		expect(std.out).toMatchInlineSnapshot(querySnapshot);
+	});
+
+	it("should handle a query with a vector-id", async () => {
+		mockVectorizeV2Request();
+		await runWrangler("vectorize query test-index --vector-id some-vector-id");
+		expect(std.out).toMatchInlineSnapshot(querySnapshot);
+
+		// No warning or error
+		expect(std.warn).toMatchInlineSnapshot(`""`);
+		expect(std.err).toMatchInlineSnapshot(`""`);
 	});
 
 	it("should handle a query on a vectorize index with all options", async () => {
 		mockVectorizeV2Request();
 		await runWrangler(
 			`vectorize query test-index --vector 1 2 3 '4' --top-k=2 --return-values=true --return-metadata=indexed --namespace=abc --filter '{ "p1": "abc", "p2": { "$ne": true }, "p3": 10, "p4": false, "nested.p5": "abcd" }'`
 		);
-		expect(std.out).toMatchInlineSnapshot(`
-			"📋 Searching for relevant vectors...
-{
-  \\"count\\": 2,
-  \\"matches\\": [
-    {
-      \\"id\\": \\"a\\",
-      \\"score\\": 0.5,
-      \\"values\\": [
-        1,
-        2,
-        3,
-        4
-      ],
-      \\"namespace\\": \\"abcd\\",
-      \\"metadata\\": {
-        \\"a\\": true,
-        \\"b\\": 123
-      }
-    },
-    {
-      \\"id\\": \\"b\\",
-      \\"score\\": 0.75,
-      \\"values\\": [
-        5,
-        6,
-        7,
-        8
-      ],
-      \\"metadata\\": {
-        \\"c\\": false,
-        \\"b\\": \\"123\\"
-      }
-    }
-  ]
-}"
-		`);
+		expect(std.out).toMatchInlineSnapshot(querySnapshot);
 
 		// No warning > Valid filter
 		expect(std.warn).toMatchInlineSnapshot(`""`);
@@ -614,43 +553,7 @@ describe("vectorize commands", () => {
 		await runWrangler(
 			"vectorize query test-index --vector 1 2 3 '4' --filter='{ 'p1': [1,2,3] }'"
 		);
-		expect(std.out).toMatchInlineSnapshot(`
-			"📋 Searching for relevant vectors...
-{
-  \\"count\\": 2,
-  \\"matches\\": [
-    {
-      \\"id\\": \\"a\\",
-      \\"score\\": 0.5,
-      \\"values\\": [
-        1,
-        2,
-        3,
-        4
-      ],
-      \\"namespace\\": \\"abcd\\",
-      \\"metadata\\": {
-        \\"a\\": true,
-        \\"b\\": 123
-      }
-    },
-    {
-      \\"id\\": \\"b\\",
-      \\"score\\": 0.75,
-      \\"values\\": [
-        5,
-        6,
-        7,
-        8
-      ],
-      \\"metadata\\": {
-        \\"c\\": false,
-        \\"b\\": \\"123\\"
-      }
-    }
-  ]
-}"
-		`);
+		expect(std.out).toMatchInlineSnapshot(querySnapshot);
 
 		expect(std.warn).toMatchInlineSnapshot(`
 		"▲ [WARNING] 🚨 Invalid query filter. Please use the recommended format.
@@ -669,6 +572,34 @@ describe("vectorize commands", () => {
 		expect(std.warn).toMatchInlineSnapshot(`
 			"▲ [WARNING] Could not find any relevant vectors
 
+"
+		`);
+	});
+
+	it("should fail query when neither vector nor vector-id is provided", async () => {
+		mockVectorizeV2RequestError();
+		await runWrangler(
+			"vectorize query test-index --top-k=2 --return-values=true"
+		);
+		expect(std.out).toMatchInlineSnapshot(`""`);
+
+		expect(std.err).toMatchInlineSnapshot(`
+			"X [ERROR] 🚨 Either vector or vector-id parameter must be provided, but not both.
+
+"
+		`);
+	});
+
+	it("should fail query when both vector and vector-id are provided", async () => {
+		mockVectorizeV2RequestError();
+		await runWrangler(
+			"vectorize query test-index --vector 1 2 3 '4' --vector-id some-vector-id"
+		);
+		expect(std.out).toMatchInlineSnapshot(`""`);
+
+		expect(std.err).toMatchInlineSnapshot(`
+			"X [ERROR] 🚨 Either vector or vector-id parameter must be provided, but not both.
+
 "
 		`);
 	});
@@ -866,6 +797,43 @@ describe("vectorize query filter", () => {
 	});
 });
 
+const querySnapshot = `
+			"📋 Searching for relevant vectors...
+{
+  \\"count\\": 2,
+  \\"matches\\": [
+    {
+      \\"id\\": \\"a\\",
+      \\"score\\": 0.5,
+      \\"values\\": [
+        1,
+        2,
+        3,
+        4
+      ],
+      \\"namespace\\": \\"abcd\\",
+      \\"metadata\\": {
+        \\"a\\": true,
+        \\"b\\": 123
+      }
+    },
+    {
+      \\"id\\": \\"b\\",
+      \\"score\\": 0.75,
+      \\"values\\": [
+        5,
+        6,
+        7,
+        8
+      ],
+      \\"metadata\\": {
+        \\"c\\": false,
+        \\"b\\": \\"123\\"
+      }
+    }
+  ]
+}"`;
+
 /** Create a mock handler for the Vectorize API */
 function mockVectorizeRequest() {
 	msw.use(

packages/wrangler/src/vectorize/client.ts

@@ -131,7 +131,7 @@ export async function upsertIntoIndex(
 	);
 }
 
-export async function queryIndex(
+export async function queryIndexByVector(
 	config: Config,
 	indexName: string,
 	vector: VectorFloatArray | number[],
@@ -153,6 +153,28 @@ export async function queryIndex(
 	);
 }
 
+export async function queryIndexByVectorId(
+	config: Config,
+	indexName: string,
+	vectorId: string,
+	options: VectorizeQueryOptions
+): Promise<VectorizeMatches> {
+	const accountId = await requireAuth(config);
+	return await fetchResult(
+		`/accounts/${accountId}/vectorize/v2/indexes/${indexName}/query`,
+		{
+			method: "POST",
+			headers: {
+				"content-type": jsonContentType,
+			},
+			body: JSON.stringify({
+				...options,
+				vectorId,
+			}),
+		}
+	);
+}
+
 export async function getByIds(
 	config: Config,
 	indexName: string,

packages/wrangler/src/vectorize/query.ts

@@ -1,12 +1,13 @@
 import { readConfig } from "../config";
 import { logger } from "../logger";
-import { queryIndex } from "./client";
+import { queryIndexByVector, queryIndexByVectorId } from "./client";
 import { vectorizeGABanner } from "./common";
 import type {
 	CommonYargsArgv,
 	StrictYargsOptionsToInterface,
 } from "../yargs-types";
 import type {
+	VectorizeMatches,
 	VectorizeMetadataFilterInnerValue,
 	VectorizeMetadataFilterValue,
 	VectorizeMetadataRetrievalLevel,
@@ -26,7 +27,6 @@ export function options(yargs: CommonYargsArgv) {
 		.options({
 			vector: {
 				type: "array",
-				demandOption: true,
 				describe: "Vector to query the Vectorize Index",
 				coerce: (arg: unknown[]) =>
 					arg
@@ -38,6 +38,11 @@ export function options(yargs: CommonYargsArgv) {
 								typeof value === "number" && !isNaN(value)
 						),
 			},
+			"vector-id": {
+				type: "string",
+				describe:
+					"Identifier for a vector in the index against which the index should be queried",
+			},
 			"top-k": {
 				type: "number",
 				default: 5,
@@ -114,10 +119,35 @@ export async function handler(
 		}
 	}
 
+	if (
+		(args.vector === undefined && args.vectorId === undefined) ||
+		(args.vector !== undefined && args.vectorId !== undefined)
+	) {
+		logger.error(
+			"🚨 Either vector or vector-id parameter must be provided, but not both."
+		);
+		return;
+	}
+
 	logger.log(`📋 Searching for relevant vectors...`);
-	const res = await queryIndex(config, args.name, args.vector, queryOptions);
+	let res: VectorizeMatches | undefined;
+	if (args.vector !== undefined) {
+		res = await queryIndexByVector(
+			config,
+			args.name,
+			args.vector,
+			queryOptions
+		);
+	} else if (args.vectorId !== undefined) {
+		res = await queryIndexByVectorId(
+			config,
+			args.name,
+			args.vectorId,
+			queryOptions
+		);
+	}
 
-	if (res.count === 0) {
+	if (res === undefined || res.count === 0) {
 		logger.warn(`Could not find any relevant vectors`);
 		return;
 	}