eigen-db 1.0.0

package/src/lib/__tests__/vector-db.test.ts

@@ -0,0 +1,443 @@
+import { describe, it, expect, beforeEach } from "vitest";
+import { readFileSync } from "fs";
+import { resolve } from "path";
+import { VectorDB } from "../vector-db";
+import { InMemoryStorageProvider } from "../storage";
+import { VectorCapacityExceededError } from "../errors";
+import { compileWatToWasm } from "../wasm-compute";
+
+const watSource = readFileSync(resolve(__dirname, "../simd.wat"), "utf-8");
+
+let wasmBinaryPromise: Promise<Uint8Array>;
+function getWasmBinary(): Promise<Uint8Array> {
+  if (!wasmBinaryPromise) {
+    wasmBinaryPromise = compileWatToWasm(watSource);
+  }
+  return wasmBinaryPromise;
+}
+
+describe("VectorDB", () => {
+  let storage: InMemoryStorageProvider;
+
+  beforeEach(() => {
+    storage = new InMemoryStorageProvider();
+  });
+
+  describe("with JS compute", () => {
+    runTestSuite(false);
+  });
+
+  describe("with WASM SIMD compute", () => {
+    runTestSuite(true);
+  });
+
+  function runTestSuite(useWasm: boolean) {
+    let wasmBinary: Uint8Array | null;
+
+    beforeEach(async () => {
+      wasmBinary = useWasm ? await getWasmBinary() : null;
+    });
+
+    // --- open ---
+    it("opens with zero records", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        storage,
+        wasmBinary,
+      });
+
+      expect(db.size).toBe(0);
+    });
+
+    // --- set and get ---
+    it("stores and retrieves a vector by key", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        normalize: false,
+        storage,
+        wasmBinary,
+      });
+
+      db.set("a", new Float32Array([1, 2, 3, 4]));
+      expect(db.size).toBe(1);
+
+      const result = db.get("a");
+      expect(result).toBeDefined();
+      expect(result![0]).toBeCloseTo(1);
+      expect(result![1]).toBeCloseTo(2);
+      expect(result![2]).toBeCloseTo(3);
+      expect(result![3]).toBeCloseTo(4);
+    });
+
+    it("last write wins when writing different value to the same key", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        normalize: false,
+        storage,
+        wasmBinary,
+      });
+
+      db.set("a", new Float32Array([1, 0, 0, 0]));
+      db.set("a", new Float32Array([0, 1, 0, 0]));
+
+      expect(db.size).toBe(1); // only one entry
+      const result = db.get("a");
+      expect(result![0]).toBeCloseTo(0);
+      expect(result![1]).toBeCloseTo(1);
+    });
+
+    it("last write wins with multiple overwrites", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        normalize: false,
+        storage,
+        wasmBinary,
+      });
+
+      db.set("key1", new Float32Array([1, 0, 0, 0]));
+      db.set("key1", new Float32Array([0, 1, 0, 0]));
+      db.set("key1", new Float32Array([0, 0, 1, 0]));
+
+      expect(db.size).toBe(1);
+      const result = db.get("key1");
+      expect(result![2]).toBeCloseTo(1);
+    });
+
+    it("returns undefined for non-existent key", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        storage,
+        wasmBinary,
+      });
+
+      expect(db.get("nonexistent")).toBeUndefined();
+    });
+
+    it("validates vector dimensions on set", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        storage,
+        wasmBinary,
+      });
+
+      expect(() => db.set("a", new Float32Array(8))).toThrow("dimension mismatch");
+    });
+
+    // --- setMany and getMany ---
+    it("stores and retrieves multiple entries with setMany/getMany", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        normalize: false,
+        storage,
+        wasmBinary,
+      });
+
+      db.setMany([
+        ["a", new Float32Array([1, 0, 0, 0])],
+        ["b", new Float32Array([0, 1, 0, 0])],
+        ["c", new Float32Array([0, 0, 1, 0])],
+      ]);
+
+      expect(db.size).toBe(3);
+
+      const results = db.getMany(["a", "b", "c", "d"]);
+      expect(results[0]![0]).toBeCloseTo(1);
+      expect(results[1]![1]).toBeCloseTo(1);
+      expect(results[2]![2]).toBeCloseTo(1);
+      expect(results[3]).toBeUndefined();
+    });
+
+    it("setMany with duplicate keys in batch uses last value", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        normalize: false,
+        storage,
+        wasmBinary,
+      });
+
+      db.setMany([
+        ["x", new Float32Array([1, 0, 0, 0])],
+        ["x", new Float32Array([0, 0, 0, 1])],
+      ]);
+
+      expect(db.size).toBe(1);
+      const result = db.get("x");
+      expect(result![3]).toBeCloseTo(1);
+    });
+
+    // --- query ---
+    it("query returns ranked results by similarity", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        storage,
+        wasmBinary,
+      });
+
+      db.set("x-axis", new Float32Array([1, 0, 0, 0]));
+      db.set("y-axis", new Float32Array([0, 1, 0, 0]));
+      db.set("xy-axis", new Float32Array([1, 1, 0, 0]));
+
+      const results = db.query(new Float32Array([1, 0, 0, 0]));
+      expect(results.length).toBe(3);
+
+      // x-axis should be the best match (identical direction)
+      expect(results.get(0).key).toBe("x-axis");
+      expect(results.get(0).score).toBeCloseTo(1.0, 2);
+
+      // xy-axis should be second (partially aligned)
+      expect(results.get(1).key).toBe("xy-axis");
+      expect(results.get(1).score).toBeGreaterThan(0);
+
+      // y-axis should be last (orthogonal)
+      expect(results.get(2).key).toBe("y-axis");
+      expect(results.get(2).score).toBeCloseTo(0.0, 2);
+    });
+
+    it("query respects topK option", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        storage,
+        wasmBinary,
+      });
+
+      db.set("a", new Float32Array([1, 0, 0, 0]));
+      db.set("b", new Float32Array([0, 1, 0, 0]));
+      db.set("c", new Float32Array([0, 0, 1, 0]));
+
+      const results = db.query(new Float32Array([1, 0, 0, 0]), { topK: 2 });
+      expect(results.length).toBe(2);
+    });
+
+    it("query on empty database returns empty result", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        storage,
+        wasmBinary,
+      });
+
+      const results = db.query(new Float32Array([1, 0, 0, 0]));
+      expect(results.length).toBe(0);
+    });
+
+    it("query validates vector dimensions", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        storage,
+        wasmBinary,
+      });
+
+      db.set("a", new Float32Array([1, 0, 0, 0]));
+      expect(() => db.query(new Float32Array(8))).toThrow("dimension mismatch");
+    });
+
+    it("query results support pagination", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        normalize: false,
+        storage,
+        wasmBinary,
+      });
+
+      for (let i = 0; i < 5; i++) {
+        const vec = new Float32Array(4);
+        vec[0] = 1 - i * 0.2;
+        db.set(`t${i}`, vec);
+      }
+
+      const results = db.query(new Float32Array([1, 0, 0, 0]), { normalize: false });
+
+      const page0 = results.getPage(0, 2);
+      expect(page0).toHaveLength(2);
+
+      const page1 = results.getPage(1, 2);
+      expect(page1).toHaveLength(2);
+
+      const page2 = results.getPage(2, 2);
+      expect(page2).toHaveLength(1);
+    });
+
+    it("query after overwrite uses updated vector", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        storage,
+        wasmBinary,
+      });
+
+      db.set("point", new Float32Array([1, 0, 0, 0]));
+      db.set("other", new Float32Array([0, 1, 0, 0]));
+
+      // Overwrite 'point' to be along y-axis
+      db.set("point", new Float32Array([0, 1, 0, 0]));
+
+      const results = db.query(new Float32Array([0, 1, 0, 0]));
+      // Both 'point' and 'other' are now along y-axis, so both should score high
+      expect(results.get(0).score).toBeCloseTo(1.0, 2);
+      expect(results.get(1).score).toBeCloseTo(1.0, 2);
+      expect(db.size).toBe(2);
+    });
+
+    // --- flush and persistence ---
+    it("flush persists data and reopen loads it", async () => {
+      const db1 = await VectorDB.open({
+        dimensions: 4,
+        normalize: false,
+        storage,
+        wasmBinary,
+      });
+
+      db1.set("alpha", new Float32Array([1, 0, 0, 0]));
+      db1.set("beta", new Float32Array([0, 1, 0, 0]));
+      await db1.flush();
+
+      // Reopen same storage
+      const db2 = await VectorDB.open({
+        dimensions: 4,
+        normalize: false,
+        storage,
+        wasmBinary,
+      });
+
+      expect(db2.size).toBe(2);
+      expect(db2.get("alpha")![0]).toBeCloseTo(1);
+      expect(db2.get("beta")![1]).toBeCloseTo(1);
+    });
+
+    it("flush persists overwritten values correctly", async () => {
+      const db1 = await VectorDB.open({
+        dimensions: 4,
+        normalize: false,
+        storage,
+        wasmBinary,
+      });
+
+      db1.set("key", new Float32Array([1, 0, 0, 0]));
+      db1.set("key", new Float32Array([0, 0, 0, 1]));
+      await db1.flush();
+
+      const db2 = await VectorDB.open({
+        dimensions: 4,
+        normalize: false,
+        storage,
+        wasmBinary,
+      });
+
+      expect(db2.size).toBe(1);
+      expect(db2.get("key")![3]).toBeCloseTo(1);
+    });
+
+    // --- close ---
+    it("close flushes and prevents further operations", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        normalize: false,
+        storage,
+        wasmBinary,
+      });
+
+      db.set("a", new Float32Array([1, 0, 0, 0]));
+      await db.close();
+
+      // Should be persisted
+      const db2 = await VectorDB.open({
+        dimensions: 4,
+        normalize: false,
+        storage,
+        wasmBinary,
+      });
+      expect(db2.size).toBe(1);
+
+      // Original instance should be closed
+      expect(() => db.set("b", new Float32Array(4))).toThrow("closed");
+    });
+
+    it("close is idempotent", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        storage,
+        wasmBinary,
+      });
+
+      await db.close();
+      await db.close(); // should not throw
+    });
+
+    // --- clear ---
+    it("clear removes all data", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        storage,
+        wasmBinary,
+      });
+
+      db.set("a", new Float32Array([1, 0, 0, 0]));
+      db.set("b", new Float32Array([0, 1, 0, 0]));
+      expect(db.size).toBe(2);
+
+      await db.clear();
+      expect(db.size).toBe(0);
+      expect(db.get("a")).toBeUndefined();
+    });
+
+    it("clear allows reuse of the database", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        normalize: false,
+        storage,
+        wasmBinary,
+      });
+
+      db.set("old", new Float32Array([1, 0, 0, 0]));
+      await db.clear();
+
+      db.set("new", new Float32Array([0, 1, 0, 0]));
+      expect(db.size).toBe(1);
+      expect(db.get("new")![1]).toBeCloseTo(1);
+    });
+
+    // --- normalization ---
+    it("normalizes vectors by default", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        storage,
+        wasmBinary,
+      });
+
+      db.set("a", new Float32Array([3, 0, 0, 0]));
+      const result = db.get("a");
+      // Should be normalized to unit length
+      expect(result![0]).toBeCloseTo(1.0);
+    });
+
+    it("skips normalization when normalize is false", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        normalize: false,
+        storage,
+        wasmBinary,
+      });
+
+      db.set("a", new Float32Array([3, 0, 0, 0]));
+      const result = db.get("a");
+      expect(result![0]).toBeCloseTo(3.0);
+    });
+
+    it("per-call normalize option overrides default", async () => {
+      const db = await VectorDB.open({
+        dimensions: 4,
+        normalize: false,
+        storage,
+        wasmBinary,
+      });
+
+      db.set("a", new Float32Array([3, 0, 0, 0]), { normalize: true });
+      const result = db.get("a");
+      expect(result![0]).toBeCloseTo(1.0);
+    });
+  }
+
+  it("throws VectorCapacityExceededError when full", async () => {
+    const err = new VectorCapacityExceededError(100);
+    expect(err).toBeInstanceOf(VectorCapacityExceededError);
+    expect(err.message).toContain("100");
+  });
+});

