@claude-flow/embeddings/src/searchable-embedding-cache.ts

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

/**
 * SearchableEmbeddingCache — ADR-121 Phase 2.
 *
 * The legacy `PersistentEmbeddingCache` and the in-process `LRUCache`
 * only support exact-key lookup. That's fine for "I've seen this
 * exact text before" but useless for "I haven't computed this exact
 * text but I have something within ε" — which is the whole point of
 * a vector cache for embedding-heavy workloads.
 *
 * This module adds an HNSW-backed cache implementation that:
 *   - Implements the same key-value contract as the legacy cache
 *     (back-compat: existing callers don't change)
 *   - Adds `search(query, k)` and `searchByText(text, k, embed)` for
 *     ANN retrieval over the cached vectors
 *   - Lazy-loads `@ruvector/core` as an optional peer dep; falls back
 *     to in-memory exact-only when the dep isn't installed
 *
 * Phase 2 acceptance:
 *   - ≥0.95 near-duplicate recall (where exact key match would be 0)
 *   - ≥10k inserts/sec on the standard test machine
 *   - Legacy PersistentEmbeddingCache unchanged
 */

import type { EmbeddingResult } from './types.js';
import { int8Encode, int8Decode, type Int8EncodedVector } from './quantization.js';

/**
 * Search result row.
 *
 * `score` is a **similarity** score: higher = more similar, identical
 * vectors → 1.0. The cache normalizes the raw distance metric returned
 * by `@ruvector/core` (which returns cosine *distance* and Euclidean
 * *distance*) into the similarity convention used elsewhere in this
 * package (e.g. `cosineSimilarity()`).
 *
 * Conversion table:
 *   cosine     : similarity = 1 - distance
 *   dot-product: similarity = -distance  (already in dot-product space)
 *   euclidean  : similarity = 1 / (1 + distance)  (monotonic, bounded [0,1])
 *   manhattan  : similarity = 1 / (1 + distance)
 */
export interface CacheSearchHit {
  /** Stable opaque id assigned to the cached entry. */
  readonly id: string;
  /** The cached embedding vector. */
  readonly vector: Float32Array;
  /** The original text the vector embeds, when available. */
  readonly text?: string;
  /** Similarity score in [0, 1] (higher = more similar). */
  readonly score: number;
}

export interface SearchableCacheOptions {
  /** Embedding dimension. Required — must match the producer service. */
  readonly dimension: number;
  /** Distance metric (default 'cosine'). */
  readonly metric?: 'cosine' | 'euclidean' | 'dot-product' | 'manhattan';
  /** Max in-memory entries (default 10,000). */
  readonly maxSize?: number;
  /**
   * Quantization mode for cached vectors (ADR-121 Phase 3, alpha.21):
   *   'none' (default) — full float32 in memory
   *   'int8'           — per-vector int8 + scale, 4× memory reduction,
   *                      recall ≥0.98 on unit-normalized vectors
   * RaBitQ (1-bit, 32× reduction) is a separate batch API; see
   * `RabitqSnapshot` for that path — it doesn't fit the streaming
   * insert/get/delete contract this cache exposes.
   */
  readonly quantize?: 'none' | 'int8';
}

/**
 * Minimal structural type for the @ruvector/core VectorDb we use.
 *
 * The real `@ruvector/core` API (probed against v0.1.31) is all-async
 * with NAPI binding bridges:
 *   - `insert({ id, vector })`  — object arg, returns Promise<void>
 *   - `search({ vector, k })`   — object arg, returns Promise<Array<{ id, score, vector }>>
 *   - `get(id)`                 — Promise<Float32Array | null>
 *   - `delete(id)`              — Promise<boolean>
 *   - `len()`                   — Promise<number>
 *   - `isEmpty()`               — Promise<boolean>
 *
 * `score` from search is a **distance** (cosine distance for the
 * 'Cosine' metric — identical vectors → ~0, orthogonal → 1). We
 * convert to similarity at the cache layer (see `distanceToSimilarity`).
 */
interface VectorDbSearchResult {
  readonly id: string;
  readonly score: number; // distance, not similarity
  readonly vector?: Float32Array | Record<number, number>;
}

interface VectorDbLike {
  insert(args: { id: string; vector: Float32Array | number[] }): Promise<void> | void;
  delete(id: string): Promise<boolean> | boolean;
  get(id: string): Promise<Float32Array | number[] | null | undefined> | Float32Array | number[] | null | undefined;
  search(args: { vector: Float32Array | number[]; k: number }): Promise<VectorDbSearchResult[]> | VectorDbSearchResult[];
  len(): Promise<number> | number;
  isEmpty(): Promise<boolean> | boolean;
}

