Add TurboQuant vector quantization algorithm

packages/ai/src/task/VectorQuantizeTask.ts

@@ -10,11 +10,19 @@ import {
   FromSchema,
   normalizeNumberArray,
   TensorType,
+  turboQuantizeToTypedArray,
   TypedArray,
   TypedArraySchema,
   TypedArraySchemaOptions,
 } from "@workglow/util/schema";
 
+export const QuantizationMethod = {
+  LINEAR: "linear",
+  TURBO: "turbo",
+} as const;
+
+export type QuantizationMethod = (typeof QuantizationMethod)[keyof typeof QuantizationMethod];
+
 const inputSchema = {
   type: "object",
   properties: {
@@ -48,6 +56,21 @@ const inputSchema = {
       description: "Normalize vector before quantization",
       default: true,
     },
+    method: {
+      type: "string",
+      enum: Object.values(QuantizationMethod),
+      title: "Method",
+      description:
+        "Quantization method: 'linear' for simple min-max scaling, 'turbo' for TurboQuant (randomized rotation + optimal scalar quantization, better distortion than linear at the same bit width). Turbo requires an integer targetType (int8, uint8, int16, uint16).",
+      default: QuantizationMethod.LINEAR,
+    },
+    turboSeed: {
+      type: "integer",
+      title: "TurboQuant Seed",
+      description:
+        "Seed for the random rotation in TurboQuant. All vectors in the same collection must use the same seed for similarity search to work.",
+      default: 42,
+    },
   },
   required: ["vector", "targetType"],
   additionalProperties: false,
@@ -117,12 +140,24 @@ export class VectorQuantizeTask extends Task<
   }
 
   override async executeReactive(input: VectorQuantizeTaskInput): Promise<VectorQuantizeTaskOutput> {
-    const { vector, targetType, normalize = true } = input;
+    const {
+      vector,
+      targetType,
+      normalize = true,
+      method = QuantizationMethod.LINEAR,
+      turboSeed = 42,
+    } = input;
     const isArray = Array.isArray(vector);
     const vectors = isArray ? vector : [vector];
     const originalType = this.getVectorType(vectors[0]);
 
-    const quantized = vectors.map((v) => this.vectorQuantize(v, targetType, normalize));
+    let quantized: TypedArray[];
+
+    if (method === QuantizationMethod.TURBO) {
+      quantized = vectors.map((v) => turboQuantizeToTypedArray(v, targetType, turboSeed));
+    } else {
+      quantized = vectors.map((v) => this.vectorQuantize(v, targetType, normalize));
+    }
 
     return {
       vector: isArray ? quantized : quantized[0],

packages/test/src/test/rag/VectorQuantizeTask.test.ts

@@ -229,4 +229,68 @@ describe("VectorQuantizeTask", () => {
     expect(result).toBeDefined();
     expect(result.vector).toBeInstanceOf(Int8Array);
   });
+
+  describe("turbo method", () => {
+    test("should return target TypedArray type directly", async () => {
+      const vector = new Float32Array([1, 2, 3, 4, 5, 6, 7, 8]);
+
+      const result = await vectorQuantize({
+        vector,
+        targetType: TensorType.INT8,
+        method: "turbo",
+        turboSeed: 42,
+      });
+
+      expect(result).toBeDefined();
+      expect(result.vector).toBeInstanceOf(Int8Array);
+      expect(result.targetType).toBe(TensorType.INT8);
+      expect(result.originalType).toBe(TensorType.FLOAT32);
+      expect((result.vector as Int8Array).length).toBe(vector.length);
+    });
+
+    test("should be deterministic for a fixed seed", async () => {
+      const vector = new Float32Array([1, 2, 3, 4, 5, 6, 7, 8]);
+
+      const r1 = await vectorQuantize({
+        vector,
+        targetType: TensorType.INT8,
+        method: "turbo",
+        turboSeed: 99,
+      });
+
+      const r2 = await vectorQuantize({
+        vector,
+        targetType: TensorType.INT8,
+        method: "turbo",
+        turboSeed: 99,
+      });
+
+      const v1 = r1.vector as Int8Array;
+      const v2 = r2.vector as Int8Array;
+      expect(v1.length).toBe(v2.length);
+      for (let i = 0; i < v1.length; i++) {
+        expect(v1[i]).toBe(v2[i]);
+      }
+    });
+
+    test("should handle array of vectors with turbo method", async () => {
+      const vectors = [
+        new Float32Array([1, 2, 3, 4]),
+        new Float32Array([5, 6, 7, 8]),
+      ];
+
+      const result = await vectorQuantize({
+        vector: vectors,
+        targetType: TensorType.INT8,
+        method: "turbo",
+        turboSeed: 42,
+      });
+
+      expect(Array.isArray(result.vector)).toBe(true);
+      const out = result.vector as Int8Array[];
+      expect(out.length).toBe(2);
+      out.forEach((v) => expect(v).toBeInstanceOf(Int8Array));
+      expect(result.targetType).toBe(TensorType.INT8);
+    });
+  });
 });