package/src/lib/__tests__/wasm-compute.test.ts

@@ -0,0 +1,152 @@
+import { describe, it, expect, beforeAll } from "vitest";
+import { readFileSync } from "fs";
+import { resolve } from "path";
+import { compileWatToWasm, instantiateWasm, type WasmExports } from "../wasm-compute";
+
+describe("WASM SIMD compute", () => {
+  let memory: WebAssembly.Memory;
+  let wasm: WasmExports;
+
+  beforeAll(async () => {
+    const watSource = readFileSync(resolve(__dirname, "../simd.wat"), "utf-8");
+    const wasmBinary = await compileWatToWasm(watSource);
+    memory = new WebAssembly.Memory({ initial: 1 }); // 64KB
+    wasm = await instantiateWasm(wasmBinary, memory);
+  });
+
+  function writeFloat32Array(offset: number, data: Float32Array): void {
+    new Float32Array(memory.buffer, offset, data.length).set(data);
+  }
+
+  function readFloat32Array(offset: number, length: number): Float32Array {
+    return new Float32Array(memory.buffer.slice(offset, offset + length * 4));
+  }
+
+  describe("normalize", () => {
+    it("normalizes a vector to unit length", () => {
+      const vec = new Float32Array([3, 4, 0, 0]);
+      const ptr = 0;
+      writeFloat32Array(ptr, vec);
+
+      wasm.normalize(ptr, 4);
+
+      const result = readFloat32Array(ptr, 4);
+      const mag = Math.sqrt(result[0] ** 2 + result[1] ** 2 + result[2] ** 2 + result[3] ** 2);
+      expect(mag).toBeCloseTo(1.0, 4);
+      expect(result[0]).toBeCloseTo(0.6, 4);
+      expect(result[1]).toBeCloseTo(0.8, 4);
+    });
+
+    it("normalizes a non-multiple-of-4 dimension vector", () => {
+      const vec = new Float32Array([1, 2, 3]);
+      const ptr = 0;
+      writeFloat32Array(ptr, vec);
+
+      wasm.normalize(ptr, 3);
+
+      const result = readFloat32Array(ptr, 3);
+      let sumSq = 0;
+      for (let i = 0; i < 3; i++) sumSq += result[i] * result[i];
+      expect(Math.sqrt(sumSq)).toBeCloseTo(1.0, 4);
+    });
+
+    it("handles zero vector", () => {
+      const vec = new Float32Array([0, 0, 0, 0]);
+      const ptr = 0;
+      writeFloat32Array(ptr, vec);
+
+      wasm.normalize(ptr, 4);
+
+      const result = readFloat32Array(ptr, 4);
+      expect(result[0]).toBe(0);
+      expect(result[1]).toBe(0);
+    });
+  });
+
+  describe("search_all", () => {
+    it("computes dot products correctly", () => {
+      const dimensions = 4;
+      const dbSize = 2;
+      const queryPtr = 0;
+      const dbPtr = dimensions * 4; // after query
+      const scoresPtr = dbPtr + dbSize * dimensions * 4; // after db
+
+      const query = new Float32Array([1, 0, 0, 0]);
+      const db = new Float32Array([
+        1, 0, 0, 0, // identical to query
+        0, 1, 0, 0, // orthogonal
+      ]);
+
+      writeFloat32Array(queryPtr, query);
+      writeFloat32Array(dbPtr, db);
+
+      wasm.search_all(queryPtr, dbPtr, scoresPtr, dbSize, dimensions);
+
+      const scores = readFloat32Array(scoresPtr, dbSize);
+      expect(scores[0]).toBeCloseTo(1.0, 4);
+      expect(scores[1]).toBeCloseTo(0.0, 4);
+    });
+
+    it("computes correct scores for normalized vectors", () => {
+      const dimensions = 4;
+      const dbSize = 3;
+      const queryPtr = 0;
+      const dbPtr = dimensions * 4;
+      const scoresPtr = dbPtr + dbSize * dimensions * 4;
+
+      // Create and normalize vectors manually
+      const q = new Float32Array([1, 2, 3, 0]);
+      const mag = Math.sqrt(q[0] ** 2 + q[1] ** 2 + q[2] ** 2);
+      for (let i = 0; i < 3; i++) q[i] /= mag;
+
+      const v1 = new Float32Array([1, 2, 3, 0]); // same direction
+      const v1mag = Math.sqrt(14);
+      for (let i = 0; i < 3; i++) v1[i] /= v1mag;
+
+      const v2 = new Float32Array([-1, -2, -3, 0]); // opposite
+      for (let i = 0; i < 3; i++) v2[i] /= v1mag;
+
+      const v3 = new Float32Array([0, 0, 0, 1]); // orthogonal
+
+      writeFloat32Array(queryPtr, q);
+      writeFloat32Array(dbPtr, new Float32Array([...v1, ...v2, ...v3]));
+
+      wasm.search_all(queryPtr, dbPtr, scoresPtr, dbSize, dimensions);
+
+      const scores = readFloat32Array(scoresPtr, dbSize);
+      expect(scores[0]).toBeCloseTo(1.0, 3); // same direction
+      expect(scores[1]).toBeCloseTo(-1.0, 3); // opposite
+      expect(scores[2]).toBeCloseTo(0.0, 3); // orthogonal
+    });
+
+    it("handles non-multiple-of-4 dimensions", () => {
+      const dimensions = 3;
+      const dbSize = 1;
+      const queryPtr = 0;
+      const dbPtr = 16; // aligned
+      const scoresPtr = dbPtr + dbSize * dimensions * 4 + 4; // extra space
+
+      const query = new Float32Array([1, 0, 0]);
+      const db = new Float32Array([0.5, 0.5, 0]);
+
+      writeFloat32Array(queryPtr, query);
+      writeFloat32Array(dbPtr, db);
+
+      wasm.search_all(queryPtr, dbPtr, scoresPtr, dbSize, dimensions);
+
+      const scores = readFloat32Array(scoresPtr, dbSize);
+      expect(scores[0]).toBeCloseTo(0.5, 4);
+    });
+
+    it("handles empty database", () => {
+      const queryPtr = 0;
+      const dbPtr = 16;
+      const scoresPtr = 32;
+
+      writeFloat32Array(queryPtr, new Float32Array([1, 0, 0, 0]));
+
+      // Should not crash with dbSize=0
+      wasm.search_all(queryPtr, dbPtr, scoresPtr, 0, 4);
+    });
+  });
+});

