@claude-flow/embeddings/src/ann-router.ts

V3 Embedding Service - OpenAI, Transformers.js, Agentic-Flow (ONNX), Mock providers with hyperbolic embeddings, normalization, and chunking

/**
 * ADR-121 Phase 8 — `AnnRouter` composition layer.
 *
 * The package now ships three ANN backings, each with a different
 * sweet spot:
 *
 *   `SearchableEmbeddingCache` (HNSW + `@ruvector/core`)
 *     - in-memory, streaming insert + delete, ANN search
 *     - sweet spot: <100k vectors, frequently mutated, low-latency
 *
 *   `RabitqSnapshot` (1-bit quantization + `@ruvector/rabitq-wasm`)
 *     - in-memory, batch-build only, 32× memory reduction with
 *       exact-L2² rerank
 *     - sweet spot: 100k-1M vectors, mostly read, memory-tight
 *
 *   `DiskannSnapshot` (DiskANN/Vamana + `@ruvector/diskann`)
 *     - on-disk via `storagePath`, batch-build (with streaming insert
 *       before build), survives process restarts
 *     - sweet spot: >1M vectors, persistent, agent-fleet-scale
 *
 * Choosing between them requires understanding the workload —
 * something a downstream caller often doesn't want to think about.
 * `AnnRouter` takes a workload descriptor (`{ corpusSize, persistent,
 * mutable }`) and routes to the right backing automatically, behind a
 * single unified `search(query, k)` interface.
 *
 * Routing rules (documented + enumerable, see `decideBacking()`):
 *
 *   persistent=true              → DiskANN (only on-disk option)
 *   corpusSize ≥ 1_000_000       → DiskANN (memory pressure)
 *   corpusSize ≥ 100_000 + !mutable → RaBitQ (batch, compressed)
 *   default (small, mutable, in-RAM)→ HNSW cache
 *
 * When the preferred backing's peer dep isn't installed, the router
 * falls back to the next-best available one and reports `degraded:
 * true` + a `degradationReason`. Operators always get *something*
 * that works (even if it's the linear-scan fallback inside the HNSW
 * cache) rather than a hard failure.
 */

import { createSearchableEmbeddingCache, type SearchableEmbeddingCache } from './searchable-embedding-cache.js';
import { RabitqSnapshot } from './quantization.js';
import { DiskannSnapshot, diskannAvailable } from './diskann-snapshot.js';

export type AnnBacking = 'hnsw' | 'rabitq' | 'diskann';

export interface AnnRouterWorkload {
  /** Approximate corpus size (vectors). Used for the routing decision. */
  readonly corpusSize: number;
  /** Embedding dimension. Required. */
  readonly dimension: number;
  /**
   * Will the corpus need to survive process restarts? When true,
   * forces DiskANN regardless of size.
   */
  readonly persistent?: boolean;
  /**
   * Will the corpus be mutated after the initial batch (streaming
   * inserts, deletes)? When true, prefers HNSW over RaBitQ/DiskANN
   * (those are batch-build).
   */
  readonly mutable?: boolean;
  /**
   * On-disk storage path. Required when `persistent: true`. The
   * snapshot lives at this path and can be reopened via
   * `AnnRouter.load(workload)`.
   */
  readonly storagePath?: string;
}

export interface AnnRouterHit {
  readonly id: string;
  /** Similarity-style score in [0, 1] for HNSW, raw distance for the snapshot backings. Lower-better unless via HNSW. */
  readonly score: number;
  /** Original vector when the backing returns it (HNSW does, snapshot backings don't always). */
  readonly vector?: Float32Array;
}

export interface AnnRouterDecision {
  /** Backing selected by `decideBacking`. */
  readonly backing: AnnBacking;
  /** Free-form rationale — useful for logs + the routing smoke. */
  readonly reason: string;
  /** Was a fallback used because the preferred backing wasn't available? */
  readonly degraded: boolean;
  /** When `degraded`, why the preferred backing was unavailable. */
  readonly degradationReason?: string;
}

/**
 * Pure function exposing the routing decision. Doesn't actually
 * instantiate anything — useful for previewing what `AnnRouter.create`
 * would do, and for tests that assert routing without booting NAPI.
 */
export function decideBacking(workload: AnnRouterWorkload): { backing: AnnBacking; reason: string } {
  if (workload.persistent) {
    return { backing: 'diskann', reason: 'persistent=true forces an on-disk backing (DiskANN)' };
  }
  if (workload.corpusSize >= 1_000_000) {
    return { backing: 'diskann', reason: `corpusSize=${workload.corpusSize} ≥1M; DiskANN preserves recall under memory pressure` };
  }
  if (workload.corpusSize >= 100_000 && !workload.mutable) {
    return { backing: 'rabitq', reason: `corpusSize=${workload.corpusSize} in 100k-1M batch range; RaBitQ gives 32× compression` };
  }
  return { backing: 'hnsw', reason: `corpusSize=${workload.corpusSize} small or mutable=true; HNSW handles streaming + low-latency` };
}

/**
 * Unified ANN-router handle. Always exposes the same `add`/`build`/
 * `search` surface regardless of which backing is live. Callers can
 * inspect `decision` to see what was picked + why.
 */
export class AnnRouter {
  readonly decision: AnnRouterDecision;
  private readonly workload: AnnRouterWorkload;
  private hnsw: SearchableEmbeddingCache | null = null;
  private rabitq: RabitqSnapshot | null = null;
  private diskann: DiskannSnapshot | null = null;
  // Buffered entries before `build()` for batch backings.
  private pending: Array<{ id: string; vector: Float32Array }> = [];

