feat: add hybrid search

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

@@ -4,6 +4,8 @@ import {
   Index,
   awaitUntilIndexed,
   newHttpClient,
+  populateHybridIndex,
+  populateSparseIndex,
   randomID,
   range,
   resetIndexes,
@@ -69,6 +71,14 @@ describe("FETCH with Index Client", () => {
     token: process.env.UPSTASH_VECTOR_REST_TOKEN!,
     url: process.env.UPSTASH_VECTOR_REST_URL!,
   });
+  const sparseIndex = new Index({
+    token: process.env.SPARSE_UPSTASH_VECTOR_REST_TOKEN!,
+    url: process.env.SPARSE_UPSTASH_VECTOR_REST_URL!,
+  });
+  const hybridIndex = new Index({
+    token: process.env.HYBRID_UPSTASH_VECTOR_REST_TOKEN!,
+    url: process.env.HYBRID_UPSTASH_VECTOR_REST_URL!,
+  });
 
   test("should fetch array of records by IDs succesfully", async () => {
     const randomizedData = Array.from({ length: 20 })
@@ -157,4 +167,98 @@ describe("FETCH with Index Client", () => {
     const { data: _data, ...mockDataWithoutData } = mockData;
     expect(fetchWithID).toEqual([mockDataWithoutData]);
   });
+
+  test("should fetch from sparse", async () => {
+    const namespace = "fetch-hybrid";
+    await populateSparseIndex(sparseIndex, namespace);
+
+    const result = await sparseIndex.fetch(["id0", "id1", "id2", "id3"], {
+      includeVectors: true,
+      includeMetadata: true,
+      includeData: true,
+      namespace,
+    });
+
+    expect(result).toEqual([
+      {
+        id: "id0",
+        metadata: undefined,
+        data: undefined,
+        vector: undefined,
+        sparseVector: [
+          [0, 1],
+          [0.1, 0.2],
+        ],
+      },
+      {
+        id: "id1",
+        metadata: { key: "value" },
+        data: undefined,
+        vector: undefined,
+        sparseVector: [
+          [0, 1],
+          [0.2, 0.3],
+        ],
+      },
+      {
+        id: "id2",
+        metadata: { key: "value" },
+        data: "data",
+        vector: undefined,
+        sparseVector: [
+          [0, 1],
+          [0.3, 0.4],
+        ],
+      },
+      // @ts-expect-error checking an index that doesn't exist
+      undefined,
+    ]);
+  });
+
+  test("should fetch from hybrid", async () => {
+    const namespace = "fetch-hybrid";
+    await populateHybridIndex(hybridIndex, namespace);
+
+    const result = await hybridIndex.fetch(["id0", "id1", "id2", "id3"], {
+      includeVectors: true,
+      includeMetadata: true,
+      includeData: true,
+      namespace,
+    });
+
+    expect(result).toEqual([
+      {
+        id: "id0",
+        metadata: undefined,
+        data: undefined,
+        vector: [0.1, 0.2],
+        sparseVector: [
+          [0, 1],
+          [0.1, 0.2],
+        ],
+      },
+      {
+        id: "id1",
+        metadata: { key: "value" },
+        data: undefined,
+        vector: [0.2, 0.3],
+        sparseVector: [
+          [0, 1],
+          [0.2, 0.3],
+        ],
+      },
+      {
+        id: "id2",
+        metadata: { key: "value" },
+        data: "data",
+        vector: [0.3, 0.4],
+        sparseVector: [
+          [0, 1],
+          [0.3, 0.4],
+        ],
+      },
+      // @ts-expect-error checking an index that doesn't exist
+      undefined,
+    ]);
+  });
 });

src/commands/client/query/index.ts

@@ -1,2 +1,3 @@
 export * from "./query-many";
 export * from "./query-single";
+export { FusionAlgorithm, QueryMode, WeightingStrategy } from "./types";

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

@@ -1,7 +1,15 @@
-import { afterAll, describe, expect, test } from "bun:test";
+import { afterAll, describe, expect, mock, test } from "bun:test";
 import { UpsertCommand } from "@commands/client/upsert";
-import { Index, awaitUntilIndexed, newHttpClient, randomID, range } from "@utils/test-utils";
+import {
+  Index,
+  awaitUntilIndexed,
+  newHttpClient,
+  populateHybridIndex,
+  randomID,
+  range,
+} from "@utils/test-utils";
 import { QueryManyCommand } from ".";
+import { FusionAlgorithm, WeightingStrategy } from "../types";
 
 const client = newHttpClient();
 
@@ -77,9 +85,14 @@ describe("QUERY", () => {
 
 describe("QUERY with Index Client", () => {
   const index = new Index();
+  const hybridIndex = new Index({
+    token: process.env.HYBRID_UPSTASH_VECTOR_REST_TOKEN!,
+    url: process.env.HYBRID_UPSTASH_VECTOR_REST_URL!,
+  });
 
   afterAll(async () => {
     await index.reset();
+    await hybridIndex.reset({ all: true });
   });
   test("should query in batches successfully", async () => {
     const ID = randomID();
@@ -144,4 +157,48 @@ describe("QUERY with Index Client", () => {
       ],
     ]);
   });
+
+  test("should query hybrid index", async () => {
+    const namespace = "query-hybrid";
+    const mockData = await populateHybridIndex(hybridIndex, namespace);
+
+    const result = await index.queryMany(
+      [
+        {
+          topK: 1,
+          vector: [0.1, 0.1],
+          sparseVector: [
+            [3, 4],
+            [0.1, 0.2],
+          ],
+          fusionAlgorithm: FusionAlgorithm.RRF,
+        },
+        {
+          topK: 1,
+          vector: [0.5, 0.1],
+          sparseVector: [
+            [0, 1],
+            [0.5, 0.1],
+          ],
+          includeVectors: true,
+        },
+        {
+          topK: 1,
+          sparseVector: [
+            [2, 3],
+            [0.5, 0.5],
+          ],
+          weightingStrategy: WeightingStrategy.IDF,
+          fusionAlgorithm: FusionAlgorithm.DBSF,
+          includeMetadata: true,
+        },
+      ],
+      {
+        namespace,
+      }
+    );
+
+    // @ts-expect-error will fix after testing with actual index
+    expect(result).toEqual("todo: fix with actual");
+  });
 });