Merge pull request #36 from upstash/DX-1015

Dx 1015

src/commands/client/query/index.ts

@@ -1,43 +1,2 @@
-import type { NAMESPACE } from "@commands/client/types";
-import { Dict } from "@commands/client/types";
-import { Command } from "@commands/command";
-
-type QueryCommandPayload = {
-  topK: number;
-  filter?: string;
-  includeVectors?: boolean;
-  includeMetadata?: boolean;
-  includeData?: boolean;
-} & ({ vector: number[]; data?: never } | { data: string; vector?: never });
-
-export type QueryResult<TMetadata = Dict> = {
-  id: number | string;
-  score: number;
-  vector?: number[];
-  metadata?: TMetadata;
-  data?: string;
-};
-
-type QueryCommandOptions = { namespace?: string };
-
-type QueryEndpointVariants =
-  | `query`
-  | `query-data`
-  | `query/${NAMESPACE}`
-  | `query-data/${NAMESPACE}`;
-
-export class QueryCommand<TMetadata> extends Command<QueryResult<TMetadata>[]> {
-  constructor(payload: QueryCommandPayload, options?: QueryCommandOptions) {
-    let endpoint: QueryEndpointVariants = "query";
-
-    if ("data" in payload) {
-      endpoint = "query-data";
-    }
-
-    if (options?.namespace) {
-      endpoint = `${endpoint}/${options.namespace}`;
-    }
-
-    super(payload, endpoint);
-  }
-}
+export * from "./query-many";
+export * from "./query-single";

src/commands/client/query/query-many/index.test.ts

@@ -0,0 +1,147 @@
+import { afterAll, describe, expect, test } from "bun:test";
+import { UpsertCommand } from "@commands/client/upsert";
+import { Index, awaitUntilIndexed, newHttpClient, randomID, range } from "@utils/test-utils";
+import { QueryManyCommand } from ".";
+
+const client = newHttpClient();
+
+describe("QUERY", () => {
+  const index = new Index();
+
+  afterAll(async () => {
+    await index.reset();
+  });
+  test("should query in batches successfully", async () => {
+    const ID = randomID();
+    const initialData = [
+      {
+        id: `1-${ID}`,
+        vector: range(0, 384),
+        metadata: {
+          animal: "elephant",
+          tags: ["mammal"],
+          diet: "herbivore",
+        },
+      },
+      {
+        id: `2-${ID}`,
+        vector: range(0, 384),
+        metadata: {
+          animal: "tiger",
+          tags: ["mammal"],
+          diet: "carnivore",
+        },
+      },
+    ];
+    await new UpsertCommand(initialData).exec(client);
+
+    await awaitUntilIndexed(client);
+
+    const res = await new QueryManyCommand<{
+      animal: string;
+      tags: string[];
+      diet: string;
+    }>([
+      {
+        vector: initialData[0].vector,
+        topK: 1,
+        filter: "tags[0] = 'mammal' AND diet = 'herbivore'",
+        includeMetadata: true,
+      },
+      {
+        vector: initialData[1].vector,
+        topK: 1,
+        filter: "tags[0] = 'mammal' AND diet = 'carnivore'",
+        includeMetadata: true,
+      },
+    ]).exec(client);
+
+    expect(res).toEqual([
+      [
+        {
+          id: `1-${ID}`,
+          score: 1,
+          metadata: { animal: "elephant", tags: ["mammal"], diet: "herbivore" },
+        },
+      ],
+      [
+        {
+          id: `2-${ID}`,
+          score: 1,
+          metadata: { animal: "tiger", tags: ["mammal"], diet: "carnivore" },
+        },
+      ],
+    ]);
+  });
+});
+
+describe("QUERY with Index Client", () => {
+  const index = new Index();
+
+  afterAll(async () => {
+    await index.reset();
+  });
+  test("should query in batches successfully", async () => {
+    const ID = randomID();
+    const initialData = [
+      {
+        id: `1-${ID}`,
+        vector: range(0, 384),
+        metadata: {
+          animal: "elephant",
+          tags: ["mammal"],
+          diet: "herbivore",
+        },
+      },
+      {
+        id: `2-${ID}`,
+        vector: range(0, 384),
+        metadata: {
+          animal: "tiger",
+          tags: ["mammal"],
+          diet: "carnivore",
+        },
+      },
+    ];
+
+    await index.upsert(initialData);
+
+    await awaitUntilIndexed(index);
+
+    const res = await index.queryMany<{
+      animal: string;
+      tags: string[];
+      diet: string;
+    }>([
+      {
+        vector: initialData[0].vector,
+        topK: 1,
+        filter: "tags[0] = 'mammal' AND diet = 'herbivore'",
+        includeMetadata: true,
+      },
+      {
+        vector: initialData[1].vector,
+        topK: 1,
+        filter: "tags[0] = 'mammal' AND diet = 'carnivore'",
+        includeMetadata: true,
+      },
+    ]);
+
+    expect(res).toEqual([
+      [
+        {
+          id: `1-${ID}`,
+          score: 1,
+          metadata: { animal: "elephant", tags: ["mammal"], diet: "herbivore" },
+        },
+      ],
+      [
+        {
+          id: `2-${ID}`,
+          score: 1,
+          metadata: { animal: "tiger", tags: ["mammal"], diet: "carnivore" },
+        },
+      ],
+    ]);
+  });
+});

src/commands/client/query/query-many/index.ts