interface RuvectorCoreModule {
  VectorDb: {
    withDimensions(dim: number, metric?: string): VectorDbLike;
  };
  JsDistanceMetric: Record<string, string>;
}

async function loadRuvectorCore(): Promise<RuvectorCoreModule | null> {
  try {
    // `@ruvector/core` is a CJS module — ESM dynamic import returns
    // an object whose `.default` is the module surface. We accept
    // either flat-named exports (vitest mocks) or the .default shape
    // (real NAPI module).
    const mod = (await import('@ruvector/core')) as unknown as Record<string, unknown>;
    const surface = (typeof (mod as { VectorDb?: unknown }).VectorDb === 'function' || typeof (mod as { VectorDb?: unknown }).VectorDb === 'object'
      ? mod
      : (mod.default as Record<string, unknown> | undefined)) as RuvectorCoreModule | undefined;
    return surface ?? null;
  } catch {
    return null;
  }
}

/** Convert @ruvector/core's distance score into a [0,1] similarity. */
function distanceToSimilarity(distance: number, metric: string): number {
  switch (metric) {
    case 'cosine':
      // Cosine distance ∈ [0, 2]; cosine similarity = 1 - distance,
      // but @ruvector/core uses the (1 - cos) convention so the
      // simple inversion gives back the conventional [-1, 1] range
      // (clamped to [0, 1] for non-pathological inputs).
      return Math.max(0, Math.min(1, 1 - distance));
    case 'dot-product':
      // Dot-product "distance" is actually negated similarity.
      // Clamp to [0, 1] for unit vectors; out-of-range inputs flow through.
      return Math.max(0, Math.min(1, -distance));
    case 'euclidean':
    case 'manhattan':
      // Unbounded distances → monotonic, bounded [0, 1] similarity.
      return 1 / (1 + Math.max(0, distance));
    default:
      return distance;
  }
}

/**
 * In-memory fallback when @ruvector/core isn't installed. Provides
 * exact key lookup + a linear-scan `search` so the contract still
 * holds end-to-end (just at O(N) instead of O(log N)). The fallback
 * is intentionally not silent — callers see `degraded: true` in the
 * cache's stats so they can opt into installing the optional dep.
 */
class LinearScanFallback implements VectorDbLike {
  private readonly store = new Map<string, Float32Array>();

  insert(args: { id: string; vector: Float32Array | number[] }): void {
    const v = args.vector instanceof Float32Array ? args.vector : new Float32Array(args.vector);
    this.store.set(args.id, v);
  }

  delete(id: string): boolean {
    return this.store.delete(id);
  }

  get(id: string): Float32Array | null {
    return this.store.get(id) ?? null;
  }

  /**
   * Returns cosine **distance** (1 - similarity) to match the score
   * convention of `@ruvector/core`. The cache layer converts to
   * similarity uniformly across backings.
   */
  search(args: { vector: Float32Array | number[]; k: number }): VectorDbSearchResult[] {
    const q = args.vector instanceof Float32Array ? args.vector : new Float32Array(args.vector);
    const out: VectorDbSearchResult[] = [];
    for (const [id, vec] of this.store) {
      const sim = cosineSimilarity(q, vec);
      // Match real backing's score convention: distance, not similarity.
      const distance = Math.max(0, 1 - sim);
      out.push({ id, score: distance, vector: vec });
    }
    out.sort((a, b) => a.score - b.score); // ascending: smaller distance first
    return out.slice(0, args.k);
  }

  len(): number { return this.store.size; }
  isEmpty(): boolean { return this.store.size === 0; }
}

function cosineSimilarity(a: Float32Array, b: Float32Array): number {
  let dot = 0, na = 0, nb = 0;
  const len = Math.min(a.length, b.length);
  for (let i = 0; i < len; i++) {
    dot += a[i]! * b[i]!;
    na += a[i]! * a[i]!;
    nb += b[i]! * b[i]!;
  }
  if (na === 0 || nb === 0) return 0;
  return dot / (Math.sqrt(na) * Math.sqrt(nb));
}

export class SearchableEmbeddingCache {
  private backing: VectorDbLike | null = null;
  private degraded = true; // until backing is loaded
  private readonly options: Required<SearchableCacheOptions>;
  /** Map from cache key (text or content hash) to the assigned vector id. */
  private readonly keyToId = new Map<string, string>();
  /** Reverse map from vector id back to the original text key. */
  private readonly idToKey = new Map<string, string>();
  /**
   * Per-id int8 store. Only populated when `quantize === 'int8'`. The
   * cache continues to hand back float32 vectors to callers (dequantized
   * on demand) — quantization is a memory-layout concern transparent
   * to the IEmbeddingService contract.
   */
  private readonly int8Store = new Map<string, Int8EncodedVector>();
  private nextId = 1;
  private hits = 0;
  private misses = 0;

