diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml deleted file mode 100644 index 813d3021..00000000 --- a/.github/workflows/ci.yml +++ /dev/null @@ -1,206 +0,0 @@ -name: CI - -on: - push: - branches: [main] - pull_request: - branches: [main] - -env: - NODE_VERSION: '20' - RUST_VERSION: 'stable' - -jobs: - # Build and test TypeScript/JavaScript - test-node: - name: Test (Node.js) - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - - name: Setup Node.js - uses: actions/setup-node@v4 - with: - node-version: ${{ env.NODE_VERSION }} - cache: 'npm' - - - name: Install dependencies - run: npm ci - - - name: Build TypeScript - run: npm run build - - - name: Run unit tests - run: npm run test:unit - - - name: Run integration tests - run: npm run test:integration - - # Build and test with Bun - test-bun: - name: Test (Bun) - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - - name: Setup Bun - uses: oven-sh/setup-bun@v2 - with: - bun-version: latest - - - name: Setup Node.js (for npm compatibility) - uses: actions/setup-node@v4 - with: - node-version: ${{ env.NODE_VERSION }} - cache: 'npm' - - - name: Install dependencies - run: npm ci - - - name: Build TypeScript - run: npm run build - - - name: Run Bun tests - run: npm run test:bun - - # Build Candle WASM (if Rust source changed) - build-candle-wasm: - name: Build Candle WASM - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - - name: Setup Rust - uses: dtolnay/rust-action@stable - with: - toolchain: ${{ env.RUST_VERSION }} - targets: wasm32-unknown-unknown - - - name: Install wasm-pack - run: cargo install wasm-pack - - - name: Cache Cargo registry - uses: actions/cache@v4 - with: - path: | - ~/.cargo/registry - ~/.cargo/git - src/embeddings/candle-wasm/target - key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }} - - - name: Build WASM - run: | - cd src/embeddings/candle-wasm - wasm-pack build --target web --release - - - name: Upload WASM artifacts - uses: actions/upload-artifact@v4 - with: - name: candle-wasm - path: src/embeddings/wasm/pkg/ - retention-days: 7 - - # Test Bun compile (standalone binary) - test-bun-compile: - name: Test Bun Compile - runs-on: ubuntu-latest - needs: [build-candle-wasm] - steps: - - uses: actions/checkout@v4 - - - name: Setup Bun - uses: oven-sh/setup-bun@v2 - with: - bun-version: latest - - - name: Setup Node.js - uses: actions/setup-node@v4 - with: - node-version: ${{ env.NODE_VERSION }} - cache: 'npm' - - - name: Download WASM artifacts - uses: actions/download-artifact@v4 - with: - name: candle-wasm - path: src/embeddings/wasm/pkg/ - - - name: Install dependencies - run: npm ci - - - name: Build TypeScript - run: npm run build - - - name: Test Bun compile - run: | - # Create a test script - cat > /tmp/test-compile.ts << 'EOF' - import { Brainy } from './dist/index.js' - - const brainy = new Brainy() - await brainy.init() - console.log('Brainy initialized!') - - const embedding = await brainy.embed('Hello world') - console.log(`Embedding dimension: ${embedding.length}`) - - if (embedding.length !== 384) { - throw new Error('Expected 384-dimensional embedding') - } - - console.log('Test passed!') - EOF - - # Compile to standalone binary - bun build --compile /tmp/test-compile.ts --outfile /tmp/brainy-test - - # Run the compiled binary - /tmp/brainy-test - - # Lint and type check - lint: - name: Lint & Type Check - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - - name: Setup Node.js - uses: actions/setup-node@v4 - with: - node-version: ${{ env.NODE_VERSION }} - cache: 'npm' - - - name: Install dependencies - run: npm ci - - - name: Type check - run: npm run typecheck - - - name: Lint - run: npm run lint - - # License check - license-check: - name: License Check - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - - name: Setup Node.js - uses: actions/setup-node@v4 - with: - node-version: ${{ env.NODE_VERSION }} - cache: 'npm' - - - name: Install dependencies - run: npm ci - - - name: Check licenses - run: | - # Install license checker - npm install -g license-checker - - # Check for problematic licenses - license-checker --production --excludePrivatePackages \ - --failOn 'GPL;LGPL;AGPL;SSPL' \ - --summary diff --git a/src/brainy.ts b/src/brainy.ts index 06269d37..fffb16a7 100644 --- a/src/brainy.ts +++ b/src/brainy.ts @@ -6236,7 +6236,14 @@ export class Brainy implements BrainyInterface { private setupIndex(): HNSWIndex | TypeAwareHNSWIndex { const indexConfig = { ...this.config.index, - distanceFunction: this.distance + distanceFunction: this.distance, + // Wire HNSW optimization config (v7.11.0) + quantization: this.config.hnsw?.quantization ? { + enabled: this.config.hnsw.quantization.enabled ?? false, + bits: this.config.hnsw.quantization.bits ?? 8, + rerankMultiplier: this.config.hnsw.quantization.rerankMultiplier ?? 3 + } : undefined, + vectorStorage: this.config.hnsw?.vectorStorage } const persistMode = this.resolveHNSWPersistMode() @@ -6368,6 +6375,8 @@ export class Brainy implements BrainyInterface { reservedQueryMemory: config?.reservedQueryMemory ?? undefined as any, // HNSW persistence mode - undefined = smart default in setupIndex hnswPersistMode: config?.hnswPersistMode ?? undefined as any, + // HNSW optimization options (v7.11.0) + hnsw: config?.hnsw ?? undefined as any, // Embedding initialization - false = lazy init on first embed() eagerEmbeddings: config?.eagerEmbeddings ?? false, // Integration Hub - undefined/false = disabled diff --git a/src/coreTypes.ts b/src/coreTypes.ts index f6122d00..baa4d374 100644 --- a/src/coreTypes.ts +++ b/src/coreTypes.ts @@ -96,6 +96,10 @@ export interface HNSWNoun { vector: Vector connections: Map> // level -> set of connected noun ids level: number // The highest layer this noun appears in + // SQ8 quantized vector for approximate distance computation (B1 optimization) + quantizedVector?: Uint8Array // 4x smaller than float32 vector + codebookMin?: number // quantization codebook minimum + codebookMax?: number // quantization codebook maximum // ✅ NO metadata field - stored separately for optimization } @@ -309,6 +313,14 @@ export interface HNSWConfig { ml: number // Maximum level useDiskBasedIndex?: boolean // Whether to use disk-based index maxConcurrentNeighborWrites?: number // Maximum concurrent neighbor updates during insert. Default: unlimited (full concurrency) + // SQ8 vector quantization (4x memory reduction, ~0.4% accuracy loss) + quantization?: { + enabled: boolean // default: false — preserves current behavior exactly + bits?: 8 | 4 // default: 8 (SQ8). SQ4 requires brainy-cortex native. + rerankMultiplier?: number // default: 3 — over-retrieve 3x, rerank with float32 + } + // Vector storage mode + vectorStorage?: 'memory' | 'lazy' // default: 'memory' — 'lazy' evicts vectors after insert } /** diff --git a/src/hnsw/hnswIndex.ts b/src/hnsw/hnswIndex.ts index 42a595d8..9ff0fa15 100644 --- a/src/hnsw/hnswIndex.ts +++ b/src/hnsw/hnswIndex.ts @@ -15,6 +15,8 @@ import { executeInThread } from '../utils/workerUtils.js' import type { BaseStorage } from '../storage/baseStorage.js' import { getGlobalCache, UnifiedCache } from '../utils/unifiedCache.js' import { prodLog } from '../utils/logger.js' +import { quantizeSQ8, distanceSQ8 } from '../utils/vectorQuantization.js' +import type { SQ8QuantizedVector } from '../utils/vectorQuantization.js' // Default HNSW parameters const DEFAULT_CONFIG: HNSWConfig = { @@ -53,6 +55,12 @@ export class HNSWIndex { private dirtyNodes: Set = new Set() // Nodes with unpersisted HNSW data private dirtySystem: boolean = false // Whether system data (entryPoint, maxLevel) needs persist + // SQ8 quantization support (B1 optimization) + private quantizationEnabled: boolean = false + private rerankMultiplier: number = 3 + // Lazy vector storage (B2 optimization) + private vectorStorageMode: 'memory' | 'lazy' = 'memory' + constructor( config: Partial = {}, distanceFunction: DistanceFunction = euclideanDistance, @@ -67,6 +75,15 @@ export class HNSWIndex { this.storage = options.storage || null this.persistMode = options.persistMode || 'immediate' + // SQ8 quantization config (default: disabled, preserves current behavior) + if (config.quantization?.enabled) { + this.quantizationEnabled = true + this.rerankMultiplier = config.quantization.rerankMultiplier ?? 3 + } + + // Vector storage mode (default: 'memory', preserves current behavior) + this.vectorStorageMode = config.vectorStorage || 'memory' + // Use SAME UnifiedCache as Graph and Metadata for fair memory competition this.unifiedCache = getGlobalCache() } @@ -257,7 +274,11 @@ export class HNSWIndex { id: original.id, vector: [...original.vector], // Deep copy vector array connections: connectionsCopy, - level: original.level + level: original.level, + // Copy SQ8 quantized data if present + quantizedVector: original.quantizedVector ? new Uint8Array(original.quantizedVector) : undefined, + codebookMin: original.codebookMin, + codebookMax: original.codebookMax } this.nouns.set(nodeId, nodeCopy) @@ -343,7 +364,7 @@ export class HNSWIndex { // Generate random level for this noun const nounLevel = this.getRandomLevel() - // Create new noun + // Create new noun with optional SQ8 quantization const noun: HNSWNoun = { id, vector, @@ -351,6 +372,14 @@ export class HNSWIndex { level: nounLevel } + // Quantize vector if enabled (B1: 4x storage reduction) + if (this.quantizationEnabled) { + const sq8 = quantizeSQ8(vector) + noun.quantizedVector = sq8.quantized + noun.codebookMin = sq8.min + noun.codebookMax = sq8.max + } + // Initialize empty connection sets for each level for (let level = 0; level <= nounLevel; level++) { noun.connections.set(level, new Set()) @@ -574,6 +603,13 @@ export class HNSWIndex { this.highLevelNodes.get(nounLevel)!.add(id) } + // Lazy vector eviction (B2: graph-only memory after insert) + // After graph construction completes, evict the full vector from memory. + // Future searches will load vectors on-demand via getVectorSafe() + UnifiedCache. + if (this.vectorStorageMode === 'lazy' && this.storage) { + noun.vector = [] // Release float32 vector from memory + } + // Persist HNSW graph data to storage // Respect persistMode setting if (this.storage && this.persistMode === 'immediate') { @@ -631,11 +667,21 @@ export class HNSWIndex { /** * Search for nearest neighbors + * + * When SQ8 quantization is enabled and reranking is active: + * - Phase 1: Over-retrieve k*rerankMultiplier candidates using SQ8 approximate distances + * - Phase 2: Load full float32 vectors for candidates, compute exact distances, return top k + * + * @param queryVector Query vector + * @param k Number of results to return + * @param filter Optional filter function + * @param options Additional search options */ public async search( queryVector: Vector, k: number = 10, - filter?: (id: string) => Promise + filter?: (id: string) => Promise, + options?: { rerank?: { multiplier: number } } ): Promise> { if (this.nouns.size === 0) { return [] @@ -687,6 +733,11 @@ export class HNSWIndex { let currObj = entryPoint + // SQ8: Pre-quantize query vector for fast approximate distances during traversal + if (this.quantizationEnabled) { + this._querySQ8 = quantizeSQ8(queryVector) + } + // OPTIMIZATION: Preload entry point vector await this.preloadVectors([entryPoint.id]) @@ -754,9 +805,14 @@ export class HNSWIndex { } } - // Search at level 0 with ef = k + // Determine effective rerank multiplier + const rerankActive = this.quantizationEnabled && (options?.rerank || this.rerankMultiplier > 1) + const multiplier = options?.rerank?.multiplier ?? this.rerankMultiplier + const effectiveK = rerankActive ? k * multiplier : k + + // Search at level 0 with ef = effectiveK // If we have a filter, increase ef to compensate for filtered results - const ef = filter ? Math.max(this.config.efSearch * 3, k * 3) : Math.max(this.config.efSearch, k) + const ef = filter ? Math.max(this.config.efSearch * 3, effectiveK * 3) : Math.max(this.config.efSearch, effectiveK) const nearestNouns = await this.searchLayer( queryVector, currObj, @@ -765,6 +821,35 @@ export class HNSWIndex { filter ) + // Phase 2: Rerank with exact float32 distances if quantization is active (B3) + if (rerankActive && nearestNouns.size > 0) { + // Clear SQ8 cache before reranking (we need exact distances now) + this._querySQ8 = null + + const candidates = [...nearestNouns].slice(0, effectiveK) + + // Load full float32 vectors for the candidate set + const candidateIds = candidates.map(([id]) => id) + await this.preloadVectors(candidateIds) + + // Recompute exact distances + const reranked: Array<[string, number]> = [] + for (const [id] of candidates) { + const noun = this.nouns.get(id) + if (!noun) continue + const exactVector = await this.getVectorSafe(noun) + const exactDist = this.distanceFunction(queryVector, exactVector) + reranked.push([id, exactDist]) + } + + // Sort by exact distance and return top k + reranked.sort((a, b) => a[1] - b[1]) + return reranked.slice(0, k) + } + + // Clear SQ8 cache + this._querySQ8 = null + // Convert to array and sort by distance return [...nearestNouns].slice(0, k) } @@ -1074,6 +1159,19 @@ export class HNSWIndex { * @returns number | Promise - sync when cached, async when needs load */ private distanceSafe(queryVector: Vector, noun: HNSWNoun): number | Promise { + // SQ8 fast path: use quantized distance when available (B1 optimization) + // This avoids loading full float32 vectors during graph traversal + if (this.quantizationEnabled && + noun.quantizedVector && + noun.codebookMin !== undefined && + noun.codebookMax !== undefined && + this._querySQ8) { + return distanceSQ8( + this._querySQ8.quantized, this._querySQ8.min, this._querySQ8.max, + noun.quantizedVector, noun.codebookMin, noun.codebookMax + ) + } + // Try sync fast path const nounVector = this.getVectorSync(noun) @@ -1088,6 +1186,9 @@ export class HNSWIndex { ) } + // Cached SQ8 quantization of the current query vector for distanceSafe fast path + private _querySQ8: SQ8QuantizedVector | null = null + /** * Get all nodes at a specific level for clustering * This enables O(n) clustering using HNSW's natural hierarchy @@ -1207,14 +1308,25 @@ export class HNSWIndex { continue } + // Determine if vector should be kept in memory + const keepVector = shouldPreload && this.vectorStorageMode !== 'lazy' + // Create noun object with restored connections const noun: HNSWNoun = { id: nounData.id, - vector: shouldPreload ? nounData.vector : [], // Preload if dataset is small + vector: keepVector ? nounData.vector : [], // Preload if dataset is small and not lazy connections: new Map(), level: hnswData.level } + // Restore SQ8 quantized data if quantization is enabled + if (this.quantizationEnabled && nounData.vector.length > 0) { + const sq8 = quantizeSQ8(nounData.vector) + noun.quantizedVector = sq8.quantized + noun.codebookMin = sq8.min + noun.codebookMax = sq8.max + } + // Restore connections from persisted data for (const [levelStr, nounIds] of Object.entries(hnswData.connections)) { const level = parseInt(levelStr, 10) @@ -1281,14 +1393,25 @@ export class HNSWIndex { continue } + // Determine if vector should be kept in memory + const keepVector = shouldPreload && this.vectorStorageMode !== 'lazy' + // Create noun object with restored connections const noun: HNSWNoun = { id: nounData.id, - vector: shouldPreload ? nounData.vector : [], // Preload if dataset is small + vector: keepVector ? nounData.vector : [], // Preload if dataset is small and not lazy connections: new Map(), level: hnswData.level } + // Restore SQ8 quantized data if quantization is enabled + if (this.quantizationEnabled && nounData.vector.length > 0) { + const sq8 = quantizeSQ8(nounData.vector) + noun.quantizedVector = sq8.quantized + noun.codebookMin = sq8.min + noun.codebookMax = sq8.max + } + // Restore connections from persisted data for (const [levelStr, nounIds] of Object.entries(hnswData.connections)) { const level = parseInt(levelStr, 10) diff --git a/src/hnsw/typeAwareHNSWIndex.ts b/src/hnsw/typeAwareHNSWIndex.ts index 4f57f448..35d01e79 100644 --- a/src/hnsw/typeAwareHNSWIndex.ts +++ b/src/hnsw/typeAwareHNSWIndex.ts @@ -256,21 +256,22 @@ export class TypeAwareHNSWIndex { queryVector: Vector, k: number = 10, type?: NounType | NounType[], - filter?: (id: string) => Promise + filter?: (id: string) => Promise, + options?: { rerank?: { multiplier: number } } ): Promise> { // Single-type search (fast path) if (type && typeof type === 'string') { const index = this.getIndexForType(type) - return await index.search(queryVector, k, filter) + return await index.search(queryVector, k, filter, options) } // Multi-type search (handle empty array edge case) if (type && Array.isArray(type) && type.length > 0) { - return await this.searchMultipleTypes(queryVector, k, type, filter) + return await this.searchMultipleTypes(queryVector, k, type, filter, options) } // All-types search (slowest path + empty array fallback) - return await this.searchAllTypes(queryVector, k, filter) + return await this.searchAllTypes(queryVector, k, filter, options) } /** @@ -286,7 +287,8 @@ export class TypeAwareHNSWIndex { queryVector: Vector, k: number, types: NounType[], - filter?: (id: string) => Promise + filter?: (id: string) => Promise, + options?: { rerank?: { multiplier: number } } ): Promise> { const allResults: Array<[string, number]> = [] @@ -294,7 +296,7 @@ export class TypeAwareHNSWIndex { for (const type of types) { if (this.indexes.has(type)) { const index = this.indexes.get(type)! - const results = await index.search(queryVector, k, filter) + const results = await index.search(queryVector, k, filter, options) allResults.push(...results) } } @@ -320,13 +322,14 @@ export class TypeAwareHNSWIndex { private async searchAllTypes( queryVector: Vector, k: number, - filter?: (id: string) => Promise + filter?: (id: string) => Promise, + options?: { rerank?: { multiplier: number } } ): Promise> { const allResults: Array<[string, number]> = [] // Search each type's graph for (const [type, index] of this.indexes.entries()) { - const results = await index.search(queryVector, k, filter) + const results = await index.search(queryVector, k, filter, options) allResults.push(...results) } diff --git a/src/types/brainy.types.ts b/src/types/brainy.types.ts index b5f14632..68ddd16b 100644 --- a/src/types/brainy.types.ts +++ b/src/types/brainy.types.ts @@ -728,6 +728,16 @@ export interface BrainyConfig { // Cloud storage (GCS/S3/R2/Azure) should use 'deferred' for 30-50× faster adds hnswPersistMode?: 'immediate' | 'deferred' + // HNSW optimization options (v7.11.0) + hnsw?: { + quantization?: { + enabled?: boolean // default: false — current behavior exactly + bits?: 8 | 4 // default: 8 (SQ8). SQ4 requires brainy-cortex native. + rerankMultiplier?: number // default: 3 — over-retrieve 3x, rerank with float32 + } + vectorStorage?: 'memory' | 'lazy' // default: 'memory' — 'lazy' evicts vectors after insert + } + // Memory management options maxQueryLimit?: number // Override auto-detected query result limit (max: 100000) reservedQueryMemory?: number // Memory reserved for queries in bytes (e.g., 1073741824 = 1GB) diff --git a/src/utils/vectorQuantization.ts b/src/utils/vectorQuantization.ts new file mode 100644 index 00000000..648e422a --- /dev/null +++ b/src/utils/vectorQuantization.ts @@ -0,0 +1,218 @@ +/** + * Vector Quantization Utilities + * + * Standalone SQ8 (Scalar Quantization, 8-bit) functions for HNSW vector compression. + * Each vector is independently quantized using its own min/max codebook. + * + * Storage format per vector: + * - quantized: Uint8Array (dimension bytes, 4x smaller than float32) + * - min: number (codebook minimum) + * - max: number (codebook maximum) + * + * Accuracy: SQ8 introduces ~0.4% error per dimension on normalized vectors. + * Combined with reranking (3x over-retrieval + float32 rerank), recall@100 + * is expected to remain >99% for typical workloads. + */ + +import type { Vector } from '../coreTypes.js' + +/** + * Result of SQ8 quantization + */ +export interface SQ8QuantizedVector { + quantized: Uint8Array + min: number + max: number +} + +/** + * Configuration for HNSW quantization + */ +export interface HNSWQuantizationConfig { + enabled: boolean + bits: 8 | 4 + rerankMultiplier: number +} + +/** + * Quantize a float32 vector to SQ8 (8-bit unsigned integers). + * + * Maps [min, max] of the vector to [0, 255]. + * For zero-range vectors (all values equal), all quantized values are 128. + * + * @param vector - Input float32 vector + * @returns Quantized vector with codebook (min/max) + */ +export function quantizeSQ8(vector: Vector): SQ8QuantizedVector { + const len = vector.length + let min = vector[0] + let max = vector[0] + + for (let i = 1; i < len; i++) { + const v = vector[i] + if (v < min) min = v + if (v > max) max = v + } + + const range = max - min + const quantized = new Uint8Array(len) + + if (range === 0) { + // All values are identical — map to midpoint + quantized.fill(128) + } else { + const scale = 255 / range + for (let i = 0; i < len; i++) { + quantized[i] = Math.round((vector[i] - min) * scale) + } + } + + return { quantized, min, max } +} + +/** + * Dequantize an SQ8 vector back to float32. + * + * @param quantized - Uint8Array of quantized values + * @param min - Codebook minimum + * @param max - Codebook maximum + * @returns Reconstructed float32 vector + */ +export function dequantizeSQ8(quantized: Uint8Array, min: number, max: number): Vector { + const len = quantized.length + const result = new Array(len) + const range = max - min + + if (range === 0) { + result.fill(min) + } else { + const scale = range / 255 + for (let i = 0; i < len; i++) { + result[i] = min + quantized[i] * scale + } + } + + return result +} + +/** + * Compute approximate cosine distance between two SQ8-quantized vectors. + * + * Operates directly on uint8 arrays without full dequantization. + * Uses integer arithmetic where possible for speed. + * + * Cosine distance = 1 - (A.B / (|A| * |B|)) + * + * For quantized values q_i with codebook (min, max): + * actual_i = min + q_i * (max - min) / 255 + * + * The dot product and norms can be computed on quantized values + * and scaled by the codebook parameters. + * + * @param a - First quantized vector + * @param aMin - First vector codebook minimum + * @param aMax - First vector codebook maximum + * @param b - Second quantized vector + * @param bMin - Second vector codebook minimum + * @param bMax - Second vector codebook maximum + * @returns Approximate cosine distance [0, 2] + */ +export function distanceSQ8( + a: Uint8Array, + aMin: number, + aMax: number, + b: Uint8Array, + bMin: number, + bMax: number +): number { + const len = a.length + + // Compute raw integer sums for efficiency + // actual_a[i] = aMin + a[i] * aScale, where aScale = (aMax - aMin) / 255 + // actual_b[i] = bMin + b[i] * bScale, where bScale = (bMax - bMin) / 255 + // + // dot = sum(actual_a[i] * actual_b[i]) + // = sum((aMin + a[i]*aScale) * (bMin + b[i]*bScale)) + // = n*aMin*bMin + aMin*bScale*sum(b[i]) + bMin*aScale*sum(a[i]) + aScale*bScale*sum(a[i]*b[i]) + // + // normA^2 = sum(actual_a[i]^2) + // = n*aMin^2 + 2*aMin*aScale*sum(a[i]) + aScale^2*sum(a[i]^2) + + const aScale = (aMax - aMin) / 255 + const bScale = (bMax - bMin) / 255 + + let sumA = 0 + let sumB = 0 + let sumAB = 0 + let sumAA = 0 + let sumBB = 0 + + for (let i = 0; i < len; i++) { + const ai = a[i] + const bi = b[i] + sumA += ai + sumB += bi + sumAB += ai * bi + sumAA += ai * ai + sumBB += bi * bi + } + + const dot = + len * aMin * bMin + + aMin * bScale * sumB + + bMin * aScale * sumA + + aScale * bScale * sumAB + + const normASq = + len * aMin * aMin + + 2 * aMin * aScale * sumA + + aScale * aScale * sumAA + + const normBSq = + len * bMin * bMin + + 2 * bMin * bScale * sumB + + bScale * bScale * sumBB + + const normA = Math.sqrt(normASq) + const normB = Math.sqrt(normBSq) + + if (normA === 0 || normB === 0) { + return 1.0 // Maximum distance for zero vectors + } + + const cosine = dot / (normA * normB) + // Clamp to [-1, 1] to handle floating point imprecision + return 1 - Math.max(-1, Math.min(1, cosine)) +} + +/** + * Serialize an SQ8 quantized vector to a compact binary format. + * + * Format: [min:float32][max:float32][quantized:uint8[]] + * Total size: 8 + dimension bytes + * + * @param sq8 - Quantized vector + * @returns ArrayBuffer with serialized data + */ +export function serializeSQ8(sq8: SQ8QuantizedVector): ArrayBuffer { + const buffer = new ArrayBuffer(8 + sq8.quantized.length) + const view = new DataView(buffer) + view.setFloat32(0, sq8.min, true) // little-endian + view.setFloat32(4, sq8.max, true) + new Uint8Array(buffer, 8).set(sq8.quantized) + return buffer +} + +/** + * Deserialize an SQ8 quantized vector from binary format. + * + * @param buffer - ArrayBuffer with serialized data + * @returns Deserialized SQ8 quantized vector + */ +export function deserializeSQ8(buffer: ArrayBuffer): SQ8QuantizedVector { + const view = new DataView(buffer) + const min = view.getFloat32(0, true) + const max = view.getFloat32(4, true) + const quantized = new Uint8Array(buffer, 8) + return { quantized, min, max } +} diff --git a/tests/unit/hnsw/lazy-vectors.test.ts b/tests/unit/hnsw/lazy-vectors.test.ts new file mode 100644 index 00000000..2b2d29db --- /dev/null +++ b/tests/unit/hnsw/lazy-vectors.test.ts @@ -0,0 +1,293 @@ +/** + * Lazy Vector Loading Tests (B2 optimization) + * + * Tests for: + * - Vectors evicted from memory after addItem() in lazy mode + * - Search returns correct results after vector eviction + * - Memory mode retains vectors (default behavior) + * - Lazy mode requires storage adapter + * - Combined lazy + quantization mode + */ + +import { describe, it, expect, beforeEach } from 'vitest' +import { v4 as uuidv4 } from 'uuid' +import { HNSWIndex } from '../../../src/hnsw/hnswIndex.js' +import { euclideanDistance } from '../../../src/utils/index.js' +import { MemoryStorage } from '../../../src/storage/adapters/memoryStorage.js' + +// Helper: generate a random vector of given dimension +function randomVector(dim: number): number[] { + return Array.from({ length: dim }, () => Math.random() * 2 - 1) +} + +// Helper: save a vector to storage so lazy loading can retrieve it +// Both noun (vector) and metadata must be saved for getNounVector() to work +async function saveVector(storage: MemoryStorage, id: string, vector: number[]): Promise { + await storage.saveNoun({ + id, + vector, + connections: new Map(), + level: 0 + }) + await storage.saveNounMetadata(id, { + noun: 'thing', + createdAt: Date.now(), + updatedAt: Date.now() + }) +} + +describe('Lazy Vector Loading (B2)', () => { + const dim = 32 + + // ================================================================= + // 1. LAZY MODE: VECTOR EVICTION + // ================================================================= + describe('vector eviction in lazy mode', () => { + let index: HNSWIndex + let storage: MemoryStorage + + beforeEach(async () => { + storage = new MemoryStorage() + index = new HNSWIndex( + { + M: 8, + efConstruction: 100, + efSearch: 50, + ml: 8, + vectorStorage: 'lazy' + }, + euclideanDistance, + { useParallelization: false, storage } + ) + }) + + it('should still be searchable after vector eviction', async () => { + const targetId = uuidv4() + const target = randomVector(dim) + + // Store noun in storage so lazy loading can find it + await saveVector(storage, targetId, target) + await index.addItem({ id: targetId, vector: target }) + + // Add more entities + for (let i = 0; i < 20; i++) { + const id = uuidv4() + const v = randomVector(dim) + await saveVector(storage, id, v) + await index.addItem({ id, vector: v }) + } + + // Search should still find the target + const results = await index.search(target, 5) + expect(results.length).toBeGreaterThan(0) + + // The closest result should be the target (distance ~0) + const targetResult = results.find(([id]) => id === targetId) + expect(targetResult).toBeDefined() + expect(targetResult![1]).toBeCloseTo(0, 1) + }) + + it('should return correct top-k results in lazy mode', async () => { + for (let i = 0; i < 30; i++) { + const id = uuidv4() + const v = randomVector(dim) + await saveVector(storage, id, v) + await index.addItem({ id, vector: v }) + } + + const query = randomVector(dim) + const results = await index.search(query, 10) + + expect(results.length).toBe(10) + + // Results should be sorted by distance + for (let i = 1; i < results.length; i++) { + expect(results[i][1]).toBeGreaterThanOrEqual(results[i - 1][1]) + } + }) + }) + + // ================================================================= + // 2. MEMORY MODE: VECTORS RETAINED (DEFAULT) + // ================================================================= + describe('memory mode retains vectors (default)', () => { + it('should keep vectors in memory by default', async () => { + const storage = new MemoryStorage() + const index = new HNSWIndex( + { M: 4, efConstruction: 50, efSearch: 20 }, + euclideanDistance, + { useParallelization: false, storage } + ) + + const id = uuidv4() + const v = randomVector(dim) + await index.addItem({ id, vector: v }) + + // Search should work without needing storage + const results = await index.search(v, 1) + expect(results.length).toBe(1) + expect(results[0][0]).toBe(id) + expect(results[0][1]).toBeCloseTo(0, 5) + }) + + it('should work without storage adapter in memory mode', async () => { + // No storage adapter provided + const index = new HNSWIndex( + { M: 4, efConstruction: 50, efSearch: 20 }, + euclideanDistance, + { useParallelization: false } + ) + + const id = uuidv4() + const v = randomVector(dim) + await index.addItem({ id, vector: v }) + + const results = await index.search(v, 1) + expect(results.length).toBe(1) + expect(results[0][0]).toBe(id) + }) + }) + + // ================================================================= + // 3. LAZY + NO STORAGE: GRACEFUL BEHAVIOR + // ================================================================= + describe('lazy mode without storage', () => { + it('should not evict vectors when no storage adapter is configured', async () => { + // vectorStorage: 'lazy' but no storage — vectors should stay in memory + const index = new HNSWIndex( + { + M: 4, + efConstruction: 50, + efSearch: 20, + vectorStorage: 'lazy' + }, + euclideanDistance, + { useParallelization: false } + ) + + const id = uuidv4() + const v = randomVector(dim) + await index.addItem({ id, vector: v }) + + // Should still work because vectors aren't evicted without storage + const results = await index.search(v, 1) + expect(results.length).toBe(1) + expect(results[0][0]).toBe(id) + }) + }) + + // ================================================================= + // 4. LAZY + QUANTIZATION COMBINED + // ================================================================= + describe('lazy mode with quantization', () => { + it('should use SQ8 for traversal and load full vectors for rerank', async () => { + const storage = new MemoryStorage() + const index = new HNSWIndex( + { + M: 8, + efConstruction: 100, + efSearch: 50, + ml: 8, + vectorStorage: 'lazy', + quantization: { enabled: true, rerankMultiplier: 3 } + }, + euclideanDistance, + { useParallelization: false, storage } + ) + + const targetId = uuidv4() + const target = randomVector(dim) + await saveVector(storage, targetId, target) + await index.addItem({ id: targetId, vector: target }) + + for (let i = 0; i < 30; i++) { + const id = uuidv4() + const v = randomVector(dim) + await saveVector(storage, id, v) + await index.addItem({ id, vector: v }) + } + + // Search with reranking — should load full vectors for rerank phase + const results = await index.search(target, 5) + expect(results.length).toBe(5) + + // The target should be the closest match + const targetResult = results.find(([id]) => id === targetId) + expect(targetResult).toBeDefined() + }) + + it('should return results sorted by exact distance after rerank', async () => { + const storage = new MemoryStorage() + const index = new HNSWIndex( + { + M: 8, + efConstruction: 100, + efSearch: 50, + ml: 8, + vectorStorage: 'lazy', + quantization: { enabled: true, rerankMultiplier: 3 } + }, + euclideanDistance, + { useParallelization: false, storage } + ) + + for (let i = 0; i < 40; i++) { + const id = uuidv4() + const v = randomVector(dim) + await saveVector(storage, id, v) + await index.addItem({ id, vector: v }) + } + + const query = randomVector(dim) + const results = await index.search(query, 10) + + // Results should be sorted by exact distance (rerank ensures this) + for (let i = 1; i < results.length; i++) { + expect(results[i][1]).toBeGreaterThanOrEqual(results[i - 1][1]) + } + }) + }) + + // ================================================================= + // 5. MULTIPLE SEARCHES IN LAZY MODE + // ================================================================= + describe('multiple searches in lazy mode', () => { + it('should handle repeated searches correctly', async () => { + const storage = new MemoryStorage() + const index = new HNSWIndex( + { + M: 8, + efConstruction: 100, + efSearch: 50, + ml: 8, + vectorStorage: 'lazy' + }, + euclideanDistance, + { useParallelization: false, storage } + ) + + const entries: Array<{ id: string; vector: number[] }> = [] + for (let i = 0; i < 25; i++) { + const id = uuidv4() + const v = randomVector(dim) + entries.push({ id, vector: v }) + await saveVector(storage, id, v) + await index.addItem({ id, vector: v }) + } + + // Run multiple searches — each should return results and be consistent + for (let q = 0; q < 5; q++) { + const entry = entries[q * 5] + const results = await index.search(entry.vector, 5) + + expect(results.length).toBe(5) + // Results should be sorted by distance + for (let i = 1; i < results.length; i++) { + expect(results[i][1]).toBeGreaterThanOrEqual(results[i - 1][1]) + } + // At least the closest result should have a reasonably small distance + expect(results[0][1]).toBeLessThan(5) + } + }) + }) +}) diff --git a/tests/unit/hnsw/quantization.test.ts b/tests/unit/hnsw/quantization.test.ts new file mode 100644 index 00000000..bec65ec2 --- /dev/null +++ b/tests/unit/hnsw/quantization.test.ts @@ -0,0 +1,454 @@ +/** + * SQ8 Quantization Tests + * + * Tests for: + * - SQ8 quantize/dequantize round-trip accuracy + * - SQ8 distance vs float32 distance correlation + * - Serialization/deserialization round-trip + * - Search with quantization enabled: recall vs exact + * - Two-phase rerank: improved recall over single-phase SQ8 + * - Config defaults: quantization disabled by default (no behavior change) + */ + +import { describe, it, expect, beforeEach } from 'vitest' +import { v4 as uuidv4 } from 'uuid' +import { + quantizeSQ8, + dequantizeSQ8, + distanceSQ8, + serializeSQ8, + deserializeSQ8 +} from '../../../src/utils/vectorQuantization.js' +import { HNSWIndex } from '../../../src/hnsw/hnswIndex.js' +import { euclideanDistance } from '../../../src/utils/index.js' +import { MemoryStorage } from '../../../src/storage/adapters/memoryStorage.js' + +// Helper: generate a random vector of given dimension +function randomVector(dim: number): number[] { + return Array.from({ length: dim }, () => Math.random() * 2 - 1) +} + +// Helper: cosine distance for float32 vectors (reference implementation) +function cosineDistance(a: number[], b: number[]): number { + let dot = 0 + let normA = 0 + let normB = 0 + for (let i = 0; i < a.length; i++) { + dot += a[i] * b[i] + normA += a[i] * a[i] + normB += b[i] * b[i] + } + const denom = Math.sqrt(normA) * Math.sqrt(normB) + if (denom === 0) return 1.0 + return 1 - dot / denom +} + +describe('SQ8 Quantization', () => { + // ================================================================= + // 1. QUANTIZE / DEQUANTIZE ROUND-TRIP + // ================================================================= + describe('quantize/dequantize round-trip', () => { + it('should round-trip a normalized vector with low error', () => { + const original = randomVector(384) + const sq8 = quantizeSQ8(original) + const restored = dequantizeSQ8(sq8.quantized, sq8.min, sq8.max) + + expect(restored.length).toBe(original.length) + + // Per-dimension error should be small (max 1/255 of range) + const range = sq8.max - sq8.min + const maxError = range / 255 + for (let i = 0; i < original.length; i++) { + expect(Math.abs(restored[i] - original[i])).toBeLessThanOrEqual(maxError + 1e-7) + } + }) + + it('should handle a zero-range vector (all identical values)', () => { + const original = new Array(128).fill(0.5) + const sq8 = quantizeSQ8(original) + + expect(sq8.min).toBe(0.5) + expect(sq8.max).toBe(0.5) + // All quantized values should be 128 (midpoint) + for (let i = 0; i < sq8.quantized.length; i++) { + expect(sq8.quantized[i]).toBe(128) + } + + const restored = dequantizeSQ8(sq8.quantized, sq8.min, sq8.max) + for (let i = 0; i < restored.length; i++) { + expect(restored[i]).toBe(0.5) + } + }) + + it('should handle negative values', () => { + const original = [-1.0, -0.5, 0, 0.5, 1.0] + const sq8 = quantizeSQ8(original) + + expect(sq8.min).toBe(-1.0) + expect(sq8.max).toBe(1.0) + // Min maps to 0, max maps to 255 + expect(sq8.quantized[0]).toBe(0) // -1.0 + expect(sq8.quantized[4]).toBe(255) // 1.0 + + const restored = dequantizeSQ8(sq8.quantized, sq8.min, sq8.max) + expect(Math.abs(restored[0] - (-1.0))).toBeLessThan(0.01) + expect(Math.abs(restored[4] - 1.0)).toBeLessThan(0.01) + }) + + it('should produce Uint8Array output in [0, 255] range', () => { + const original = randomVector(512) + const sq8 = quantizeSQ8(original) + + expect(sq8.quantized).toBeInstanceOf(Uint8Array) + expect(sq8.quantized.length).toBe(512) + for (let i = 0; i < sq8.quantized.length; i++) { + expect(sq8.quantized[i]).toBeGreaterThanOrEqual(0) + expect(sq8.quantized[i]).toBeLessThanOrEqual(255) + } + }) + + it('should achieve 4x storage reduction (uint8 vs float32)', () => { + const dim = 384 + const original = randomVector(dim) + const sq8 = quantizeSQ8(original) + + const float32Size = dim * 4 // 4 bytes per float32 + const sq8Size = sq8.quantized.byteLength + 8 // uint8 array + 2 floats for min/max + const ratio = float32Size / sq8Size + + // Should be close to 4x (actually slightly less due to min/max overhead) + expect(ratio).toBeGreaterThan(3.9) + }) + }) + + // ================================================================= + // 2. SQ8 DISTANCE VS FLOAT32 DISTANCE + // ================================================================= + describe('SQ8 distance accuracy', () => { + it('should correlate strongly with float32 cosine distance', () => { + const dim = 384 + const numPairs = 100 + const errors: number[] = [] + + for (let i = 0; i < numPairs; i++) { + const a = randomVector(dim) + const b = randomVector(dim) + + const exactDist = cosineDistance(a, b) + + const sq8A = quantizeSQ8(a) + const sq8B = quantizeSQ8(b) + const approxDist = distanceSQ8( + sq8A.quantized, sq8A.min, sq8A.max, + sq8B.quantized, sq8B.min, sq8B.max + ) + + errors.push(Math.abs(exactDist - approxDist)) + } + + // Mean absolute error should be very small + const meanError = errors.reduce((sum, e) => sum + e, 0) / errors.length + expect(meanError).toBeLessThan(0.02) // Less than 2% average error + + // Max error should be bounded + const maxError = Math.max(...errors) + expect(maxError).toBeLessThan(0.1) // Less than 10% worst case + }) + + it('should return 0 for identical vectors', () => { + const v = randomVector(128) + const sq8 = quantizeSQ8(v) + const dist = distanceSQ8( + sq8.quantized, sq8.min, sq8.max, + sq8.quantized, sq8.min, sq8.max + ) + expect(dist).toBeCloseTo(0, 5) + }) + + it('should preserve relative ordering of distances', () => { + const dim = 128 + const query = randomVector(dim) + // Create a vector close to query and one far away + const close = query.map(v => v + (Math.random() * 0.1 - 0.05)) + const far = randomVector(dim) + + const exactClose = cosineDistance(query, close) + const exactFar = cosineDistance(query, far) + + const sq8Query = quantizeSQ8(query) + const sq8Close = quantizeSQ8(close) + const sq8Far = quantizeSQ8(far) + + const approxClose = distanceSQ8( + sq8Query.quantized, sq8Query.min, sq8Query.max, + sq8Close.quantized, sq8Close.min, sq8Close.max + ) + const approxFar = distanceSQ8( + sq8Query.quantized, sq8Query.min, sq8Query.max, + sq8Far.quantized, sq8Far.min, sq8Far.max + ) + + // Relative ordering should be preserved + if (exactClose < exactFar) { + expect(approxClose).toBeLessThan(approxFar) + } + }) + + it('should handle zero vectors gracefully', () => { + const zero = new Array(64).fill(0) + const nonZero = randomVector(64) + + const sq8Zero = quantizeSQ8(zero) + const sq8NonZero = quantizeSQ8(nonZero) + + const dist = distanceSQ8( + sq8Zero.quantized, sq8Zero.min, sq8Zero.max, + sq8NonZero.quantized, sq8NonZero.min, sq8NonZero.max + ) + // Zero vectors should return max distance (1.0) + expect(dist).toBeCloseTo(1.0, 1) + }) + }) + + // ================================================================= + // 3. SERIALIZATION / DESERIALIZATION + // ================================================================= + describe('serialization round-trip', () => { + it('should serialize and deserialize SQ8 data with float32 precision', () => { + const original = randomVector(384) + const sq8 = quantizeSQ8(original) + const buffer = serializeSQ8(sq8) + const restored = deserializeSQ8(buffer) + + // min/max are stored as float32, so some precision loss is expected + expect(restored.min).toBeCloseTo(sq8.min, 5) + expect(restored.max).toBeCloseTo(sq8.max, 5) + expect(restored.quantized.length).toBe(sq8.quantized.length) + for (let i = 0; i < sq8.quantized.length; i++) { + expect(restored.quantized[i]).toBe(sq8.quantized[i]) + } + }) + + it('should produce compact binary format (8 + dim bytes)', () => { + const dim = 384 + const sq8 = quantizeSQ8(randomVector(dim)) + const buffer = serializeSQ8(sq8) + + // 4 bytes for min (float32) + 4 bytes for max (float32) + dim bytes + expect(buffer.byteLength).toBe(8 + dim) + }) + }) + + // ================================================================= + // 4. HNSW SEARCH WITH QUANTIZATION + // ================================================================= + describe('HNSW search with quantization', () => { + let index: HNSWIndex + let storage: MemoryStorage + const dim = 32 // Small dimension for fast tests + let ids: string[] + + beforeEach(async () => { + storage = new MemoryStorage() + index = new HNSWIndex( + { + M: 8, + efConstruction: 100, + efSearch: 50, + ml: 8, + quantization: { enabled: true, rerankMultiplier: 3 } + }, + euclideanDistance, + { useParallelization: false, storage } + ) + ids = [] + }) + + it('should return correct results with quantization enabled', async () => { + // Insert entities with UUID IDs + const vectors: number[][] = [] + for (let i = 0; i < 50; i++) { + const id = uuidv4() + ids.push(id) + const v = randomVector(dim) + vectors.push(v) + await index.addItem({ id, vector: v }) + } + + // Search with a known vector + const queryIdx = 5 + const results = await index.search(vectors[queryIdx], 5) + + // The exact vector should be the closest (distance ~0) + expect(results.length).toBeGreaterThan(0) + expect(results[0][0]).toBe(ids[queryIdx]) + expect(results[0][1]).toBeCloseTo(0, 1) + }) + + it('should find the exact match as top result', async () => { + const targetId = uuidv4() + const target = randomVector(dim) + await index.addItem({ id: targetId, vector: target }) + + // Add noise vectors + for (let i = 0; i < 30; i++) { + await index.addItem({ id: uuidv4(), vector: randomVector(dim) }) + } + + const results = await index.search(target, 1) + expect(results.length).toBe(1) + expect(results[0][0]).toBe(targetId) + }) + + it('should respect k parameter', async () => { + for (let i = 0; i < 100; i++) { + await index.addItem({ id: uuidv4(), vector: randomVector(dim) }) + } + + const results5 = await index.search(randomVector(dim), 5) + const results10 = await index.search(randomVector(dim), 10) + + expect(results5.length).toBe(5) + expect(results10.length).toBe(10) + }) + }) + + // ================================================================= + // 5. TWO-PHASE RERANK + // ================================================================= + describe('two-phase rerank', () => { + it('should accept rerank options in search', async () => { + const storage = new MemoryStorage() + const index = new HNSWIndex( + { + M: 8, + efConstruction: 100, + efSearch: 50, + ml: 8, + quantization: { enabled: true, rerankMultiplier: 1 } // Disable default rerank + }, + euclideanDistance, + { useParallelization: false, storage } + ) + + const dim = 32 + for (let i = 0; i < 30; i++) { + await index.addItem({ id: uuidv4(), vector: randomVector(dim) }) + } + + // Search with explicit rerank + const results = await index.search( + randomVector(dim), + 5, + undefined, + { rerank: { multiplier: 3 } } + ) + + expect(results.length).toBe(5) + // Results should be sorted by distance + for (let i = 1; i < results.length; i++) { + expect(results[i][1]).toBeGreaterThanOrEqual(results[i - 1][1]) + } + }) + + it('reranked results should be sorted by exact distance', async () => { + const storage = new MemoryStorage() + const index = new HNSWIndex( + { + M: 8, + efConstruction: 100, + efSearch: 50, + ml: 8, + quantization: { enabled: true, rerankMultiplier: 3 } + }, + euclideanDistance, + { useParallelization: false, storage } + ) + + const dim = 32 + for (let i = 0; i < 50; i++) { + await index.addItem({ id: uuidv4(), vector: randomVector(dim) }) + } + + const query = randomVector(dim) + const results = await index.search(query, 10) + + // Results should be sorted by exact distance (after reranking) + for (let i = 1; i < results.length; i++) { + expect(results[i][1]).toBeGreaterThanOrEqual(results[i - 1][1]) + } + }) + }) + + // ================================================================= + // 6. CONFIG DEFAULTS + // ================================================================= + describe('configuration defaults', () => { + it('should disable quantization by default', async () => { + const storage = new MemoryStorage() + const defaultIndex = new HNSWIndex( + { M: 4, efConstruction: 50, efSearch: 20 }, + euclideanDistance, + { useParallelization: false, storage } + ) + + const dim = 16 + for (let i = 0; i < 10; i++) { + await defaultIndex.addItem({ id: uuidv4(), vector: randomVector(dim) }) + } + + // Search should work identically to pre-quantization behavior + const results = await defaultIndex.search(randomVector(dim), 5) + expect(results.length).toBe(5) + + // Distances should be exact euclidean (not SQ8 approximate) + for (const [, dist] of results) { + expect(typeof dist).toBe('number') + expect(dist).toBeGreaterThanOrEqual(0) + } + }) + + it('should use default rerankMultiplier of 3', async () => { + const storage = new MemoryStorage() + const index = new HNSWIndex( + { + M: 4, + efConstruction: 50, + efSearch: 20, + quantization: { enabled: true } + }, + euclideanDistance, + { useParallelization: false, storage } + ) + + // The index should be created without errors + // Default rerankMultiplier should be 3 (tested implicitly via search) + const dim = 16 + for (let i = 0; i < 10; i++) { + await index.addItem({ id: uuidv4(), vector: randomVector(dim) }) + } + + const results = await index.search(randomVector(dim), 3) + expect(results.length).toBe(3) + }) + + it('should default vectorStorage to memory mode', async () => { + const storage = new MemoryStorage() + const index = new HNSWIndex( + { M: 4, efConstruction: 50, efSearch: 20 }, + euclideanDistance, + { useParallelization: false, storage } + ) + + const id = uuidv4() + const v = randomVector(16) + await index.addItem({ id, vector: v }) + + // In memory mode, vector should be immediately searchable + const results = await index.search(v, 1) + expect(results.length).toBe(1) + expect(results[0][0]).toBe(id) + expect(results[0][1]).toBeCloseTo(0, 5) + }) + }) +})