@convex-dev/rag 0.1.7

package/src/component/embeddings/importance.test.ts

@@ -0,0 +1,249 @@
+import { describe, it, expect } from "vitest";
+import {
+  searchVector,
+  vectorWithImportance,
+  modifyImportance,
+  getImportance,
+  vectorWithImportanceDimension,
+} from "./importance.js";
+
+describe("importance.ts", () => {
+  describe("searchVector", () => {
+    it("should add a 0 to the end of a normal embedding", () => {
+      const embedding = [0.1, 0.2, 0.3, 0.4];
+      const result = searchVector(embedding);
+
+      expect(result).toEqual([0.1, 0.2, 0.3, 0.4, 0]);
+      expect(result).toHaveLength(embedding.length + 1);
+    });
+
+    it("should handle 4096 dimension embeddings by slicing to 4095 and adding 0", () => {
+      const embedding = new Array(4096).fill(0).map((_, i) => i / 4096);
+      const result = searchVector(embedding);
+
+      expect(result).toHaveLength(4096);
+      expect(result[4095]).toBe(0);
+      expect(result.slice(0, 4095)).toEqual(embedding.slice(0, 4095));
+    });
+  });
+
+  function normalizeVector(vector: number[]) {
+    const sumOfSquares = vector.reduce((acc, v) => acc + v * v, 0);
+    const magnitude = Math.sqrt(sumOfSquares);
+    return magnitude === 0
+      ? vector.map(() => 0)
+      : vector.map((v) => v / magnitude);
+  }
+
+  describe("vectorWithImportance", () => {
+    it("should return normalized vector relative to importance", () => {
+      const embedding = [0.6, 0.8]; // magnitude = 1.0
+      const importance = 0.5;
+      const result = vectorWithImportance(embedding, importance);
+
+      expect(result).toHaveLength(3);
+      expect(result[0]).toBeCloseTo(embedding[0] * importance);
+      expect(result[1]).toBeCloseTo(embedding[1] * importance);
+      expect(
+        Math.sqrt(result[0] ** 2 + result[1] ** 2 + result[2] ** 2)
+      ).toBeCloseTo(1);
+    });
+
+    it("should handle maximum importance", () => {
+      const embedding = [0.6, 0.8];
+      const importance = 1.0;
+      const result = vectorWithImportance(embedding, importance);
+
+      expect(result).toHaveLength(3);
+      expect(result[0]).toBeCloseTo(0.6);
+      expect(result[1]).toBeCloseTo(0.8);
+    });
+
+    it("should handle minimum importance", () => {
+      const embedding = [0.6, 0.8];
+      const importance = 0.0;
+      const result = vectorWithImportance(embedding, importance);
+
+      expect(result).toHaveLength(3);
+      expect(result[0]).toBeCloseTo(0);
+      expect(result[1]).toBeCloseTo(0);
+    });
+
+    it("should handle 4096 dimension embedding by slicing to 4095", () => {
+      const embedding = new Array(4096).fill(0.1);
+      const importance = 0.5;
+      const result = vectorWithImportance(embedding, importance);
+
+      expect(result).toHaveLength(4096);
+      expect(getImportance(result)).toBeCloseTo(0.5);
+    });
+
+    it("should properly normalize non-unit vectors", () => {
+      const embedding = [3, 4]; // magnitude = 5
+      const importance = 1;
+      const result = vectorWithImportance(embedding, importance);
+
+      // After normalization: [3/5, 4/5] = [0.6, 0.8]
+      expect(result).toHaveLength(3);
+      expect(result[0]).toBeCloseTo(0.6);
+      expect(result[1]).toBeCloseTo(0.8);
+    });
+  });
+
+  describe("getImportance", () => {
+    it("should correctly extract importance from vector", () => {
+      const vector = vectorWithImportance([0.1, 0.2, 0.3], 0.49);
+      const importance = getImportance(vector);
+
+      expect(importance).toBeCloseTo(0.49);
+    });
+
+    it("should handle zero importance", () => {
+      const vector = [0.1, 0.2, 0.3, 0];
+      const importance = getImportance(vector);
+
+      expect(importance).toBe(1);
+    });
+
+    it("should handle maximum importance", () => {
+      const vector = [0.1, 0.2, 0.3, 0];
+      const importance = getImportance(vector);
+
+      expect(importance).toBe(1);
+    });
+  });
+
+  describe("modifyImportance", () => {
+    it("should modify importance of existing vector", () => {
+      const originalVector = [0.6, 0.8, Math.sqrt(0.75)]; // original importance = 0.25
+      const newImportance = 0.64;
+      const result = modifyImportance(originalVector, newImportance);
+
+      expect(result).toHaveLength(3);
+      expect(result[0]).toBeCloseTo(0.6 * newImportance);
+      expect(result[1]).toBeCloseTo(0.8 * newImportance);
+      expect(getImportance(result)).toBeCloseTo(0.64);
+    });
+
+    it("should handle zero importance modification", () => {
+      const originalVector = [0.6, 0.8, 0.5];
+      const newImportance = 0;
+      const result = modifyImportance(originalVector, newImportance);
+
+      expect(result).toHaveLength(3);
+      expect(result[0]).toBe(0);
+      expect(result[1]).toBe(0);
+      expect(result[2]).toBe(1);
+      expect(getImportance(result)).toBe(0);
+    });
+  });
+
+  describe("vectorWithImportanceDimension", () => {
+    it("should return dimensions + 1 for normal dimensions", () => {
+      expect(vectorWithImportanceDimension(128)).toBe(129);
+      expect(vectorWithImportanceDimension(256)).toBe(257);
+      expect(vectorWithImportanceDimension(512)).toBe(513);
+      expect(vectorWithImportanceDimension(1024)).toBe(1025);
+      expect(vectorWithImportanceDimension(1536)).toBe(1537);
+    });
+
+    it("should return 4096 for 4096 input (respecting global limit)", () => {
+      expect(vectorWithImportanceDimension(4096)).toBe(4096);
+    });
+
+    it("should handle edge cases", () => {
+      expect(vectorWithImportanceDimension(0)).toBe(1);
+      expect(vectorWithImportanceDimension(1)).toBe(2);
+    });
+  });
+
+  describe("round-trip importance testing", () => {
+    const testCases = [
+      { importance: 0.0, tolerance: 0.001 },
+      { importance: 0.1, tolerance: 0.001 },
+      { importance: 0.25, tolerance: 0.001 },
+      { importance: 0.5, tolerance: 0.001 },
+      { importance: 0.75, tolerance: 0.001 },
+      { importance: 1.0, tolerance: 0.001 },
+    ];
+
+    testCases.forEach(({ importance, tolerance }) => {
+      it(`should round-trip importance value ${importance} approximately`, () => {
+        const embedding = [0.6, 0.8]; // unit vector
+        const vectorWithImp = vectorWithImportance(embedding, importance);
+        const retrievedImportance = getImportance(vectorWithImp);
+
+        expect(retrievedImportance).toBeCloseTo(importance, 3);
+        expect(Math.abs(retrievedImportance - importance)).toBeLessThan(
+          tolerance
+        );
+      });
+    });
+
+    it("should round-trip with non-unit vectors", () => {
+      const embedding = [3, 4, 5]; // magnitude = sqrt(50)
+      const originalImportance = 0.36;
+
+      const vectorWithImp = vectorWithImportance(embedding, originalImportance);
+      const retrievedImportance = getImportance(vectorWithImp);
+
+      expect(retrievedImportance).toBeCloseTo(originalImportance, 3);
+    });
+
+    it("should round-trip after modifyImportance", () => {
+      const embedding = [0.1, 0.2, 0.3];
+      const initialImportance = 0.5;
+      const newImportance = 0.8;
+
+      // Create vector with initial importance
+      const vectorWithInitialImp = vectorWithImportance(
+        embedding,
+        initialImportance
+      );
+
+      // Modify importance
+      const vectorWithModifiedImp = modifyImportance(
+        vectorWithInitialImp,
+        newImportance
+      );
+
+      // Retrieve and verify
+      const retrievedImportance = getImportance(vectorWithModifiedImp);
+      expect(retrievedImportance).toBeCloseTo(newImportance, 3);
+    });
+  });
+
+  describe("edge cases and error conditions", () => {
+    it("should handle very small importance values", () => {
+      const embedding = [1, 0];
+      const importance = 1e-10;
+      const result = vectorWithImportance(embedding, importance);
+
+      expect(getImportance(result)).toBeCloseTo(importance);
+    });
+
+    it("should handle very large embeddings", () => {
+      const embedding = new Array(2048).fill(0.1);
+      const importance = 0.234;
+      const result = vectorWithImportance(embedding, importance);
+
+      expect(result).toHaveLength(2049);
+    });
+
+    it("should maintain vector properties after importance weighting", () => {
+      const embedding = [0.6, 0.8]; // unit vector
+      const importance = 0.25;
+      const result = vectorWithImportance(embedding, importance);
+
+      const normalized = normalizeVector(result.slice(0, 2));
+      // The first two components should be normalized versions of original
+      expect(normalized[0]).toBeCloseTo(0.6);
+      expect(normalized[1]).toBeCloseTo(0.8);
+
+      // When used in search (ignoring importance), should behave correctly
+      const searchVec = searchVector(embedding);
+      expect(searchVec[0]).toBeCloseTo(0.6);
+      expect(searchVec[1]).toBeCloseTo(0.8);
+    });
+  });
+});