@@ -0,0 +1,22 @@
+import { Command } from "@commands/command";
+import {
+  QueryCommandOptions,
+  QueryCommandPayload,
+  QueryEndpointVariants,
+  QueryResult,
+} from "../types";
+
+export class QueryManyCommand<TMetadata> extends Command<QueryResult<TMetadata>[][]> {
+  constructor(payload: QueryCommandPayload[], options?: QueryCommandOptions) {
+    let endpoint: QueryEndpointVariants = "query";
+
+    const hasData = payload.some((p) => p.data);
+    endpoint = hasData ? "query-data" : "query";
+
+    if (options?.namespace) {
+      endpoint = `${endpoint}/${options.namespace}`;
+    }
+
+    super(payload, endpoint);
+  }
+}

src/commands/client/query/query-single/index.ts

@@ -0,0 +1,23 @@
+import { Command } from "@commands/command";
+import {
+  QueryCommandOptions,
+  QueryCommandPayload,
+  QueryEndpointVariants,
+  QueryResult,
+} from "../types";
+
+export class QueryCommand<TMetadata> extends Command<QueryResult<TMetadata>[]> {
+  constructor(payload: QueryCommandPayload, options?: QueryCommandOptions) {
+    let endpoint: QueryEndpointVariants = "query";
+
+    if ("data" in payload) {
+      endpoint = "query-data";
+    }
+
+    if (options?.namespace) {
+      endpoint = `${endpoint}/${options.namespace}`;
+    }
+
+    super(payload, endpoint);
+  }
+}

src/commands/client/query/types.ts

@@ -0,0 +1,25 @@
+import { Dict, NAMESPACE } from "../types";
+
+export type QueryCommandPayload = {
+  topK: number;
+  filter?: string;
+  includeVectors?: boolean;
+  includeMetadata?: boolean;
+  includeData?: boolean;
+} & ({ vector: number[]; data?: never } | { data: string; vector?: never });
+
+export type QueryResult<TMetadata = Dict> = {
+  id: number | string;
+  score: number;
+  vector?: number[];
+  metadata?: TMetadata;
+  data?: string;
+};
+
+export type QueryCommandOptions = { namespace?: string };
+
+export type QueryEndpointVariants =
+  | `query`
+  | `query-data`
+  | `query/${NAMESPACE}`
+  | `query-data/${NAMESPACE}`;

src/commands/types.ts

@@ -1,5 +1,5 @@
 export { type Vector } from "./client/types";
 export { type RangeResult } from "./client/range/index";
 export { type FetchResult } from "./client/fetch/index";
-export { type QueryResult } from "./client/query/index";
+export { type QueryResult } from "./client/query/types";
 export { type InfoResult } from "./client/info/index";

src/vector.ts

@@ -4,6 +4,7 @@ import {
   InfoCommand,
   Namespace,
   QueryCommand,
+  QueryManyCommand,
   RangeCommand,
   ResetCommand,
   UpdateCommand,
@@ -73,14 +74,53 @@ export class Index<TIndexMetadata extends Dict = Dict> {
    * @param {string} [args.filter] - An optional filter string to be used in the query. The filter string is used to narrow down the query results.
    * @param {boolean} [args.includeVectors=false] - When set to true, includes the feature vectors of the returned items in the response.
    * @param {boolean} [args.includeMetadata=false] - When set to true, includes additional metadata of the returned items in the response.
+   * @param {boolean} [args.includeData=false] - When set to true, includes data - string - of the returned items in the response.
    *
-   * @returns A promise that resolves with an array of query result objects when the request to query the index is completed.
+   *  A promise that resolves with an array of query result objects when the request to query the index is completed.
    */
   query = <TMetadata extends Dict = TIndexMetadata>(
     args: CommandArgs<typeof QueryCommand>,
     options?: { namespace?: string }
   ) => new QueryCommand<TMetadata>(args, options).exec(this.client);
 
+  /**
+   * Queries an index with specified parameters.
+   * This method creates and executes a query command on an index based on the provided arguments.
+   *
+   * @example
+   * ```js
+   * await index.queryMany([
+   * {
+   *     topK: 3,
+   *     vector: [0.22, 0.66],
+   *     filter: "age >= 23 and (type = 'turtle' OR type = 'cat')",
+   * },
+   * {
+   *     topK: 3,
+   *     vector: [0.45, 0.52],
+   *     filter: "age >= 27 and (type = 'rabbit' OR type = 'dog')",
+   * },
+   * ]);
+   *
+   * ```
+   *
+   * @param {Object} args - The arguments for the query command.
+   * @param {number[]} args.vector - An array of numbers representing the feature vector for the query.
+   *                                This vector is utilized to find the most relevant items in the index.
+   * @param {number} args.topK - The desired number of top results to be returned, based on relevance or similarity to the query vector.
+   * @param {string} [args.filter] - An optional filter string to be used in the query. The filter string is used to narrow down the query results.
+   * @param {boolean} [args.includeVectors=false] - When set to true, includes the feature vectors of the returned items in the response.
+   * @param {boolean} [args.includeMetadata=false] - When set to true, includes additional metadata of the returned items in the response.
+   * @param {boolean} [args.includeData=false] - When set to true, includes data - string - of the returned items in the response.
+   *
+   *  A promise that resolves with an array of arrays of query result objects,
+   *  where each inner array represents a group of results matching a specific query condition.
+   */
+  queryMany = <TMetadata extends Dict = TIndexMetadata>(
+    args: CommandArgs<typeof QueryManyCommand>,
+    options?: { namespace?: string }
+  ) => new QueryManyCommand<TMetadata>(args, options).exec(this.client);
+
   /**
    * Upserts (Updates and Inserts) specific items into the index.
    * It's used for adding new items to the index or updating existing ones.