  constructor(options: SearchableCacheOptions) {
    this.options = {
      dimension: options.dimension,
      metric: options.metric ?? 'cosine',
      maxSize: options.maxSize ?? 10_000,
      quantize: options.quantize ?? 'none',
    };
  }

  /** Lazy-initialize the HNSW backing. Idempotent. */
  async initialize(): Promise<void> {
    if (this.backing) return;
    const mod = await loadRuvectorCore();
    if (mod) {
      try {
        // Map our string metric onto the JsDistanceMetric enum.
        const metricMap: Record<string, string> = {
          'cosine': mod.JsDistanceMetric?.Cosine ?? 'Cosine',
          'euclidean': mod.JsDistanceMetric?.Euclidean ?? 'Euclidean',
          'dot-product': mod.JsDistanceMetric?.DotProduct ?? 'DotProduct',
          'manhattan': mod.JsDistanceMetric?.Manhattan ?? 'Manhattan',
        };
        this.backing = mod.VectorDb.withDimensions(this.options.dimension, metricMap[this.options.metric]);
        this.degraded = false;
        return;
      } catch {
        // Fall through to the linear-scan fallback on construction failure.
      }
    }
    this.backing = new LinearScanFallback();
    this.degraded = true;
  }

  /** Cache the embedding under the given text key. */
  async set(key: string, vector: Float32Array, opts?: { id?: string }): Promise<void> {
    await this.initialize();
    if (vector.length !== this.options.dimension) {
      throw new Error(
        `SearchableEmbeddingCache: vector dimension ${vector.length} ` +
        `doesn't match cache dimension ${this.options.dimension}`,
      );
    }
    // Reuse an existing id when the key was seen before, so updates
    // don't bloat the index.
    let id = this.keyToId.get(key);
    if (id) {
      await this.backing!.delete(id);
    } else {
      id = opts?.id ?? `e${this.nextId++}`;
      this.keyToId.set(key, id);
      this.idToKey.set(id, key);
    }
    // ADR-121 Phase 3 — int8 quantization mode. The ANN backing still
    // receives the float32 vector (its index lookup needs precision),
    // but we also stash the int8 encoding so get(key) can return the
    // smaller representation when memory pressure is the binding
    // constraint. The float32 round-trip stays the default.
    if (this.options.quantize === 'int8') {
      this.int8Store.set(id, int8Encode(vector));
    }
    await this.backing!.insert({ id, vector });
    await this.maybeEvict();
  }

  /** Exact-key lookup. Returns the cached vector or undefined. */
  async get(key: string): Promise<Float32Array | undefined> {
    await this.initialize();
    const id = this.keyToId.get(key);
    if (!id) { this.misses++; return undefined; }
    // int8 path: dequantize from the int8 store. This is faster than
    // hitting the ANN backing and is the whole point of quantize mode
    // (smaller memory footprint, deterministic round-trip).
    if (this.options.quantize === 'int8') {
      const enc = this.int8Store.get(id);
      if (!enc) { this.misses++; return undefined; }
      this.hits++;
      return int8Decode(enc);
    }
    const v = await this.backing!.get(id);
    if (!v) { this.misses++; return undefined; }
    this.hits++;
    return this.toFloat32Array(v);
  }

  /** Whether the exact key is cached. */
  async has(key: string): Promise<boolean> {
    await this.initialize();
    return this.keyToId.has(key);
  }

  /** Remove a cached entry by key. */
  async delete(key: string): Promise<boolean> {
    await this.initialize();
    const id = this.keyToId.get(key);
    if (!id) return false;
    await this.backing!.delete(id);
    this.keyToId.delete(key);
    this.idToKey.delete(id);
    this.int8Store.delete(id);
    return true;
  }

  /**
   * ANN search over the cached embeddings. Returns top-k neighbors
   * by similarity, sorted descending by score (higher = more similar).
   */
  async search(query: Float32Array, k: number = 10): Promise<CacheSearchHit[]> {
    await this.initialize();
    if (query.length !== this.options.dimension) {
      throw new Error(
        `SearchableEmbeddingCache.search: query dimension ${query.length} ` +
        `doesn't match cache dimension ${this.options.dimension}`,
      );
    }
    if (await this.backing!.isEmpty()) return [];
    const raw = await this.backing!.search({ vector: query, k });
    const out: CacheSearchHit[] = [];
    for (const row of raw) {
      // Search results from @ruvector/core include the vector inline;
      // the linear-scan fallback also includes it. Fall back to a
      // separate get() call only when the backing didn't return it.
      let vec: Float32Array | null = null;
      if (row.vector) {
        vec = this.toFloat32Array(row.vector);
      } else {
        const fetched = await this.backing!.get(row.id);
        if (fetched) vec = this.toFloat32Array(fetched);
      }
      if (!vec) continue;
      out.push({
        id: row.id,
        score: distanceToSimilarity(row.score, this.options.metric),
        vector: vec,
        text: this.idToKey.get(row.id),
      });
    }
    // Already sorted by distance ascending; after similarity conversion
    // that's similarity descending — but be defensive in case the
    // backing returned unsorted results.
    out.sort((a, b) => b.score - a.score);
    return out;
  }

