Merge pull request #28 from upstash/dx-905-vectorjs-update-command

DX-905: Update Command

src/commands/client/index.ts

@@ -6,3 +6,4 @@ export * from "./range";
 export * from "./reset";
 export * from "./info";
 export * from "./namespace";
+export * from "./update";

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

@@ -73,4 +73,36 @@ describe("NAMESPACE", () => {
 
     expect(res2.length).toEqual(0);
   });
+
+  test("should update vector in namespace", async () => {
+    const index = new Index({
+      url: process.env.UPSTASH_VECTOR_REST_URL!,
+      token: process.env.UPSTASH_VECTOR_REST_TOKEN!,
+    });
+
+    const namespace = index.namespace("test-namespace-update");
+
+    await namespace.upsert([
+      {
+        id: "test-1",
+        vector: range(0, 384),
+        metadata: { upstash: "test-1-not-updated" },
+      },
+    ]);
+
+    await awaitUntilIndexed(index);
+
+    const res = await namespace.update({
+      id: "test-1",
+      metadata: { upstash: "test-1-updated" },
+    });
+
+    expect(res).toEqual({ updated: 1 });
+
+    await awaitUntilIndexed(index);
+
+    const fetchData = await namespace.fetch(["test-1"], { includeMetadata: true });
+
+    expect(fetchData[0]?.metadata?.upstash).toBe("test-1-updated");
+  });
 });

src/commands/client/namespace/index.ts

@@ -4,6 +4,7 @@ import {
   QueryCommand,
   RangeCommand,
   ResetCommand,
+  UpdateCommand,
   UpsertCommand,
 } from "@commands/client";
 import { Dict } from "@commands/client/types";
@@ -58,6 +59,31 @@ export class Namespace<TIndexMetadata extends Dict = Dict> {
     args: CommandArgs<typeof UpsertCommand<TMetadata>>
   ) => new UpsertCommand<TMetadata>(args, { namespace: this.namespace }).exec(this.client);
 
+  /*
+   * Updates specific items in the index.
+   * It's used for updating existing items in the index.
+   *
+   * @example
+   * ```js
+   * const updateArgs = {
+   *   id: '123',
+   *   metadata: { updatedProperty: 'value1' }
+   * };
+   * const updateResult = await index.update(updateArgs);
+   * console.log(updateResult); // Outputs the result of the update operation
+   * ```
+   *
+   * @param {CommandArgs<typeof UpdateCommand>} args - The arguments for the update command.
+   * @param {number|string} args.id - The unique identifier for the item being updated.
+   * @param {number[]} args.vector - The feature vector associated with the item.
+   * @param {Record<string, unknown>} [args.metadata] - Optional metadata to be associated with the item.
+   *
+   * @returns {Promise<{updated: number}>} A promise that returns the number of items successfully updated.
+   */
+  update = <TMetadata extends Dict = TIndexMetadata>(
+    args: CommandArgs<typeof UpdateCommand<TMetadata>>
+  ) => new UpdateCommand<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.

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

