@claude-flow/embeddings/src/ruvector-onnx-embedding-service.ts

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

/**
 * RuvectorOnnxEmbeddingService — ADR-121 Phase 1.
 *
 * Drop-in replacement for the legacy `TransformersEmbeddingService` that
 * uses `ruvector-onnx-embeddings-wasm` instead of the retired
 * `@xenova/transformers@2.x`. Matches the existing IEmbeddingService
 * contract so consumers don't migrate: opt in via
 * `provider: 'ruvector-onnx'` on `createEmbeddingService`.
 *
 * Why this exists:
 *  - `@xenova/transformers` was renamed upstream to
 *    `@huggingface/transformers@4.x` and the v2 line is unmaintained.
 *    Carrying it in the runtime tree of every embedding install is a
 *    growing liability.
 *  - `ruvector-onnx-embeddings-wasm` is pure WASM with SIMD + parallel
 *    workers, smaller install footprint (no native bindings, no
 *    Python-style model-conversion tooling shipped), and runs
 *    identically in Node, browsers, Workers, and Deno.
 *  - Same model catalog (all-MiniLM-L6-v2, BGE, E5, GTE) so output
 *    parity with the Xenova path is achievable to >0.999 cosine
 *    correlation on a fixed test corpus — that's the Phase 1
 *    acceptance gate.
 *
 * Optional peer dep: when `ruvector-onnx-embeddings-wasm` isn't
 * installed, construction throws a clear error naming the dep and the
 * single install command. Callers that don't opt into `ruvector-onnx`
 * are unaffected.
 */

import type {
  EmbeddingProvider,
  RuvectorOnnxEmbeddingConfig,
  EmbeddingResult,
  BatchEmbeddingResult,
} from './types.js';
import { BaseEmbeddingService } from './embedding-service-base.js';

// Minimal structural type for the embedder we get from the WASM module.
// Mirrors the public surface of WasmEmbedder we use — we don't reach for
// anything else so a future minor of the WASM package can extend the
// class without breaking us.
interface WasmEmbedderLike {
  embedBatch(texts: string[]): Float32Array;
  dimension(): number;
  maxLength(): number;
  similarity(text1: string, text2: string): number;
  free(): void;
}

// Structural type for the loader.js entry — the package exposes a
// `createEmbedder(modelName)` helper that handles model download +
// WASM init + WasmEmbedder construction.
interface RuvectorLoaderModule {
  readonly DEFAULT_MODEL: string;
  readonly MODELS: Record<string, { dimension: number; maxLength: number }>;
  createEmbedder(modelName?: string): Promise<WasmEmbedderLike>;
}

/**
 * Lazy loader for `ruvector-onnx-embeddings-wasm/loader`. Direct
 * dynamic `import()` so vitest / mocks can intercept; wrapped in
 * try/catch so the module-not-found case is a clean named error.
 */
async function loadRuvectorOnnxLoader(): Promise<RuvectorLoaderModule> {
  try {
    return (await import('ruvector-onnx-embeddings-wasm/loader')) as unknown as RuvectorLoaderModule;
  } catch (err) {
    const detail = err instanceof Error ? err.message : String(err);
    throw new Error(
      'ruvector-onnx-embeddings-wasm is not installed. Run: ' +
      'npm install ruvector-onnx-embeddings-wasm (or add it to your ' +
      'project\'s optional peer deps). Original error: ' + detail,
    );
  }
}

export class RuvectorOnnxEmbeddingService extends BaseEmbeddingService {
  readonly provider: EmbeddingProvider = 'ruvector-onnx';
  private embedder: WasmEmbedderLike | null = null;
  private readonly modelName: string;
  private readonly normalizeOutput: boolean;
  private dimensionCache: number | null = null;

  constructor(config: RuvectorOnnxEmbeddingConfig) {
    super(config);
    this.modelName = config.model ?? 'all-MiniLM-L6-v2';
    this.normalizeOutput = config.normalize ?? true;
  }

  private async initialize(): Promise<void> {
    if (this.embedder) return;
    const loader = await loadRuvectorOnnxLoader();
    try {
      // createEmbedder handles model download (~23MB for MiniLM-L6),
      // WASM init, tokenizer load, and WasmEmbedder construction.
      this.embedder = await loader.createEmbedder(this.modelName);
      this.dimensionCache = this.embedder.dimension();
    } catch (err) {
      const detail = err instanceof Error ? err.message : String(err);
      // Distinguish between missing-package and broken-package paths
      // so operators don't waste time `npm install`ing something
      // that's already there. The most common runtime breakage is the
      // wasm-bindgen ESM/CJS mismatch — surface the dep name + the
      // upstream issue path so the operator can find it.
      throw new Error(
        `ruvector-onnx-embeddings-wasm runtime initialization failed for ` +
        `model '${this.modelName}'. The package is installed but createEmbedder ` +
        `threw. Most common cause: WASM/ESM packaging mismatch in the published ` +
        `tarball (the .js wrapper references CommonJS \`module\` from an ESM ` +
        `context). Check ruvector-onnx-embeddings-wasm@latest for fixes, or ` +
        `fall back to provider: 'transformers' until upstream is patched. ` +
        `Original error: ${detail}`,
      );
    }
  }