package/src/component/embeddings/importance.ts

@@ -0,0 +1,75 @@
+/**
+ * This file contains functions for modifying embeddings to include importance.
+ * Terminology is roughly: a "vector" is an "embedding" + importance.
+ * Users pass in embeddings, the tables and vector search deal with vectors.
+ */
+
+/**
+ * For a search, we need to add a 0 to the end of the embedding so we ignore
+ * the weight value.
+ */
+export function searchVector(embedding: number[]) {
+  if (embedding.length === 4096) {
+    return [...embedding.slice(0, 4095), 0];
+  }
+  return [...embedding, 0];
+}
+/**
+ * For an importance of x (0 to 1):
+ * @param embedding - The vector to modify with an importance weight.
+ * @param importance - 0 - 1, where 0 is no importance and 1 is full importance.
+ * @returns The vector with the importance added.
+ */
+export function vectorWithImportance(embedding: number[], importance: number) {
+  /*
+   * Goal: add a weighting that reduces the magnitude of the target vector after
+   * normalization.
+   * 1. Scale the existing vector by importance and add a weight so |v| = 1.
+   * 2. Search with [...embedding, 0].
+   * e.g.:
+   * Say we have an embedding of 2 numbers [.6, .8]
+   * For 50% importance: [.6 * i, .8 * i, sqrt(1 - i^2)] ->  [.3, .4, .866]
+   * For [.6, .8] we used to get 1.0.
+   * Now we get .6*.3 + .8*.4 + 0*.866 = 0.5
+   */
+  // We drop the final dimension if it'd make it larger than 4096.
+  // Unfortunate current limitation of Convex vector search.
+  const vectorToModify = normalizeVector(
+    embedding.length === 4096 ? embedding.slice(0, 4095) : embedding
+  );
+  const scaled = scaleVector(vectorToModify, importance);
+
+  // |embedding| == 1
+  // weight^2 + importance^2(|embedding|^2) == 1
+  // weight = sqrt(1 - importance^2)
+  const weight = Math.sqrt(1 - importance ** 2);
+  return [...scaled, weight];
+}
+
+function normalizeVector(vector: number[]) {
+  const sumOfSquares = vector.reduce((acc, v) => acc + v * v, 0);
+  const magnitude = Math.sqrt(sumOfSquares);
+  return magnitude === 0
+    ? vector.map(() => 0)
+    : vector.map((v) => v / magnitude);
+}
+
+function scaleVector(vector: number[], scale: number) {
+  return vector.map((v) => v * scale);
+}
+
+export function modifyImportance(vector: number[], importance: number) {
+  // Note: we don't need to handle 4096 explicitly here
+  // vectorWithImportance will turn it from 4095 to 4096.
+  const vectorToModify = vector.slice(0, vector.length - 1);
+  return vectorWithImportance(vectorToModify, importance);
+}
+
+export function getImportance(vector: number[]) {
+  return Math.sqrt(1 - vector[vector.length - 1] ** 2);
+}
+
+export function vectorWithImportanceDimension(dimensions: number) {
+  // +1 for the importance weighting, but respect global limit
+  return dimensions === 4096 ? 4096 : dimensions + 1;
+}