package/src/lib/compute.ts

@@ -0,0 +1,48 @@
+/**
+ * Pure JavaScript compute functions for vector operations.
+ * These serve as the reference implementation and fallback when WASM SIMD is unavailable.
+ */
+
+/**
+ * Normalizes a vector in-place to unit length.
+ * After normalization, cosine similarity reduces to a simple dot product.
+ */
+export function normalize(vec: Float32Array): void {
+  let sumSq = 0;
+  for (let i = 0; i < vec.length; i++) {
+    sumSq += vec[i] * vec[i];
+  }
+  const mag = Math.sqrt(sumSq);
+  if (mag === 0) return;
+  const invMag = 1 / mag;
+  for (let i = 0; i < vec.length; i++) {
+    vec[i] *= invMag;
+  }
+}
+
+/**
+ * Computes dot products of query against all vectors in the database.
+ * Writes scores to the output array.
+ *
+ * @param query - Normalized query vector (length = dimensions)
+ * @param db - Contiguous flat array of normalized vectors (length = dbSize * dimensions)
+ * @param scores - Output array for dot product scores (length = dbSize)
+ * @param dbSize - Number of vectors in the database
+ * @param dimensions - Dimensionality of each vector
+ */
+export function searchAll(
+  query: Float32Array,
+  db: Float32Array,
+  scores: Float32Array,
+  dbSize: number,
+  dimensions: number,
+): void {
+  for (let i = 0; i < dbSize; i++) {
+    let dot = 0;
+    const offset = i * dimensions;
+    for (let j = 0; j < dimensions; j++) {
+      dot += query[j] * db[offset + j];
+    }
+    scores[i] = dot;
+  }
+}

package/src/lib/errors.ts

@@ -0,0 +1,10 @@
+/**
+ * Thrown when the database exceeds the 4GB WebAssembly 32-bit memory limit,
+ * or the browser's available RAM.
+ */
+export class VectorCapacityExceededError extends Error {
+  constructor(maxVectors: number) {
+    super(`Capacity exceeded. Max vectors for this dimension size is ~${maxVectors}.`);
+    this.name = "VectorCapacityExceededError";
+  }
+}

package/src/lib/index.ts

@@ -0,0 +1,14 @@
+/**
+ * web-vector-base: In-Browser Vector Compute Engine
+ *
+ * High-performance vector search entirely on the client side,
+ * utilizing OPFS for persistent storage and WASM SIMD for computation.
+ */
+
+export { VectorDB } from "./vector-db";
+export { ResultSet } from "./result-set";
+export type { ResultItem } from "./result-set";
+export { VectorCapacityExceededError } from "./errors";
+export type { OpenOptions, OpenOptionsInternal, SetOptions, QueryOptions } from "./types";
+export { InMemoryStorageProvider, OPFSStorageProvider } from "./storage";
+export type { StorageProvider } from "./storage";