  async embed(text: string): Promise<EmbeddingResult> {
    await this.initialize();

    // Cache check first — same shape as TransformersEmbeddingService so
    // the cache hit-rate metrics stay comparable between providers.
    const cached = this.cache.get(text);
    if (cached) {
      this.emitEvent({ type: 'cache_hit', text });
      return { embedding: cached, latencyMs: 0, cached: true };
    }

    this.emitEvent({ type: 'embed_start', text });
    const startTime = performance.now();

    try {
      // The WASM API only exposes embedBatch — single-text embedding is
      // a one-element batch. The returned Float32Array is flat
      // (length === N_texts * dimension); we slice the first row out.
      const flat = this.embedder!.embedBatch([text]);
      const dim = this.dimensionCache!;
      const embedding = this.normalizeOutput
        ? this.l2NormalizeView(flat.subarray(0, dim))
        : new Float32Array(flat.subarray(0, dim));

      this.cache.set(text, embedding);
      const latencyMs = performance.now() - startTime;
      this.emitEvent({ type: 'embed_complete', text, latencyMs });
      return { embedding, latencyMs };
    } catch (error) {
      const message = error instanceof Error ? error.message : 'Unknown error';
      this.emitEvent({ type: 'embed_error', text, error: message });
      throw new Error(`ruvector-onnx embedding failed: ${message}`);
    }
  }

  async embedBatch(texts: string[]): Promise<BatchEmbeddingResult> {
    await this.initialize();

    this.emitEvent({ type: 'batch_start', count: texts.length });
    const startTime = performance.now();

    // Split into cached vs uncached. Uncached texts go through the
    // WASM batch call once, which is the whole reason for using this
    // backend — the WasmEmbedder is dramatically faster batched than
    // one-at-a-time.
    const embeddings: Float32Array[] = new Array(texts.length);
    const uncachedIndices: number[] = [];
    const uncachedTexts: string[] = [];
    let cacheHits = 0;

    for (let i = 0; i < texts.length; i++) {
      const cached = this.cache.get(texts[i]!);
      if (cached) {
        embeddings[i] = cached;
        cacheHits++;
        this.emitEvent({ type: 'cache_hit', text: texts[i]! });
      } else {
        uncachedIndices.push(i);
        uncachedTexts.push(texts[i]!);
      }
    }

    if (uncachedTexts.length > 0) {
      const flat = this.embedder!.embedBatch(uncachedTexts);
      const dim = this.dimensionCache!;
      for (let j = 0; j < uncachedTexts.length; j++) {
        const row = flat.subarray(j * dim, (j + 1) * dim);
        const embedding = this.normalizeOutput
          ? this.l2NormalizeView(row)
          : new Float32Array(row);
        const text = uncachedTexts[j]!;
        const idx = uncachedIndices[j]!;
        this.cache.set(text, embedding);
        embeddings[idx] = embedding;
      }
    }

    const totalLatencyMs = performance.now() - startTime;
    this.emitEvent({ type: 'batch_complete', count: texts.length, latencyMs: totalLatencyMs });

    return {
      embeddings,
      totalLatencyMs,
      avgLatencyMs: totalLatencyMs / texts.length,
      cacheStats: { hits: cacheHits, misses: texts.length - cacheHits },
    };
  }

  /**
   * L2-normalize a Float32Array view into a freshly-allocated
   * Float32Array. The WASM embedder can normalize internally via
   * WasmEmbedderConfig, but we do it here so the cached vector is
   * always normalized regardless of model defaults — keeps similarity
   * scoring predictable across providers.
   */
  private l2NormalizeView(view: Float32Array): Float32Array {
    let sum = 0;
    for (let i = 0; i < view.length; i++) sum += view[i]! * view[i]!;
    const norm = Math.sqrt(sum);
    if (norm === 0) return new Float32Array(view);
    const out = new Float32Array(view.length);
    for (let i = 0; i < view.length; i++) out[i] = view[i]! / norm;
    return out;
  }

  /**
   * Release the WASM embedder. Idempotent; safe to call multiple times.
   */
  async close(): Promise<void> {
    if (this.embedder) {
      try { this.embedder.free(); } catch { /* best-effort */ }
      this.embedder = null;
      this.dimensionCache = null;
    }
  }
}

/**
 * Factory helper — mirrors the createEmbeddingService convention so
 * callers can opt in without instantiating directly.
 */
export function createRuvectorOnnxEmbeddingService(
  config: Omit<RuvectorOnnxEmbeddingConfig, 'provider'> = {},
): RuvectorOnnxEmbeddingService {
  return new RuvectorOnnxEmbeddingService({ ...config, provider: 'ruvector-onnx' });
}