  private constructor(workload: AnnRouterWorkload, decision: AnnRouterDecision) {
    this.workload = workload;
    this.decision = decision;
  }

  /**
   * Pick a backing, instantiate it, return the handle. The router
   * never throws on missing peer deps — it falls back to a less-
   * preferred backing and reports `degraded: true`.
   */
  static async create(workload: AnnRouterWorkload): Promise<AnnRouter> {
    if (workload.persistent && !workload.storagePath) {
      throw new Error('AnnRouter.create: persistent=true requires storagePath');
    }
    const initial = decideBacking(workload);
    let backing: AnnBacking = initial.backing;
    let reason: string = initial.reason;
    let degraded = false;
    let degradationReason: string | undefined;

    // Try to instantiate the preferred backing; degrade if peer dep
    // isn't reachable. The HNSW cache has its own internal linear-
    // scan fallback so it always succeeds — that's the floor.
    const router = new AnnRouter(workload, { backing, reason, degraded, degradationReason });

    if (backing === 'diskann') {
      const ok = await diskannAvailable();
      if (ok) {
        try {
          router.diskann = await DiskannSnapshot.create({
            dimension: workload.dimension,
            storagePath: workload.storagePath,
          });
        } catch (err) {
          degraded = true;
          degradationReason = `DiskANN create failed: ${err instanceof Error ? err.message : String(err)}`;
          backing = 'hnsw';
          reason = 'DiskANN unavailable; falling back to HNSW';
        }
      } else {
        degraded = true;
        degradationReason = '@ruvector/diskann not installed';
        backing = 'hnsw';
        reason = 'DiskANN unavailable; falling back to HNSW';
      }
    }

    if (backing === 'rabitq') {
      // RabitqSnapshot is batch-build — we defer instantiation until
      // build() is called (it needs the corpus to construct).
      // No peer-dep probe here; the build() call will surface the
      // named error if the wasm module isn't installable.
    }

    if (backing === 'hnsw') {
      router.hnsw = createSearchableEmbeddingCache({
        dimension: workload.dimension,
        // The HNSW backing caps at maxSize; pass corpusSize as a
        // generous upper bound.
        maxSize: Math.max(1000, workload.corpusSize),
      });
    }

    // Re-write the decision now that we know the final routing.
    (router as { decision: AnnRouterDecision }).decision = {
      backing,
      reason,
      degraded,
      degradationReason,
    };
    return router;
  }

  /**
   * Insert vectors. For HNSW this is immediate. For RaBitQ this
   * buffers until `build()`. For DiskANN this calls `add()` on the
   * underlying snapshot (also requires `build()` before search).
   */
  async add(entries: ReadonlyArray<{ id: string; vector: Float32Array }>): Promise<void> {
    if (this.hnsw) {
      for (const e of entries) {
        await this.hnsw.set(e.id, e.vector);
      }
      return;
    }
    if (this.diskann) {
      this.diskann.add(entries);
      return;
    }
    // RaBitQ path — buffer for build().
    for (const e of entries) this.pending.push(e);
  }

  /**
   * Build the index. No-op for HNSW (streaming); required for
   * RaBitQ + DiskANN before the first `search`.
   */
  async build(): Promise<void> {
    if (this.hnsw) return; // streaming, no build pass
    if (this.diskann) {
      await this.diskann.build();
      return;
    }
    // RaBitQ — construct from buffered entries.
    if (this.pending.length === 0) {
      throw new Error('AnnRouter.build: no entries buffered for the RaBitQ snapshot');
    }
    this.rabitq = await RabitqSnapshot.build(this.pending, { dimension: this.workload.dimension });
    this.pending = []; // free the buffer
  }

  /**
   * Search the active backing. Always returns `AnnRouterHit[]` sorted
   * such that the more-relevant result comes first — HNSW returns
   * similarity-descending; RaBitQ + DiskANN return distance-ascending
   * (the lower the score, the better the match). Consumers reading
   * `score` should interpret it relative to `decision.backing`.
   */
  async search(query: Float32Array, k: number): Promise<AnnRouterHit[]> {
    if (this.hnsw) {
      const raw = await this.hnsw.search(query, k);
      // The HNSW cache returns its own internal id in `h.id` (auto-
      // generated as e1/e2/... by SearchableEmbeddingCache) and the
      // original caller-supplied key in `h.text`. The AnnRouter
      // contract says callers see what they put in, so map .text → .id.
      return raw.map(h => ({ id: h.text ?? h.id, score: h.score, vector: h.vector }));
    }
    if (this.rabitq) {
      const raw = await this.rabitq.search(query, k);
      return raw.map(h => ({ id: h.id, score: h.distance, vector: h.vector }));
    }
    if (this.diskann) {
      const raw = await this.diskann.search(query, k);
      return raw.map(h => ({ id: h.id, score: h.distance }));
    }
    throw new Error('AnnRouter.search: backing not initialized — call build() first?');
  }

  /**
   * Current vector count. Differs slightly by backing (HNSW returns
   * cache size; RaBitQ returns batch size; DiskANN includes tombstones).
   */
  count(): number {
    if (this.hnsw) return this.hnsw.stats().size;
    if (this.rabitq) return this.rabitq.size;
    if (this.diskann) return this.diskann.count();
    return this.pending.length;
  }

  /**
   * Persist the snapshot. Only meaningful for DiskANN (the other
   * backings are in-memory). No-op + returns false otherwise.
   */
  save(): boolean {
    if (this.diskann && this.workload.storagePath) {
      this.diskann.save(this.workload.storagePath);
      return true;
    }
    return false;
  }
}