  /**
   * Normalize whatever the backing returns (Float32Array, number[],
   * or an object-as-array from a NAPI bridge) into a Float32Array.
   */
  private toFloat32Array(v: Float32Array | number[] | Record<number, number>): Float32Array {
    if (v instanceof Float32Array) return v;
    if (Array.isArray(v)) return new Float32Array(v);
    // NAPI bridges sometimes return `{ 0: x, 1: y, ... length: n }` —
    // copy element-wise into a sized buffer.
    const dim = this.options.dimension;
    const out = new Float32Array(dim);
    for (let i = 0; i < dim; i++) {
      const slot = (v as Record<number, number>)[i];
      if (typeof slot === 'number') out[i] = slot;
    }
    return out;
  }

  /**
   * Embed a text query (via the supplied embedder), then ANN search
   * for nearest cached neighbors. Convenience helper so callers
   * don't have to wire embed → search manually.
   */
  async searchByText(
    text: string,
    k: number,
    embed: (text: string) => Promise<EmbeddingResult>,
  ): Promise<CacheSearchHit[]> {
    const { embedding } = await embed(text);
    // EmbeddingResult.embedding is Float32Array | number[]; normalize
    // to Float32Array so the search call has a concrete typed view.
    const query = embedding instanceof Float32Array ? embedding : new Float32Array(embedding);
    return this.search(query, k);
  }

  /**
   * Snapshot of cache stats — size, hit rate, backing status.
   *
   * Size is read from the local id-map (synchronous) instead of the
   * backing's `len()` because the real `@ruvector/core` returns a
   * Promise for that. The id-map stays in lockstep with the backing
   * through `set`/`delete`/`clear`/`maybeEvict`, so this is exact.
   */
  stats(): { size: number; hits: number; misses: number; hitRate: number; degraded: boolean; metric: string; dimension: number; quantize: 'none' | 'int8'; estimatedBytes: number } {
    const total = this.hits + this.misses;
    // Per-vector memory: float32 = 4 bytes/dim, int8 = 1 byte/dim + 4 scale.
    const bytesPerVec = this.options.quantize === 'int8'
      ? (this.options.dimension + 4)
      : (this.options.dimension * 4);
    return {
      size: this.idToKey.size,
      hits: this.hits,
      misses: this.misses,
      hitRate: total > 0 ? this.hits / total : 0,
      degraded: this.degraded,
      metric: this.options.metric,
      dimension: this.options.dimension,
      quantize: this.options.quantize,
      estimatedBytes: this.idToKey.size * bytesPerVec,
    };
  }

  /** Clear all cached entries. */
  async clear(): Promise<void> {
    if (!this.backing) {
      this.keyToId.clear();
      this.idToKey.clear();
      this.int8Store.clear();
      this.hits = 0;
      this.misses = 0;
      return;
    }
    for (const id of this.idToKey.keys()) {
      await this.backing.delete(id);
    }
    this.keyToId.clear();
    this.idToKey.clear();
    this.int8Store.clear();
    this.hits = 0;
    this.misses = 0;
  }

  /**
   * FIFO eviction when over maxSize. Async because the backing's
   * `delete` may be async (the real NAPI binding returns a Promise).
   */
  private async maybeEvict(): Promise<void> {
    if (!this.backing) return;
    while (this.idToKey.size > this.options.maxSize) {
      let oldestId: string | null = null;
      let oldestNum = Number.POSITIVE_INFINITY;
      for (const id of this.idToKey.keys()) {
        const n = parseInt(id.replace(/^e/, ''), 10);
        if (!Number.isNaN(n) && n < oldestNum) { oldestNum = n; oldestId = id; }
      }
      if (!oldestId) break;
      const key = this.idToKey.get(oldestId)!;
      await this.backing.delete(oldestId);
      this.keyToId.delete(key);
      this.idToKey.delete(oldestId);
      this.int8Store.delete(oldestId);
    }
  }
}

/**
 * Factory — keeps the constructor a one-liner and matches the
 * createXxxService convention used elsewhere in the package.
 */
export function createSearchableEmbeddingCache(
  options: SearchableCacheOptions,
): SearchableEmbeddingCache {
  return new SearchableEmbeddingCache(options);
}