From 2e3f840b549d431db2327cbc9f2f9eeb79400113 Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Tue, 7 Jul 2026 17:23:09 -0400 Subject: [PATCH] feat(extensions): add LatticeWasmEmbeddings provider Adds LatticeWasmEmbeddings, a local embeddings provider backed by the published @khive-ai/lattice-embed-wasm package (a pure-Rust BERT-family text embedder compiled to WebAssembly). It plugs into the existing EmbeddingProvider seam alongside OpenAI/Cohere/Anthropic/HuggingFace, giving a native-Rust local embedder option with no @xenova/transformers dependency. - getMaxBatchSize() honestly returns 1: the wasm package has no batch API, so this does not simulate a larger batch. - getDimension() returns 384 for both supported models (minilm, bge-small). - embedTexts() dynamic-imports the optional peer dependency (mirrors the existing HuggingFaceEmbeddings pattern) and treats a null result from the wasm layer as a real error for an explicitly-selected provider, never a silent skip. - Added as an optional peerDependency (peerDependenciesMeta.optional), matching the existing openai/cohere-ai convention. - Tests mirror the existing per-provider style: construction, dimension, batch size, and a live-embed assertion (exact 384 dim, L2 norm) gated behind package/model-weight availability. --- npm/packages/ruvector-extensions/package.json | 6 +- .../ruvector-extensions/src/embeddings.ts | 133 ++++++++++++++++++ npm/packages/ruvector-extensions/src/index.ts | 2 + .../tests/embeddings.test.ts | 70 +++++++++ 4 files changed, 210 insertions(+), 1 deletion(-) diff --git a/npm/packages/ruvector-extensions/package.json b/npm/packages/ruvector-extensions/package.json index a5854f242a..91d3361e92 100644 --- a/npm/packages/ruvector-extensions/package.json +++ b/npm/packages/ruvector-extensions/package.json @@ -36,7 +36,8 @@ }, "peerDependencies": { "openai": "^4.0.0", - "cohere-ai": "^7.0.0" + "cohere-ai": "^7.0.0", + "@khive-ai/lattice-embed-wasm": "^0.1.0" }, "peerDependenciesMeta": { "openai": { @@ -44,6 +45,9 @@ }, "cohere-ai": { "optional": true + }, + "@khive-ai/lattice-embed-wasm": { + "optional": true } }, "devDependencies": { diff --git a/npm/packages/ruvector-extensions/src/embeddings.ts b/npm/packages/ruvector-extensions/src/embeddings.ts index ac1aafbd01..29c1d7b49e 100644 --- a/npm/packages/ruvector-extensions/src/embeddings.ts +++ b/npm/packages/ruvector-extensions/src/embeddings.ts @@ -763,6 +763,138 @@ export class HuggingFaceEmbeddings extends EmbeddingProvider { } } +// ============================================================================ +// Lattice WASM Embeddings Provider +// ============================================================================ + +/** + * Configuration for lattice-embed-wasm local embeddings + */ +export interface LatticeWasmEmbeddingsConfig { + /** Model name (default: 'minilm'). See the supported-model list in the class doc. */ + model?: string; + /** Retry configuration */ + retryConfig?: Partial; +} + +/** + * Known lattice-embed-wasm models and their output dimensionality. Both are + * symmetric BERT-family encoders (no query/passage prefix at the wasm layer). + */ +const LATTICE_WASM_MODEL_DIMENSIONS: Record = { + minilm: 384, + 'bge-small': 384, +}; + +/** + * Local embeddings provider backed by `@khive-ai/lattice-embed-wasm`, a + * pure-Rust BERT-family text embedder compiled to WebAssembly. + * + * Node-only: the wasm adapter resolves model weights via `node:fs`, so this + * provider does not run in a browser. Single-text only: the wasm package + * exposes no batch API, so {@link getMaxBatchSize} honestly reports `1` + * rather than simulating a larger batch. Output embeddings are + * L2-normalized and symmetric (no query/passage asymmetry at the wasm + * layer, unlike some hosted embedding APIs). + * + * This is a convenience local-embedder option at parity with the existing + * `HuggingFaceEmbeddings` (transformers.js) local path above, not a faster + * or higher-quality alternative to it. + */ +export class LatticeWasmEmbeddings extends EmbeddingProvider { + private model: string; + private dimension: number; + + /** + * Creates a new lattice-embed-wasm local embeddings provider + * @param config - Configuration options + * @throws Error if the configured model is not one lattice-embed-wasm supports + */ + constructor(config: LatticeWasmEmbeddingsConfig = {}) { + super(config.retryConfig); + + const model = config.model || 'minilm'; + const dimension = LATTICE_WASM_MODEL_DIMENSIONS[model]; + + if (dimension === undefined) { + throw new Error( + `Unknown lattice-embed-wasm model "${model}". Supported models: ${Object.keys( + LATTICE_WASM_MODEL_DIMENSIONS + ).join(', ')}` + ); + } + + this.model = model; + this.dimension = dimension; + } + + getMaxBatchSize(): number { + // lattice-embed-wasm embeds one text per call; there is no batch API to + // wrap, so this honestly reports 1 rather than faking a larger batch. + return 1; + } + + getDimension(): number { + return this.dimension; + } + + async embedTexts(texts: string[]): Promise { + if (texts.length === 0) { + return { embeddings: [] }; + } + + let lattice: any; + try { + // Dynamic import to support optional peer dependency. The specifier is + // read from a variable (rather than passed as a literal) so + // TypeScript does not attempt to statically resolve module types for + // a peer that may not be installed. + const specifier = '@khive-ai/lattice-embed-wasm'; + lattice = await import(specifier); + } catch (error) { + throw new Error( + 'lattice-embed-wasm not found. Install it with: npm install @khive-ai/lattice-embed-wasm' + ); + } + + // No batch API: embed one text at a time. Node caches the dynamic + // import above by specifier, so repeated calls do not re-import. + const allResults: EmbeddingResult[] = []; + + for (let i = 0; i < texts.length; i++) { + const text = texts[i]; + + const vector = await this.withRetry(async () => { + const result = await lattice.embed(text, this.model); + + if (result === null) { + // For an explicitly-selected provider, a null result is a real + // failure (unknown model, unsupported over wasm, or weights that + // could not be resolved/verified) -- never a silent skip. + throw new Error( + `lattice-embed-wasm returned null embedding for model "${this.model}": the model is unknown, unsupported over wasm, or its weights could not be resolved` + ); + } + + return result as Float32Array; + }, `lattice-embed-wasm embedding for text ${i + 1}/${texts.length}`); + + allResults.push({ + embedding: Array.from(vector), + index: i, + }); + } + + return { + embeddings: allResults, + metadata: { + model: this.model, + provider: 'lattice-wasm', + }, + }; + } +} + // ============================================================================ // Helper Functions // ============================================================================ @@ -919,6 +1051,7 @@ export default { CohereEmbeddings, AnthropicEmbeddings, HuggingFaceEmbeddings, + LatticeWasmEmbeddings, // Helper functions embedAndInsert, diff --git a/npm/packages/ruvector-extensions/src/index.ts b/npm/packages/ruvector-extensions/src/index.ts index 5dbd9b0e76..180b02d11a 100644 --- a/npm/packages/ruvector-extensions/src/index.ts +++ b/npm/packages/ruvector-extensions/src/index.ts @@ -19,6 +19,7 @@ export { CohereEmbeddings, AnthropicEmbeddings, HuggingFaceEmbeddings, + LatticeWasmEmbeddings, // Helper functions embedAndInsert, @@ -34,6 +35,7 @@ export { type CohereEmbeddingsConfig, type AnthropicEmbeddingsConfig, type HuggingFaceEmbeddingsConfig, + type LatticeWasmEmbeddingsConfig, } from './embeddings.js'; // Re-export default for convenience diff --git a/npm/packages/ruvector-extensions/tests/embeddings.test.ts b/npm/packages/ruvector-extensions/tests/embeddings.test.ts index 09016cc37e..ac2b061137 100644 --- a/npm/packages/ruvector-extensions/tests/embeddings.test.ts +++ b/npm/packages/ruvector-extensions/tests/embeddings.test.ts @@ -13,6 +13,7 @@ import { CohereEmbeddings, AnthropicEmbeddings, HuggingFaceEmbeddings, + LatticeWasmEmbeddings, type BatchEmbeddingResult, type EmbeddingError, } from '../src/embeddings.js'; @@ -224,6 +225,75 @@ describe('HuggingFaceEmbeddings', () => { }); }); +// ============================================================================ +// Tests for Lattice WASM Provider +// ============================================================================ + +describe('LatticeWasmEmbeddings', () => { + it('should throw for an unknown model', () => { + assert.throws( + () => { + new LatticeWasmEmbeddings({ model: 'not-a-real-model' }); + }, + /Unknown lattice-embed-wasm model/ + ); + }); + + it('should create with default config', () => { + const lattice = new LatticeWasmEmbeddings(); + assert.strictEqual(lattice.getDimension(), 384); + assert.strictEqual(lattice.getMaxBatchSize(), 1); + }); + + it('should create with bge-small config', () => { + const lattice = new LatticeWasmEmbeddings({ model: 'bge-small' }); + assert.strictEqual(lattice.getDimension(), 384); + }); + + it('should not throw on construction (no eager load)', () => { + const lattice = new LatticeWasmEmbeddings(); + assert.ok(lattice); + }); + + it('should produce a 384-dim, L2-normalized embedding when the wasm package and its model weights are available', async (t) => { + let lattice: any; + try { + // Read from a variable (not a literal) so TypeScript does not attempt + // to statically resolve module types for an optional peer that may + // not be installed -- same rationale as in LatticeWasmEmbeddings itself. + const specifier = '@khive-ai/lattice-embed-wasm'; + lattice = await import(specifier); + } catch { + t.skip('@khive-ai/lattice-embed-wasm is not installed (optional peer dependency)'); + return; + } + void lattice; + + const provider = new LatticeWasmEmbeddings(); + let result: BatchEmbeddingResult; + try { + result = await provider.embedTexts(['Hello, world!']); + } catch (error: any) { + // Model weights are resolved from a local cache or a pinned release + // asset; neither is guaranteed to be present in every environment + // (e.g. a fresh CI checkout with no local model cache). This is an + // environment-availability gate, not a code-correctness failure. + t.skip(`lattice-embed-wasm model weights unavailable: ${error.message}`); + return; + } + + assert.strictEqual(result.embeddings.length, 1); + const embedding = result.embeddings[0].embedding; + // Mutation-sensitive: an exact dimension check, not just "is an array". + assert.strictEqual(embedding.length, 384); + + let sumSquares = 0; + for (const value of embedding) sumSquares += value * value; + const norm = Math.sqrt(sumSquares); + assert.ok(Math.abs(norm - 1.0) < 1e-3, `expected L2 norm near 1.0, got ${norm}`); + }); +}); + // ============================================================================ // Tests for Retry Logic // ============================================================================