Merge pull request #25 from upstash/dx-883-vector-sdk-namespace-support

DX-883: Namespace Support

index.ts

@@ -8,6 +8,7 @@ import {
 import * as core from "./src/vector";
 
 export type * from "@commands/types";
+import { Dict } from "@commands/client/types";
 
 export type { Requester, UpstashRequest, UpstashResponse };
 
@@ -36,10 +37,8 @@ export type IndexConfig = {
  * Serverless vector client for upstash.
  */
 export class Index<
-  TIndexMetadata extends Record<
-    string,
-    unknown
-  > = Record<string, unknown>
+  TIndexMetadata extends Dict
+  = Dict
 > extends core.Index<TIndexMetadata> {
   /**
    * Create a new vector client by providing the url and token
@@ -116,8 +115,8 @@ export class Index<
       baseUrl: url,
       retry: configOrRequester?.retry,
       headers: { authorization: `Bearer ${token}` },
-      cache: configOrRequester?.cache === false 
-        ? undefined 
+      cache: configOrRequester?.cache === false
+        ? undefined
         : configOrRequester?.cache || "no-store",
       signal: configOrRequester?.signal,
     });

src/commands/client/delete/index.ts

@@ -1,13 +1,21 @@
+import type { NAMESPACE } from "@commands/client/types";
 import { Command } from "@commands/command";
 
+type DeleteEndpointVariants = `delete` | `delete/${NAMESPACE}`;
 export class DeleteCommand extends Command<{ deleted: number }> {
-  constructor(id: (number[] | string[]) | number | string) {
+  constructor(id: (number[] | string[]) | number | string, options?: { namespace?: string }) {
+    let endpoint: DeleteEndpointVariants = "delete";
+
+    if (options?.namespace) {
+      endpoint = `${endpoint}/${options.namespace}`;
+    }
+
     const finalArr = [];
     if (Array.isArray(id)) {
       finalArr.push(...id);
     } else {
       finalArr.push(id);
     }
-    super(finalArr, "delete");
+    super(finalArr, endpoint);
   }
 }

src/commands/client/fetch/index.test-d.ts

@@ -1,3 +1,4 @@
+import { Dict } from "@commands/client/types";
 import { expectTypeOf, test } from "vitest";
 import { Index } from "../../../../index";
 
@@ -10,7 +11,7 @@ test("case 1: no metadata is provided, any object should be expected", () => {
 
   type RetrievedMetadata = NonNullable<NonNullable<RetrievedFetchVector>["metadata"]>;
 
-  expectTypeOf<RetrievedMetadata>().toEqualTypeOf<Record<string, unknown>>();
+  expectTypeOf<RetrievedMetadata>().toEqualTypeOf<Dict>();
 });
 
 test("case 2: index-level metadata is provided, index-level metadata should be expected", () => {

src/commands/client/fetch/index.ts

@@ -1,15 +1,27 @@
+import type { NAMESPACE, Vector } from "@commands/client/types";
+import { Dict } from "@commands/client/types";
 import { Command } from "@commands/command";
-import { Vector } from "../types";
 
 type FetchCommandOptions = {
   includeMetadata?: boolean;
   includeVectors?: boolean;
 };
 
-export type FetchResult<TMetadata = Record<string, unknown>> = Vector<TMetadata> | null;
+export type FetchResult<TMetadata = Dict> = Vector<TMetadata> | null;
+
+type FetchEndpointVariants = `fetch` | `fetch/${NAMESPACE}`;
 
 export class FetchCommand<TMetadata> extends Command<FetchResult<TMetadata>[]> {
-  constructor([ids, opts]: [ids: number[] | string[], opts?: FetchCommandOptions]) {
-    super({ ids, ...opts }, "fetch");
+  constructor(
+    [ids, opts]: [ids: number[] | string[], opts?: FetchCommandOptions],
+    options?: { namespace?: string }
+  ) {
+    let endpoint: FetchEndpointVariants = "fetch";
+
+    if (options?.namespace) {
+      endpoint = `${endpoint}/${options.namespace}`;
+    }
+
+    super({ ids, ...opts }, endpoint);
   }
 }

src/commands/client/index.ts

@@ -4,3 +4,5 @@ export * from "./upsert";
 export * from "./fetch";
 export * from "./range";
 export * from "./reset";
+export * from "./info";
+export * from "./namespace";

src/commands/client/info/index.ts

@@ -1,15 +1,26 @@
 import { Command } from "@commands/command";
 
+type NamespaceTitle = string;
+type NamespaceInfo = {
+  vectorCount: number;
+  pendingVectorCount: number;
+};
+
 export type InfoResult = {
   vectorCount: number;
   pendingVectorCount: number;
   indexSize: number;
   dimension: number;
   similarityFunction: "COSINE" | "EUCLIDEAN" | "DOT_PRODUCT";
+  namespaces: Record<NamespaceTitle, NamespaceInfo>;
 };
 
+type InfoEndpointVariants = `info`;
+
 export class InfoCommand extends Command<InfoResult> {
   constructor() {
-    super([], "info");
+    const endpoint: InfoEndpointVariants = "info";
+
+    super([], endpoint);
   }
 }

src/commands/client/namespace/index.test.ts

@@ -0,0 +1,79 @@
+import { afterAll, describe, expect, test } from "bun:test";
+import { range, resetIndexes } from "@utils/test-utils";
+
+import { sleep } from "bun";
+import { Index } from "../../../../index";
+
+describe("NAMESPACE", () => {
+  afterAll(async () => await resetIndexes());
+
+  const index = new Index({
+    url: process.env.UPSTASH_VECTOR_REST_URL!,
+    token: process.env.UPSTASH_VECTOR_REST_TOKEN!,
+  });
+
+  test("should append to specific namespace", async () => {
+    const namespace1 = index.namespace("test-namespace-1");
+    const namespace2 = index.namespace("test-namespace-2");
+
+    await namespace1.upsert({
+      id: 1,
+      vector: range(0, 384),
+      metadata: { namespace: "namespace1" },
+    });
+    await namespace2.upsert({
+      id: 2,
+      vector: range(0, 384),
+      metadata: { namespace: "namespace2" },
+    });
+
+    sleep(10000);
+
+    const query1 = await namespace1.query({
+      vector: range(0, 384),
+      topK: 3,
+      includeMetadata: true,
+    });
+    const query2 = await namespace2.query({
+      vector: range(0, 384),
+      topK: 3,
+      includeMetadata: true,
+    });
+
+    expect(query1.length).toEqual(1);
+    expect(query2.length).toEqual(1);
+    expect(query1[0].metadata?.namespace).toEqual("namespace1");
+    expect(query2[0].metadata?.namespace).toEqual("namespace2");
+  });
+
+  test("should reset namespace", async () => {
+    const namespace = index.namespace("test-namespace-reset");
+
+    await namespace.upsert({
+      id: 1,
+      vector: range(0, 384),
+      metadata: { namespace: "test-namespace-reset" },
+    });
+
+    sleep(1000);
+
+    const res = await namespace.query({
+      vector: range(0, 384),
+      topK: 3,
+      includeMetadata: true,
+    });
+    expect(res.length).toEqual(1);
+
+    await namespace.reset();
+
+    sleep(1000);
+
+    const res2 = await namespace.query({
+      vector: range(0, 384),
+      topK: 3,
+      includeMetadata: true,
+    });
+
+    expect(res2.length).toEqual(0);
+  });
+});

src/commands/client/namespace/index.ts

@@ -0,0 +1,162 @@
+import {
+  DeleteCommand,
+  FetchCommand,
+  QueryCommand,
+  RangeCommand,
+  ResetCommand,
+  UpsertCommand,
+} from "@commands/client";
+import { Dict } from "@commands/client/types";
+import { Requester } from "@http";
+import { CommandArgs } from "../../../vector";
+
+export class Namespace<TIndexMetadata extends Dict = Dict> {
+  protected client: Requester;
+  protected namespace: string;
+
+  /**
+   * Create a new index namespace client
+   *
+   * @example
+   * ```typescript
+   * const index = new Index({
+   *  url: "<UPSTASH_VECTOR_REST_URL>",
+   *  token: "<UPSTASH_VECTOR_REST_TOKEN>",
+   * });
+   *
+   * const namespace = index.namespace("ns");
+   * ```
+   */
+  constructor(client: Requester, namespace: string) {
+    this.client = client;
+    this.namespace = namespace;
+  }
+
+  /**
+   * Queries an index namespace with specified parameters.
+   * This method creates and executes a query command on an index based on the provided arguments.
+   *
+   * @example
+   * ```js
+   * await index.namespace("ns").query({
+   *  topK: 3,
+   *  vector: [ 0.22, 0.66 ],
+   *  filter: "age >= 23 and (type = \'turtle\' OR type = \'cat\')"
+   * });
+   * ```
+   *
+   * @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.
+   *
+   * @returns A promise that resolves with an array of query result objects when the request to query the index is completed.
+   */
+  upsert = <TMetadata extends Dict = TIndexMetadata>(
+    args: CommandArgs<typeof UpsertCommand<TMetadata>>
+  ) => new UpsertCommand<TMetadata>(args, { namespace: this.namespace }).exec(this.client);
+
+  /**
+   * Upserts (Updates and Inserts) specific items into the index namespace.
+   * It's used for adding new items to the index namespace or updating existing ones.
+   *
+   * @example
+   * ```js
+   * const upsertArgs = {
+   *   id: '123',
+   *   vector: [0.42, 0.87, ...],
+   *   metadata: { property1: 'value1', property2: 'value2' }
+   * };
+   * const upsertResult = await index.namespace("ns").upsert(upsertArgs);
+   * console.log(upsertResult); // Outputs the result of the upsert operation
+   * ```
+   *
+   * @param {CommandArgs<typeof UpsertCommand>} args - The arguments for the upsert command.
+   * @param {number|string} args.id - The unique identifier for the item being upserted.
+   * @param {number[]} args.vector - The feature vector associated with the item.
+   * @param {Dict} [args.metadata] - Optional metadata to be associated with the item.
+   *
+   * @returns {string} A promise that resolves with the result of the upsert operation after the command is executed.
+   */
+  fetch = <TMetadata extends Dict = TIndexMetadata>(...args: CommandArgs<typeof FetchCommand>) =>
+    new FetchCommand<TMetadata>(args, { namespace: this.namespace }).exec(this.client);
+
+  /**
+   * It's used for retrieving specific items from the index namespace, optionally including
+   * their metadata and feature vectors.
+   *
+   * @example
+   * ```js
+   * const fetchIds = ['123', '456'];
+   * const fetchOptions = { includeMetadata: true, includeVectors: false };
+   * const fetchResults = await index.namespace("ns").fetch(fetchIds, fetchOptions);
+   * console.log(fetchResults); // Outputs the fetched items
+   * ```
+   *
+   * @param {...CommandArgs<typeof FetchCommand>} args - The arguments for the fetch command.
+   * @param {(number[]|string[])} args[0] - An array of IDs of the items to be fetched.
+   * @param {FetchCommandOptions} args[1] - Options for the fetch operation.
+   * @param {boolean} [args[1].includeMetadata=false] - Optionally include metadata of the fetched items.
+   * @param {boolean} [args[1].includeVectors=false] - Optionally include feature vectors of the fetched items.
+   *
+   * @returns {Promise<FetchReturnResponse<TMetadata>[]>} A promise that resolves with an array of fetched items or null if not found, after the command is executed.
+   */
+  query = <TMetadata extends Dict = TIndexMetadata>(args: CommandArgs<typeof QueryCommand>) =>
+    new QueryCommand<TMetadata>(args, { namespace: this.namespace }).exec(this.client);
+
+  /**
+   * Deletes a specific item or items from the index namespace by their ID(s).   *
+   *
+   * @example
+   * ```js
+   * await index.namespace("ns").delete('test-id')
+   * ```
+   *
+   * @param id - List of ids or single id
+   * @returns A promise that resolves when the request to delete the index is completed.
+   */
+  delete = (args: CommandArgs<typeof DeleteCommand>) =>
+    new DeleteCommand(args, { namespace: this.namespace }).exec(this.client);
+
+  /**
+   * Retrieves a range of items from the index.
+   *
+   * @example
+   * ```js
+   * const rangeArgs = {
+   *   cursor: 0,
+   *   limit: 10,
+   *   includeVectors: true,
+   *   includeMetadata: false
+   * };
+   * const rangeResults = await index.namespace("ns").range(rangeArgs);
+   * console.log(rangeResults); // Outputs the result of the range operation
+   * ```
+   *
+   * @param {CommandArgs<typeof RangeCommand>} args - The arguments for the range command.
+   * @param {number|string} args.cursor - The starting point (cursor) for the range query.
+   * @param {number} args.limit - The maximum number of items to return in this range.
+   * @param {boolean} [args.includeVectors=false] - Optionally include the feature vectors of the items in the response.
+   * @param {boolean} [args.includeMetadata=false] - Optionally include additional metadata of the items in the response.
+   *
+   * @returns {Promise<RangeReturnResponse<TMetadata>>} A promise that resolves with the response containing the next cursor and an array of vectors, after the command is executed.
+   */
+  range = <TMetadata extends Dict = TIndexMetadata>(args: CommandArgs<typeof RangeCommand>) =>
+    new RangeCommand<TMetadata>(args, { namespace: this.namespace }).exec(this.client);
+
+  /**
+   * It's used for wiping all the vectors in a index namespace.
+   *
+   * @example
+   * ```js
+   * await index.namespace("ns").reset();
+   * console.log('Index namespace has been reset');
+   * ```
+   *
+   * @returns {Promise<string>} A promise that resolves with the result of the reset operation after the command is executed.
+   */
+  reset = () => new ResetCommand({ namespace: this.namespace }).exec(this.client);
+}

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

@@ -1,6 +1,6 @@
+import { Dict } from "@commands/client/types";
 import { expectTypeOf, test } from "vitest";
 import { Index } from "../../../../index";
-
 type Metadata = { genre: string; year: number };
 
 test("case 1: no metadata is provided, any object should be expected", () => {
@@ -10,7 +10,7 @@ test("case 1: no metadata is provided, any object should be expected", () => {
 
   type RetrievedMetadata = NonNullable<NonNullable<RetrievedQueryVector>["metadata"]>;
 
-  expectTypeOf<RetrievedMetadata>().toEqualTypeOf<Record<string, unknown>>();
+  expectTypeOf<RetrievedMetadata>().toEqualTypeOf<Dict>();
 });
 
 test("case 2: index-level metadata is provided, index-level metadata should be expected", () => {