@@ -126,15 +126,15 @@ describe("QUERY", () => {
       await new UpsertCommand([
         {
           id: "hello-world",
-          data: "Test1-2-3-4-5",
+          data: "testing-plan-text",
           metadata: { upstash: "test" },
         },
       ]).exec(embeddingClient);
 
       await awaitUntilIndexed(embeddingClient);
 
       const res = await new QueryCommand({
-        data: "Test1-2-3-4-5",
+        data: "testing-plain-text",
         topK: 1,
         includeVectors: true,
         includeMetadata: true,
@@ -156,20 +156,20 @@ describe("QUERY", () => {
       await new UpsertCommand([
         {
           id: "hello-world",
-          data: "Test1-2-3-4-5",
+          data: "testing-bulk-data-original",
           metadata: { upstash: "Cookie" },
         },
         {
           id: "hello-world1",
-          data: "Test1-2-3-4-5-6",
+          data: "testing-bulk-data-secondary",
           metadata: { upstash: "Monster" },
         },
       ]).exec(embeddingClient);
 
       await awaitUntilIndexed(embeddingClient);
 
       const res = await new QueryCommand({
-        data: "Test1-2-3-4-5",
+        data: "testing-bulk-data-original",
         topK: 1,
         includeVectors: true,
         includeMetadata: true,

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

@@ -0,0 +1,33 @@
+import { afterAll, describe, expect, test } from "bun:test";
+import { FetchCommand, UpdateCommand, UpsertCommand } from "@commands/index";
+import { awaitUntilIndexed, newHttpClient, range, resetIndexes } from "@utils/test-utils";
+
+const client = newHttpClient();
+
+describe("UPDATE", () => {
+  afterAll(async () => await resetIndexes());
+
+  test("should update vector metadata", async () => {
+    await new UpsertCommand({
+      id: 1,
+      vector: range(0, 384),
+      metadata: { upstash: "test-simple" },
+    }).exec(client);
+
+    const res = await new UpdateCommand({
+      id: 1,
+      metadata: { upstash: "test-update" },
+    }).exec(client);
+
+    expect(res).toEqual({ updated: 1 });
+
+    await awaitUntilIndexed(client, 5000);
+
+    const fetchData = await new FetchCommand<{ upstash: string }>([
+      ["1"],
+      { includeMetadata: true },
+    ]).exec(client);
+
+    expect(fetchData[0]?.metadata?.upstash).toBe("test-update");
+  });
+});

src/commands/client/update/index.ts

@@ -0,0 +1,41 @@
+import type { NAMESPACE } from "@commands/client/types";
+import { Command } from "@commands/command";
+
+type NoInfer<T> = T extends infer U ? U : never;
+
+type MetadataUpdatePayload<TMetadata> = {
+  id: string | number;
+  metadata: NoInfer<TMetadata>;
+};
+
+type VectorUpdatePayload = {
+  id: string | number;
+  vector: number[];
+};
+
+type DataUpdatePayload = {
+  id: string | number;
+  data: string;
+};
+
+type Payload<TMetadata> =
+  | MetadataUpdatePayload<TMetadata>
+  | VectorUpdatePayload
+  | DataUpdatePayload;
+
+type UpdateCommandOptions = { namespace?: string };
+
+type UpdateEndpointVariants = `update` | `update/${NAMESPACE}`;
+
+type UpdateCommandResponse = { updated: number };
+export class UpdateCommand<TMetadata> extends Command<UpdateCommandResponse> {
+  constructor(payload: Payload<TMetadata>, opts?: UpdateCommandOptions) {
+    let endpoint: UpdateEndpointVariants = "update";
+
+    if (opts?.namespace) {
+      endpoint = `${endpoint}/${opts.namespace}`;
+    }
+
+    super(payload, endpoint);
+  }
+}

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

@@ -1,10 +1,15 @@
 import { afterAll, describe, expect, test } from "bun:test";
 import { FetchCommand, UpsertCommand } from "@commands/index";
-import { newHttpClient, range, resetIndexes } from "@utils/test-utils";
+import { newHttpClient, randomID, range, resetIndexes } from "@utils/test-utils";
+import { Index } from "../../../../index";
 
 const client = newHttpClient();
 
 describe("UPSERT", () => {
+  const index = new Index({
+    url: process.env.UPSTASH_VECTOR_REST_URL!,
+    token: process.env.UPSTASH_VECTOR_REST_TOKEN!,
+  });
   afterAll(async () => await resetIndexes());
 
   test("should add record successfully", async () => {
@@ -73,13 +78,13 @@ describe("UPSERT", () => {
       await new UpsertCommand([
         {
           id: "hello-world",
-          //@ts-expect-error Mixed usage of vector and data in the same upsert command is not allowed.
+
           data: "Test1-2-3-4-5",
           metadata: { upstash: "test" },
         },
         {
           id: "hello-world",
-
+          //@ts-expect-error Mixed usage of vector and data in the same upsert command is not allowed.
           vector: [1, 2, 3, 4],
           metadata: { upstash: "test" },
         },
@@ -110,4 +115,22 @@ describe("UPSERT", () => {
 
     expect(resUpsert).toEqual("Success");
   });
+
+  test("should run with index.upsert in bulk", async () => {
+    const upsertData = [
+      {
+        id: randomID(),
+        vector: range(0, 384),
+        metadata: { upstash: "test-simple-1" },
+      },
+      {
+        id: randomID(),
+        vector: range(0, 384),
+        metadata: { upstash: "test-simple-2" },
+      },
+    ];
+    const resUpsert = await index.upsert(upsertData, { namespace: "test-namespace" });
+
+    expect(resUpsert).toEqual("Success");
+  });
 });

src/commands/client/upsert/index.ts

@@ -3,43 +3,23 @@ import { Command } from "@commands/command";
 
 type NoInfer<T> = T extends infer U ? U : never;
 
-type BasePayload = {
+type VectorPayload<TMetadata> = {
   id: number | string;
+  vector: number[];
+  metadata?: NoInfer<TMetadata>;
 };
 
-type ExtendedVectorPayload<TMetadata> = BasePayload &
-  (
-    | {
-        metadata: NoInfer<TMetadata>;
-        vector?: number[];
-        data?: never;
-      }
-    | {
-        metadata?: NoInfer<TMetadata>;
-        vector: number[];
-        data?: never;
-      }
-  );
-
-type ExtendedDataPayload<TMetadata> = BasePayload &
-  (
-    | {
-        metadata: NoInfer<TMetadata>;
-        data: string;
-        vector?: never;
-      }
-    | {
-        metadata?: NoInfer<TMetadata>;
-        data: string;
-        vector?: never;
-      }
-  );
+type DataPayload<TMetadata> = {
+  id: number | string;
+  data: string;
+  metadata?: NoInfer<TMetadata>;
+};
 
 type Payload<TMetadata> =
-  | ExtendedDataPayload<TMetadata>
-  | ExtendedVectorPayload<TMetadata>
-  | ExtendedDataPayload<TMetadata>[]
-  | ExtendedVectorPayload<TMetadata>[];
+  | VectorPayload<TMetadata>
+  | DataPayload<TMetadata>
+  | VectorPayload<TMetadata>[]
+  | DataPayload<TMetadata>[];
 
 type UpsertCommandOptions = { namespace?: string };
 

src/commands/command.ts

@@ -5,6 +5,7 @@ import { NAMESPACE } from "./client/types";
 
 const ENDPOINTS = [
   "upsert",
+  "update",
   "query",
   "delete",
   "fetch",

src/commands/index.ts

@@ -3,6 +3,7 @@ export * from "./client/delete";
 export * from "./client/fetch";
 export * from "./client/query";
 export * from "./client/upsert";
+export * from "./client/update";
 export * from "./client/reset";
 export * from "./client/range";
 export * from "./client/info";

src/vector.ts

@@ -6,6 +6,7 @@ import {
   QueryCommand,
   RangeCommand,
   ResetCommand,
+  UpdateCommand,
   UpsertCommand,
 } from "@commands/client";
 import { Dict } from "@commands/client/types";
@@ -107,6 +108,34 @@ export class Index<TIndexMetadata extends Dict = Dict> {
     options?: { namespace?: string }
   ) => new UpsertCommand<TMetadata>(args, options).exec(this.client);
 
+  /*
+   * Updates specific items in the index.
+   * It's used for updating existing items in the index.
+   *
+   * @example
+   * ```js
+   * const updateArgs = {
+   *   id: '123',
+   *   vector: [0.42, 0.87, ...],
+   *   metadata: { property1: 'value1', property2: 'value2' }
+   * };
+   * const updateResult = await index.update(updateArgs);
+   * console.log(updateResult); // Outputs the result of the update operation
+   * ```
+   *
+   * @param {CommandArgs<typeof UpdateCommand>} args - The arguments for the update command.
+   * @param {number|string} args.id - The unique identifier for the item being updated.
+   * @param {number[]} args.vector - The feature vector associated with the item.
+   * @param {Record<string, unknown>} [args.metadata] - Optional metadata to be associated with the item.
+   * @param {string} [args.namespace] - The namespace to update the item in.
+   *
+   * @returns {Promise<{updated: number}>} A promise that returns the number of items successfully updated.
+   */
+  update = <TMetadata extends Dict = TIndexMetadata>(
+    args: CommandArgs<typeof UpdateCommand<TMetadata>>,
+    options?: { namespace?: string }
+  ) => new UpdateCommand<TMetadata>(args, options).exec(this.client);
+
   /**
    * It's used for retrieving specific items from the index, optionally including
    * their metadata and feature vectors.