diff --git a/src/commands/client/delete/index.test.ts b/src/commands/client/delete/index.test.ts index 845fd02..53d7c63 100644 --- a/src/commands/client/delete/index.test.ts +++ b/src/commands/client/delete/index.test.ts @@ -1,6 +1,7 @@ import { afterAll, describe, expect, test } from "bun:test"; import { DeleteCommand, UpsertCommand } from "@commands/index"; import { newHttpClient, randomID, range, resetIndexes } from "@utils/test-utils"; +import { Index } from "@utils/test-utils"; const client = newHttpClient(); @@ -47,3 +48,40 @@ describe("DELETE", () => { }); }); }); + +describe("DELETE with Index Client", () => { + const index = new Index({ + token: process.env.UPSTASH_VECTOR_REST_TOKEN!, + url: process.env.UPSTASH_VECTOR_REST_URL!, + }); + afterAll(async () => { + await index.reset(); + }); + + test("should delete single record succesfully", async () => { + const initialVector = range(0, 384); + const id = randomID(); + + index.upsert({ id, vector: initialVector }); + + const deletionResult = await index.delete(id); + + expect(deletionResult).toEqual({ + deleted: 1 + }); + }); + + test("should delete array of records successfully", async () => { + const initialVector = range(0, 384); + const idsToUpsert = [randomID(), randomID(), randomID()]; + + const upsertPromises = idsToUpsert.map((id) => index.upsert({ id, vector: initialVector })); + + await Promise.all(upsertPromises); + + const deletionResult = await index.delete(idsToUpsert); + expect(deletionResult).toEqual({ + deleted: 3, + }) + }); +}); diff --git a/src/commands/client/fetch/index.test.ts b/src/commands/client/fetch/index.test.ts index ebdd0b5..75544fc 100644 --- a/src/commands/client/fetch/index.test.ts +++ b/src/commands/client/fetch/index.test.ts @@ -55,20 +55,30 @@ describe("FETCH", () => { expect(res).toEqual([mockData]); }); +}); - test("should fetch succesfully by index.fetch", async () => { - const index = new Index({ - url: process.env.UPSTASH_VECTOR_REST_URL!, - token: process.env.UPSTASH_VECTOR_REST_TOKEN!, - }); +describe("FETCH with Index Client", () => { + const index = new Index({ + token: process.env.UPSTASH_VECTOR_REST_TOKEN!, + url: process.env.UPSTASH_VECTOR_REST_URL!, + }); - const randomFetch = await index.fetch([randomID()], { - includeMetadata: true, - namespace: "test", - }); + test("should fetch array of records by IDs succesfully", async () => { + const randomizedData = new Array(20) + .fill("") + .map(() => ({ id: randomID(), vector: range(0, 384) })); - expect(randomFetch).toEqual([null]); + await index.upsert(randomizedData); + + await awaitUntilIndexed(index); + const IDs = randomizedData.map((x) => x.id); + const res = await index.fetch(IDs, { includeVectors: true }); + + expect(res).toEqual(randomizedData); + }); + + test("should fetch single record by ID", async () => { const mockData = { id: randomID(), vector: range(0, 384), @@ -87,4 +97,13 @@ describe("FETCH", () => { expect(fetchWithID).toEqual([mockData]); }); + + test("should return null when ID does not exist", async () => { + const randomFetch = await index.fetch([randomID()], { + includeMetadata: true, + namespace: "test", + }); + + expect(randomFetch).toEqual([null]); + }); }); diff --git a/src/commands/client/query/index.test.ts b/src/commands/client/query/index.test.ts index cb019ea..eb56625 100644 --- a/src/commands/client/query/index.test.ts +++ b/src/commands/client/query/index.test.ts @@ -1,11 +1,25 @@ import { afterAll, describe, expect, test } from "bun:test"; import { QueryCommand, UpsertCommand } from "@commands/index"; -import { awaitUntilIndexed, newHttpClient, range, resetIndexes } from "@utils/test-utils"; +import { awaitUntilIndexed, newHttpClient, randomID, range } from "@utils/test-utils"; +import { Index } from "@utils/test-utils"; const client = newHttpClient(); describe("QUERY", () => { - afterAll(async () => await resetIndexes()); + const index = new Index({ + token: process.env.UPSTASH_VECTOR_REST_TOKEN!, + url: process.env.UPSTASH_VECTOR_REST_URL!, + }); + + const embeddingIndex = new Index({ + token: process.env.EMBEDDING_UPSTASH_VECTOR_REST_TOKEN!, + url: process.env.EMBEDDING_UPSTASH_VECTOR_REST_URL!, + }); + + afterAll(async () => { + await index.reset(); + await embeddingIndex.reset(); + }); test("should query records successfully", async () => { const initialVector = range(0, 384); const initialData = { id: 33, vector: initialVector }; @@ -180,3 +194,117 @@ describe("QUERY", () => { { timeout: 20000 } ); }); + +describe("with Index Client", () => { + const index = new Index({ + token: process.env.UPSTASH_VECTOR_REST_TOKEN!, + url: process.env.UPSTASH_VECTOR_REST_URL!, + }); + const embeddingIndex = new Index({ + token: process.env.EMBEDDING_UPSTASH_VECTOR_REST_TOKEN!, + url: process.env.EMBEDDING_UPSTASH_VECTOR_REST_URL!, + }); + + afterAll(async () => { + await index.reset(); + await embeddingIndex.reset(); + }); + + test("should query records successfully", async () => { + const ID = randomID(); + const initialVector = range(0, 384); + const initialData = { id: ID, vector: initialVector }; + await index.upsert(initialData); + + await awaitUntilIndexed(index); + + const res = await index.query<{ hello: "World" }>({ + includeVectors: true, + vector: initialVector, + topK: 1, + }); + + expect(res).toEqual([ + { + id: ID, + score: 1, + vector: initialVector, + }, + ]); + }); + + test( + "should query with plain text successfully", + async () => { + await embeddingIndex.upsert([ + { + id: "hello-world", + data: "with-index-plain-text-query-test", + metadata: { upstash: "test" }, + }, + ]); + + await awaitUntilIndexed(embeddingIndex); + + const res = await embeddingIndex.query({ + data: "with-index-plain-text-query-test", + topK: 1, + includeVectors: true, + includeMetadata: true, + }); + + expect(res[0].metadata).toEqual({ upstash: "test" }); + }, + { timeout: 20000 } + ); + + test("should narrow down the query results with filter", async () => { + const ID = randomID(); + const initialVector = range(0, 384); + const initialData = [ + { + id: `1-${ID}`, + vector: initialVector, + metadata: { + animal: "elephant", + tags: ["mammal"], + diet: "herbivore", + }, + }, + { + id: `2-${ID}`, + vector: initialVector, + metadata: { + animal: "tiger", + tags: ["mammal"], + diet: "carnivore", + }, + }, + ]; + + await index.upsert(initialData); + + await awaitUntilIndexed(index); + + const res = await index.query<{ + animal: string; + tags: string[]; + diet: string; + }>({ + vector: initialVector, + topK: 1, + filter: "tags[0] = 'mammal' AND diet = 'carnivore'", + includeVectors: true, + includeMetadata: true, + }); + + expect(res).toEqual([ + { + id: `2-${ID}`, + score: 1, + vector: initialVector, + metadata: { animal: "tiger", tags: ["mammal"], diet: "carnivore" }, + }, + ]); + }); +}); diff --git a/src/commands/client/range/index.test.ts b/src/commands/client/range/index.test.ts index 299ed10..8786e3d 100644 --- a/src/commands/client/range/index.test.ts +++ b/src/commands/client/range/index.test.ts @@ -1,13 +1,20 @@ import { afterAll, describe, expect, test } from "bun:test"; import { RangeCommand, UpsertCommand } from "@commands/index"; -import { newHttpClient, randomID, range, resetIndexes } from "@utils/test-utils"; +import { + Index, + awaitUntilIndexed, + newHttpClient, + randomID, + range, + resetIndexes, +} from "@utils/test-utils"; const client = newHttpClient(); describe("RANGE", () => { afterAll(async () => await resetIndexes()); - test("should query records successfully", async () => { + test("should paginate records successfully", async () => { const randomizedData = new Array(20) .fill("") .map(() => ({ id: randomID(), vector: range(0, 384) })); @@ -15,6 +22,8 @@ describe("RANGE", () => { const payloads = randomizedData.map((data) => new UpsertCommand(data).exec(client)); await Promise.all(payloads); + await awaitUntilIndexed(client); + const res = await new RangeCommand({ cursor: 0, limit: 5, @@ -23,3 +32,31 @@ describe("RANGE", () => { expect(res.nextCursor).toBe("5"); }); }); + +describe("RANGE with Index Client", () => { + const index = new Index({ + token: process.env.UPSTASH_VECTOR_REST_TOKEN!, + url: process.env.UPSTASH_VECTOR_REST_URL!, + }); + + afterAll(async () => { + await index.reset(); + }); + test("should paginate records successfully", async () => { + const randomizedData = new Array(20) + .fill("") + .map(() => ({ id: randomID(), vector: range(0, 384) })); + + await index.upsert(randomizedData); + + await awaitUntilIndexed(index); + + const res = await index.range({ + cursor: 0, + limit: 5, + includeVectors: true, + }); + + expect(res.nextCursor).toBe("5"); + }); +}); diff --git a/src/commands/client/update/index.test.ts b/src/commands/client/update/index.test.ts index 75d6c19..7907c94 100644 --- a/src/commands/client/update/index.test.ts +++ b/src/commands/client/update/index.test.ts @@ -1,6 +1,6 @@ import { afterAll, describe, expect, test } from "bun:test"; import { FetchCommand, UpdateCommand, UpsertCommand } from "@commands/index"; -import { awaitUntilIndexed, newHttpClient, range, resetIndexes } from "@utils/test-utils"; +import { Index, awaitUntilIndexed, newHttpClient, range, resetIndexes } from "@utils/test-utils"; const client = newHttpClient(); @@ -31,3 +31,36 @@ describe("UPDATE", () => { expect(fetchData[0]?.metadata?.upstash).toBe("test-update"); }); }); + +describe("UPDATE with Index Client", () => { + const index = new Index({ + token: process.env.UPSTASH_VECTOR_REST_TOKEN!, + url: process.env.UPSTASH_VECTOR_REST_URL!, + }); + afterAll(async () => { + await index.reset(); + }); + + test("should update vector metadata", async () => { + await index.upsert({ + id: 1, + vector: range(0, 384), + metadata: { upstash: "test-simple" }, + }); + + await awaitUntilIndexed(index); + + const res = await index.update({ + id: 1, + metadata: { upstash: "test-update" }, + }); + + expect(res).toEqual({ updated: 1 }); + + await awaitUntilIndexed(client, 5000); + + const fetchData = await index.fetch(["1"], { includeMetadata: true }); + + expect(fetchData[0]?.metadata?.upstash).toBe("test-update"); + }); +}); diff --git a/src/commands/client/upsert/index.test.ts b/src/commands/client/upsert/index.test.ts index 033c7b1..27eb084 100644 --- a/src/commands/client/upsert/index.test.ts +++ b/src/commands/client/upsert/index.test.ts @@ -1,15 +1,10 @@ import { afterAll, describe, expect, test } from "bun:test"; import { FetchCommand, UpsertCommand } from "@commands/index"; -import { newHttpClient, randomID, range, resetIndexes } from "@utils/test-utils"; -import { Index } from "../../../../index"; +import { Index, newHttpClient, randomID, range, resetIndexes } from "@utils/test-utils"; 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 () => { @@ -115,6 +110,38 @@ describe("UPSERT", () => { expect(resUpsert).toEqual("Success"); }); +}); + +describe("UPSERT with Index Client", () => { + const index = new Index({ + token: process.env.EMBEDDING_UPSTASH_VECTOR_REST_TOKEN!, + url: process.env.EMBEDDING_UPSTASH_VECTOR_REST_URL!, + }); + afterAll(async () => await resetIndexes()); + + test("should add record successfully", async () => { + const res = await index.upsert({ id: 1, vector: range(0, 384) }); + expect(res).toEqual("Success"); + }); + + // biome-ignore lint/nursery/useAwait: required to test bad payloads + test("should return an error when vector is missing", async () => { + const throwable = async () => { + //@ts-ignore + await new UpsertCommand({ id: 1 }).exec(client); + }; + expect(throwable).toThrow(); + }); + + test("should add data successfully with a metadata", async () => { + //@ts-ignore + const res = await new UpsertCommand({ + id: 1, + vector: range(0, 384), + metadata: { upstash: "test" }, + }).exec(client); + expect(res).toEqual("Success"); + }); test("should run with index.upsert in bulk", async () => { const upsertData = [ @@ -133,4 +160,36 @@ describe("UPSERT", () => { expect(resUpsert).toEqual("Success"); }); + + test("should add plain text as data successfully", async () => { + const res = await index.upsert([ + { + id: "hello-world", + data: "Test1-2-3-4-5", + metadata: { upstash: "test" }, + }, + ]); + expect(res).toEqual("Success"); + }); + + test("should fail to upsert due to mixed usage of vector and plain text", () => { + const throwable = async () => { + await index.upsert([ + { + id: "hello-world", + + 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" }, + }, + ]); + }; + + expect(throwable).toThrow(); + }); }); diff --git a/src/utils/test-utils.ts b/src/utils/test-utils.ts index 3562ee0..73f7ab1 100644 --- a/src/utils/test-utils.ts +++ b/src/utils/test-utils.ts @@ -3,6 +3,7 @@ import { InfoCommand } from "../commands/client/info"; import { ResetCommand } from "../commands/client/reset"; import { HttpClient, RetryConfig } from "../http"; import { Index } from "../vector"; +export * from "../../index"; export type NonArrayType = T extends Array ? U : T; @@ -57,7 +58,8 @@ export const resetIndexes = async () => await new ResetCommand().exec(newHttpCli export const range = (start: number, end: number, step = 1) => { const result = []; for (let i = start; i < end; i += step) { - result.push(i); + const randomNum = Math.floor(Math.random() * (end - start + 1)) + start; + result.push(randomNum); } return result; };