From e79b2919b6b1d209efb46e9b71fc56a7653224f8 Mon Sep 17 00:00:00 2001 From: "Robert Kent Jr." Date: Tue, 11 Aug 2026 16:53:57 -0400 Subject: [PATCH 1/3] feat(search): make the embedding model and device configurable `getDb` called `transformersJs()` with no arguments, pinning every index and query to retriv's smallest default (bge-small-en-v1.5, 384d) on whatever device transformers.js chose, which is the CPU under Node. Neither was reachable through config, a flag, or an env var. That default thins out as skills accumulate: search builds one sqlite-vec DB per package and pools scores across all of them at query time, so cross-corpus ranking depends directly on embedding quality. Adds `embedModel` and `embedDevice` config keys, matching entries in `skilld config`, and `SKILLD_EMBED_MODEL` / `SKILLD_EMBED_DEVICE` overrides for single runs. Precedence is env, then config, then default. Both defaults are unchanged: `bge-small-en-v1.5`, and a device of `auto` that resolves to undefined so the option is omitted entirely. Device measurements on an Apple M5 Max, 120 documents, best of 3 (docs/sec): model cpu coreml webgpu bge-small-en-v1.5 664 198 1713 bge-base-en-v1.5 198 68 580 Xenova/bge-large-en-v1.5 71 9 201 webgpu is 2.6-2.9x faster than cpu at every size, 4.4x end to end through the index pipeline; coreml is consistently slower. The ranking is hardware-specific, so the device is offered rather than defaulted, and the picker leads with that caveat. Two correctness details: The bge-large entry pins the full repo id `Xenova/bge-large-en-v1.5`. retriv's bare `bge-large-en-v1.5` preset maps to `onnx-community/bge-large-en-v1.5`, whose weights return 401, so selecting it would fail at first index. The embedding cache keyed vectors by text hash and validated only dimensions. That was safe while the model was fixed; selecting one makes it reachable, since bge-large-en-v1.5 and bge-m3 are both 1024d. Switching kept every cached vector and served one model's embeddings against another's queries. No crash, just silently wrong ranking, with the correct answer dropping out of the top 3 on a 43-document corpus. Cache identity is now `@`, cleared on change, because the same model on a different backend can differ numerically. --- README.md | 51 ++++++++ src/commands/config.ts | 90 ++++++++++++++ src/core/config.ts | 12 ++ src/retriv/embedding-cache.ts | 23 +++- src/retriv/index.ts | 16 ++- src/retriv/models.ts | 122 +++++++++++++++++++ test/unit/embed-models.test.ts | 134 +++++++++++++++++++++ test/unit/embedding-cache-identity.test.ts | 103 ++++++++++++++++ 8 files changed, 547 insertions(+), 4 deletions(-) create mode 100644 src/retriv/models.ts create mode 100644 test/unit/embed-models.test.ts create mode 100644 test/unit/embedding-cache-identity.test.ts diff --git a/README.md b/README.md index f93ebc15..b69911c5 100644 --- a/README.md +++ b/README.md @@ -239,6 +239,57 @@ Generation runs locally: free, offline, no API key. Unlike the CLI and API backe The large default context can exceed memory for big models on constrained hardware (Ollama returns a 500). Lower `OLLAMA_NUM_CTX` or pick a smaller model if generation fails to load. +### Embedding Model + +`skilld search` is powered by a local embedding model. It runs offline through transformers.js — no API key, and no network traffic after the first download. Pick one under **Embedding model** in `skilld config`: + +| Model | Dimensions | Notes | +|-------|-----------:|-------| +| `bge-small-en-v1.5` | 384 | Default. Fastest to index, smallest download. | +| `bge-base-en-v1.5` | 768 | Balanced accuracy and speed. | +| `bge-m3` | 1024 | Multilingual, 8192-token context. | + +Larger models retrieve more accurately but cost more time and memory when indexing. Set `SKILLD_EMBED_MODEL` to override the saved setting for a single run: + +```bash +SKILLD_EMBED_MODEL=bge-m3 skilld add npm:vue +``` + +Search indexes store fixed-width vectors, so changing to a model with different dimensions strands existing indexes. Rebuild them after switching: + +```bash +skilld update --force +``` + +### Embedding Device + +The embedding model runs on the CPU by default. **Embedding device** in `skilld config` moves it onto a GPU backend, which can be substantially faster: + +| Device | Notes | +|--------|-------| +| `auto` | Default. Lets transformers.js choose — CPU under Node. | +| `cpu` | Always available, predictable. | +| `webgpu` | Fastest on Apple Silicon in testing. | +| `coreml` | Apple Neural Engine. Measured slower than CPU for these models. | + +Measured on an Apple M5 Max, 120 documents, best of 3 after warm-up (docs/sec): + +| Model | `cpu` | `coreml` | `webgpu` | +|-------|------:|---------:|---------:| +| `bge-small-en-v1.5` | 664 | 198 | **1713** | +| `bge-base-en-v1.5` | 198 | 68 | **580** | +| `Xenova/bge-large-en-v1.5` | 71 | 9 | **201** | + +WebGPU was 2.6-2.9x faster than CPU at every size, which means `bge-large` on WebGPU indexes faster than `bge-base` does on CPU — better retrieval for less wall-clock. CoreML was consistently slower. + +The ranking is hardware-specific, so benchmark before trusting a device on other machines. Override for a single run with `SKILLD_EMBED_DEVICE`: + +```bash +SKILLD_EMBED_DEVICE=cpu skilld update --force +``` + +If a backend is unavailable, indexing fails to start — switch back to `auto`. + ### Eject Export a skill as a portable, self-contained directory for sharing via git repos: diff --git a/src/commands/config.ts b/src/commands/config.ts index 20d76282..bb7b8e43 100644 --- a/src/commands/config.ts +++ b/src/commands/config.ts @@ -11,6 +11,7 @@ import { guard, menuLoop } from '../cli/menu.ts' import { NO_MODELS_MESSAGE, OAUTH_NOTE, pickModel } from '../cli/model-picker.ts' import { defaultFeatures, readConfig, updateConfig } from '../core/config.ts' import { getProjectState } from '../core/skills.ts' +import { DEFAULT_EMBED_DEVICE, DEFAULT_EMBED_MODEL, EMBED_DEVICES, EMBED_MODELS, getEmbedModelInfo, resolveEmbedModel } from '../retriv/models.ts' export async function configCommand(): Promise { const initConfig = readConfig() @@ -46,8 +47,15 @@ export async function configCommand(): Promise { const oauthHint = connectedOAuth > 0 ? `${connectedOAuth} connected` : 'none' options.push({ label: 'OAuth providers', value: 'oauth', hint: `${oauthHint} · ⚠ may violate provider ToS` }) } + const embedModel = resolveEmbedModel(config.embedModel) + const embedHint = features.search + ? `${embedModel} · local model powering skilld search` + : `${embedModel} · search is disabled in Data sources` + const embedDevice = config.embedDevice || DEFAULT_EMBED_DEVICE options.push( { label: 'Enhancement model', value: 'model', hint: `${modelHint} · rewrites SKILL.md with best practices` }, + { label: 'Embedding model', value: 'embedModel', hint: embedHint }, + { label: 'Embedding device', value: 'embedDevice', hint: `${embedDevice} · where the embedding model runs` }, { label: 'Target agent', value: 'agent', hint: `${config.agent || 'auto-detect'} · where skills are installed` }, ) return options @@ -92,6 +100,16 @@ export async function configCommand(): Promise { break } + case 'embedModel': { + await configureEmbedModel() + break + } + + case 'embedDevice': { + await configureEmbedDevice() + break + } + case 'agent': { const config = readConfig() const agentChoice = guard(await p.select({ @@ -230,6 +248,78 @@ async function configureModel(): Promise { } } +// ── Embedding model selection ──────────────────────────────────────── + +async function configureEmbedModel(): Promise { + const config = readConfig() + const current = resolveEmbedModel(config.embedModel) + const envOverride = process.env.SKILLD_EMBED_MODEL?.trim() + + if (envOverride) { + p.log.warn(`SKILLD_EMBED_MODEL is set to ${envOverride} and overrides this setting for the current shell.`) + } + + const choice = guard(await p.select({ + message: 'Embedding model — indexes and queries docs for skilld search', + options: EMBED_MODELS.map(m => ({ + label: m.label, + value: m.id, + hint: `${m.dimensions}d · ${m.hint}`, + })), + initialValue: current, + })) + + if (choice === config.embedModel || (choice === DEFAULT_EMBED_MODEL && !config.embedModel)) { + p.log.info(`Embedding model unchanged (${choice})`) + return + } + + const previous = getEmbedModelInfo(current) + const next = getEmbedModelInfo(choice as string) + updateConfig({ embedModel: choice === DEFAULT_EMBED_MODEL ? undefined : choice as string }) + p.log.success(`Embedding model set to ${choice}`) + + // sqlite-vec columns are fixed-width, so a dimension change strands existing + // indexes: they stay queryable at the old width but new docs cannot join them. + if (previous && next && previous.dimensions !== next.dimensions) { + p.log.warn( + `Vector width changed ${previous.dimensions}d → ${next.dimensions}d. ` + + 'Existing search indexes must be rebuilt: skilld update --force', + ) + } +} + +// ── Embedding device selection ─────────────────────────────────────── + +async function configureEmbedDevice(): Promise { + const config = readConfig() + const current = config.embedDevice || DEFAULT_EMBED_DEVICE + const envOverride = process.env.SKILLD_EMBED_DEVICE?.trim() + + if (envOverride) + p.log.warn(`SKILLD_EMBED_DEVICE is set to ${envOverride} and overrides this setting for the current shell.`) + + p.note( + 'The fastest backend depends on your hardware. On an Apple M5 Max, WebGPU\n' + + 'ran 2.6-2.9x faster than CPU across every model size, while CoreML ran\n' + + '3-8x slower. Benchmark before trusting a device on other machines.', + 'Choosing a device', + ) + + const choice = guard(await p.select({ + message: 'Embedding device — where the model runs', + options: EMBED_DEVICES.map(d => ({ label: d.label, value: d.id, hint: d.hint })), + initialValue: current, + })) + + updateConfig({ embedDevice: choice === DEFAULT_EMBED_DEVICE ? undefined : choice as string }) + p.log.success(`Embedding device set to ${choice}`) + + if (choice !== DEFAULT_EMBED_DEVICE && choice !== 'cpu') { + p.log.info('If indexing fails to start, the backend is unavailable on this machine — switch back to Auto.') + } +} + export const configCommandDef = defineCommand({ meta: { name: 'config', description: 'Edit settings' }, args: {}, diff --git a/src/core/config.ts b/src/core/config.ts index 98093c11..d59c861c 100644 --- a/src/core/config.ts +++ b/src/core/config.ts @@ -32,6 +32,10 @@ export function getActiveFeatures(overrides?: Partial): Features export interface SkilldConfig { model?: OptimizeModel agent?: string + /** Local embedding model used to build and query the search index */ + embedModel?: string + /** Execution device for the embedding model (auto, cpu, webgpu, coreml) */ + embedDevice?: string features?: FeaturesConfig projects?: string[] skipLlm?: boolean @@ -102,6 +106,10 @@ export function readConfig(): SkilldConfig { config.model = value as OptimizeModel if (key === 'agent' && value) config.agent = value + if (key === 'embedModel' && value) + config.embedModel = value + if (key === 'embedDevice' && value) + config.embedDevice = value if (key === 'skipLlm') config.skipLlm = value === 'true' } @@ -122,6 +130,10 @@ export function writeConfig(config: SkilldConfig): void { yaml += `model: ${config.model}\n` if (config.agent) yaml += `agent: ${config.agent}\n` + if (config.embedModel) + yaml += `embedModel: ${config.embedModel}\n` + if (config.embedDevice) + yaml += `embedDevice: ${config.embedDevice}\n` if (config.skipLlm) yaml += `skipLlm: true\n` if (config.features) { diff --git a/src/retriv/embedding-cache.ts b/src/retriv/embedding-cache.ts index 041ac1fd..126d3dde 100644 --- a/src/retriv/embedding-cache.ts +++ b/src/retriv/embedding-cache.ts @@ -50,7 +50,16 @@ function createSqliteStorage(db: DatabaseSync) { } } -export async function cachedEmbeddings(config: EmbeddingConfig): Promise { +/** + * Wrap an embedding provider with the on-disk vector cache. + * + * `model` identifies which embedder produced the cached vectors. Entries are + * keyed by text hash alone, so vectors from a different model would be served + * for the same text — two models of equal width (bge-large and + * qwen3-embedding:0.6b are both 1024d) would silently mix embedding spaces and + * destroy ranking. Dimensions alone cannot catch that; the model id can. + */ +export async function cachedEmbeddings(config: EmbeddingConfig, model?: string): Promise { const { cachedEmbeddings: retrivCached } = await import('retriv/embeddings/cached') const db = await openDb() const storage = createSqliteStorage(db) @@ -63,10 +72,18 @@ export async function cachedEmbeddings(config: EmbeddingConfig): Promise) { throw new SearchDepsUnavailableError(err) throw err } - const embeddings = await cachedEmbeddings(transformersJs()) + const userConfig = readConfig() + const embedModel = resolveEmbedModel(userConfig.embedModel) + const device = resolveEmbedDevice(userConfig.embedDevice) + // Cache identity pairs model with device: cached vectors are only valid for + // the embedder that produced them, and backends can differ numerically. + const embeddings = await cachedEmbeddings( + transformersJs({ + model: embedModel, + // Omitted when `auto` so transformers.js keeps its own device resolution. + ...(device ? { device } : {}), + }), + `${embedModel}@${device ?? 'auto'}`, + ) return createRetriv({ driver: sqliteMod.default({ path: config.dbPath, diff --git a/src/retriv/models.ts b/src/retriv/models.ts new file mode 100644 index 00000000..ee2d7eda --- /dev/null +++ b/src/retriv/models.ts @@ -0,0 +1,122 @@ +/** + * Local embedding models available to the search index. + * + * Every model here runs offline through transformers.js — no API key, no + * network after the initial download. Larger models retrieve more accurately + * but cost more time and memory to index with. + * + * Dimensions are fixed per model and sqlite-vec columns are fixed-width, so + * switching model invalidates existing indexes. Rebuild with + * `skilld update --force`. + */ +export interface EmbedModelInfo { + /** Model id passed to retriv (resolved to a Hugging Face repo internally) */ + id: string + label: string + /** Vector width — determines index layout */ + dimensions: number + hint: string +} + +export const DEFAULT_EMBED_MODEL = 'bge-small-en-v1.5' + +export const EMBED_MODELS: readonly EmbedModelInfo[] = [ + { + id: 'bge-small-en-v1.5', + label: 'BGE small (English)', + dimensions: 384, + hint: 'fastest to index, smallest download', + }, + { + id: 'bge-base-en-v1.5', + label: 'BGE base (English)', + dimensions: 768, + hint: 'balanced accuracy and speed', + }, + { + // Pinned to the full repo id on purpose: retriv's `bge-large-en-v1.5` + // preset maps to `onnx-community/bge-large-en-v1.5`, which returns 401. + // The Xenova repo carries the same weights and resolves correctly. + id: 'Xenova/bge-large-en-v1.5', + label: 'BGE large (English)', + dimensions: 1024, + hint: 'most accurate English retrieval, slowest to index', + }, + { + id: 'bge-m3', + label: 'BGE m3 (multilingual)', + dimensions: 1024, + hint: 'multilingual, 8192-token context', + }, +] + +export function getEmbedModelInfo(id: string): EmbedModelInfo | undefined { + return EMBED_MODELS.find(m => m.id === id) +} + +/** + * Resolve the embedding model to index and query with. + * + * `SKILLD_EMBED_MODEL` wins so a single run can be overridden without touching + * saved config; otherwise the configured value, otherwise the default. + */ +export function resolveEmbedModel(configured?: string): string { + const fromEnv = process.env.SKILLD_EMBED_MODEL?.trim() + if (fromEnv) + return fromEnv + return configured || DEFAULT_EMBED_MODEL +} + +/** + * Execution device for the embedding model. + * + * `auto` means "let transformers.js decide", which resolves to CPU under Node. + * Everything else is opt-in because the fastest backend is hardware-specific: + * on an Apple M5 Max `webgpu` measured 2.6-2.9x faster than CPU across every + * bge size, while `coreml` measured 3-8x slower (it falls back to CPU for + * unsupported ops and pays for graph partitioning). + */ +export interface EmbedDeviceInfo { + id: string + label: string + hint: string +} + +export const DEFAULT_EMBED_DEVICE = 'auto' + +export const EMBED_DEVICES: readonly EmbedDeviceInfo[] = [ + { + id: 'auto', + label: 'Auto', + hint: 'let transformers.js choose — CPU under Node', + }, + { + id: 'cpu', + label: 'CPU', + hint: 'always available, predictable', + }, + { + id: 'webgpu', + label: 'GPU (WebGPU)', + hint: 'fastest on Apple Silicon in testing — verify on your hardware', + }, + { + id: 'coreml', + label: 'CoreML', + hint: 'Apple Neural Engine — measured slower than CPU for these models', + }, +] + +export function getEmbedDeviceInfo(id: string): EmbedDeviceInfo | undefined { + return EMBED_DEVICES.find(d => d.id === id) +} + +/** + * Resolve the execution device. Returns `undefined` for `auto` so the option + * is omitted entirely and transformers.js keeps its own default resolution. + */ +export function resolveEmbedDevice(configured?: string): string | undefined { + const fromEnv = process.env.SKILLD_EMBED_DEVICE?.trim() + const value = fromEnv || configured || DEFAULT_EMBED_DEVICE + return value === DEFAULT_EMBED_DEVICE ? undefined : value +} diff --git a/test/unit/embed-models.test.ts b/test/unit/embed-models.test.ts new file mode 100644 index 00000000..88a64786 --- /dev/null +++ b/test/unit/embed-models.test.ts @@ -0,0 +1,134 @@ +import { getModelDimensions, resolveModelForPreset } from 'retriv/embeddings/model-info' +import { afterEach, beforeEach, describe, expect, it } from 'vitest' +import { DEFAULT_EMBED_DEVICE, DEFAULT_EMBED_MODEL, EMBED_DEVICES, EMBED_MODELS, getEmbedDeviceInfo, getEmbedModelInfo, resolveEmbedDevice, resolveEmbedModel } from '../../src/retriv/models.ts' + +describe('resolveEmbedModel', () => { + let original: string | undefined + + beforeEach(() => { + original = process.env.SKILLD_EMBED_MODEL + delete process.env.SKILLD_EMBED_MODEL + }) + + afterEach(() => { + if (original === undefined) + delete process.env.SKILLD_EMBED_MODEL + else + process.env.SKILLD_EMBED_MODEL = original + }) + + it('falls back to the default when nothing is configured', () => { + expect(resolveEmbedModel(undefined)).toBe(DEFAULT_EMBED_MODEL) + }) + + it('uses the configured model', () => { + expect(resolveEmbedModel('bge-base-en-v1.5')).toBe('bge-base-en-v1.5') + }) + + it('lets the env var override configured and default', () => { + process.env.SKILLD_EMBED_MODEL = 'bge-m3' + expect(resolveEmbedModel('bge-base-en-v1.5')).toBe('bge-m3') + expect(resolveEmbedModel(undefined)).toBe('bge-m3') + }) + + it('ignores a blank env var', () => { + process.env.SKILLD_EMBED_MODEL = ' ' + expect(resolveEmbedModel('bge-base-en-v1.5')).toBe('bge-base-en-v1.5') + }) +}) + +describe('embed model registry', () => { + it('includes the default model', () => { + expect(EMBED_MODELS.map(m => m.id)).toContain(DEFAULT_EMBED_MODEL) + }) + + it('has no duplicate ids', () => { + const ids = EMBED_MODELS.map(m => m.id) + expect(new Set(ids).size).toBe(ids.length) + }) + + it('looks up known models and rejects unknown ones', () => { + expect(getEmbedModelInfo(DEFAULT_EMBED_MODEL)?.dimensions).toBe(384) + expect(getEmbedModelInfo('not-a-model')).toBeUndefined() + }) + + // retriv's bare `bge-large-en-v1.5` preset maps to + // `onnx-community/bge-large-en-v1.5`, whose weights return 401. We pin the + // Xenova repo instead, so the bare id must never creep back in. + it('avoids the bge-large preset that resolves to unavailable weights', () => { + const ids = EMBED_MODELS.map(m => m.id) + expect(ids).not.toContain('bge-large-en-v1.5') + expect(ids).toContain('Xenova/bge-large-en-v1.5') + }) + + // Guards against drift: every id must resolve to a real transformers.js repo + // and our declared width must match retriv's registry, since the declared + // width is what warns users about rebuilding indexes. + it('matches retriv model resolution and dimensions', () => { + for (const model of EMBED_MODELS) { + const resolved = resolveModelForPreset(model.id, 'transformers.js') + expect(resolved, `${model.id} should resolve`).toBeTruthy() + expect(resolved, `${model.id} should map to a namespaced repo`).toContain('/') + expect(getModelDimensions(model.id), `${model.id} dimensions`).toBe(model.dimensions) + } + }) +}) + +describe('resolveEmbedDevice', () => { + let original: string | undefined + + beforeEach(() => { + original = process.env.SKILLD_EMBED_DEVICE + delete process.env.SKILLD_EMBED_DEVICE + }) + + afterEach(() => { + if (original === undefined) + delete process.env.SKILLD_EMBED_DEVICE + else + process.env.SKILLD_EMBED_DEVICE = original + }) + + // `auto` must resolve to undefined so the option is omitted entirely and + // transformers.js keeps its own device resolution. + it('returns undefined for auto so the option is omitted', () => { + expect(resolveEmbedDevice(undefined)).toBeUndefined() + expect(resolveEmbedDevice(DEFAULT_EMBED_DEVICE)).toBeUndefined() + }) + + it('returns the configured device', () => { + expect(resolveEmbedDevice('webgpu')).toBe('webgpu') + }) + + it('lets the env var override configured and default', () => { + process.env.SKILLD_EMBED_DEVICE = 'cpu' + expect(resolveEmbedDevice('webgpu')).toBe('cpu') + expect(resolveEmbedDevice(undefined)).toBe('cpu') + }) + + it('ignores a blank env var', () => { + process.env.SKILLD_EMBED_DEVICE = ' ' + expect(resolveEmbedDevice('webgpu')).toBe('webgpu') + }) + + it('treats an env var of auto as unset', () => { + process.env.SKILLD_EMBED_DEVICE = 'auto' + expect(resolveEmbedDevice('webgpu')).toBeUndefined() + }) +}) + +describe('embed device registry', () => { + it('includes the default device', () => { + expect(EMBED_DEVICES.map(d => d.id)).toContain(DEFAULT_EMBED_DEVICE) + }) + + it('has no duplicate ids', () => { + const ids = EMBED_DEVICES.map(d => d.id) + expect(new Set(ids).size).toBe(ids.length) + }) + + it('looks up known devices and rejects unknown ones', () => { + expect(getEmbedDeviceInfo('webgpu')?.label).toBe('GPU (WebGPU)') + expect(getEmbedDeviceInfo('not-a-device')).toBeUndefined() + }) +}) diff --git a/test/unit/embedding-cache-identity.test.ts b/test/unit/embedding-cache-identity.test.ts new file mode 100644 index 00000000..00b6c569 --- /dev/null +++ b/test/unit/embedding-cache-identity.test.ts @@ -0,0 +1,103 @@ +import { DatabaseSync } from 'node:sqlite' +import { describe, expect, it } from 'vitest' + +/** + * Guards the cache-invalidation rule in `src/retriv/embedding-cache.ts`. + * + * Vectors are keyed by text hash alone, so the only thing preventing one + * model's vectors being served to another is the stored identity. Dimensions + * are not enough: `Xenova/bge-large-en-v1.5` and `ollama:qwen3-embedding:0.6b` + * are both 1024d, so switching between them would silently mix embedding + * spaces and wreck ranking. + * + * This reimplements the decision against an in-memory database so the rule is + * pinned without touching the user's real cache. + */ +function applyIdentity(db: DatabaseSync, dimensions: number, model?: string): void { + const get = db.prepare('SELECT value FROM meta WHERE key = ?') + const set = db.prepare('INSERT OR REPLACE INTO meta (key, value) VALUES (?, ?)') + + const storedDims = get.get('dimensions') as { value: string } | undefined + const storedModel = get.get('model') as { value: string } | undefined + const dimsChanged = storedDims && Number(storedDims.value) !== dimensions + const modelChanged = model !== undefined && storedModel?.value !== model + + if (dimsChanged || modelChanged) + db.exec('DELETE FROM embeddings') + + set.run('dimensions', String(dimensions)) + if (model !== undefined) + set.run('model', model) +} + +function makeDb(): DatabaseSync { + const db = new DatabaseSync(':memory:') + db.exec('CREATE TABLE embeddings (text_hash TEXT PRIMARY KEY, embedding BLOB NOT NULL)') + db.exec('CREATE TABLE meta (key TEXT PRIMARY KEY, value TEXT NOT NULL)') + return db +} + +function seed(db: DatabaseSync, n = 3): void { + const stmt = db.prepare('INSERT OR IGNORE INTO embeddings (text_hash, embedding) VALUES (?, ?)') + for (let i = 0; i < n; i++) + stmt.run(`hash-${i}`, Buffer.from(new Float32Array([i, i, i]).buffer)) +} + +function count(db: DatabaseSync): number { + return (db.prepare('SELECT COUNT(*) c FROM embeddings').get() as { c: number }).c +} + +describe('embedding cache identity', () => { + it('keeps cached vectors when model and dimensions are unchanged', () => { + const db = makeDb() + applyIdentity(db, 1024, 'model-a') + seed(db) + applyIdentity(db, 1024, 'model-a') + expect(count(db)).toBe(3) + db.close() + }) + + // The regression: equal width, different model. + it('clears cached vectors when the model changes at identical dimensions', () => { + const db = makeDb() + applyIdentity(db, 1024, 'Xenova/bge-large-en-v1.5@webgpu') + seed(db) + expect(count(db)).toBe(3) + + applyIdentity(db, 1024, 'ollama:qwen3-embedding:0.6b') + expect(count(db)).toBe(0) + db.close() + }) + + it('clears cached vectors when dimensions change', () => { + const db = makeDb() + applyIdentity(db, 384, 'model-a') + seed(db) + applyIdentity(db, 1024, 'model-a') + expect(count(db)).toBe(0) + db.close() + }) + + // Same model on a different backend: numeric output can differ, so vectors + // are only interchangeable within a device. + it('clears cached vectors when only the device changes', () => { + const db = makeDb() + applyIdentity(db, 1024, 'Xenova/bge-large-en-v1.5@cpu') + seed(db) + applyIdentity(db, 1024, 'Xenova/bge-large-en-v1.5@webgpu') + expect(count(db)).toBe(0) + db.close() + }) + + // A cache written before the model key existed has unknown provenance. + it('clears a legacy cache that has no stored model', () => { + const db = makeDb() + applyIdentity(db, 1024) + seed(db) + expect(count(db)).toBe(3) + + applyIdentity(db, 1024, 'model-a') + expect(count(db)).toBe(0) + db.close() + }) +}) From 473d13510f6407181c43226ba698dd09270b07d2 Mon Sep 17 00:00:00 2001 From: Harlan Wilton Date: Wed, 12 Aug 2026 15:28:58 +1000 Subject: [PATCH 2/3] fix(search): enforce embedding identity Forward the selected device to Transformers.js and record the model-device identity in each search index. Reject incompatible indexes before queries can mix embedding spaces. Replace duplicated cache tests with API-level regression coverage. --- README.md | 20 ++-- src/commands/config.ts | 37 ++++---- src/retriv/embedding-cache.ts | 2 +- src/retriv/index-embedding-identity.ts | 57 ++++++++++++ src/retriv/index.ts | 44 ++++++--- src/retriv/models.ts | 45 ++++++--- src/retriv/pool.ts | 7 +- src/retriv/transformers-embeddings.ts | 63 +++++++++++++ test/unit/embed-models.test.ts | 13 ++- test/unit/embedding-cache-identity.test.ts | 103 --------------------- test/unit/embedding-cache.test.ts | 50 ++++++++++ test/unit/index-embedding-identity.test.ts | 79 ++++++++++++++++ test/unit/transformers-embeddings.test.ts | 34 +++++++ 13 files changed, 390 insertions(+), 164 deletions(-) create mode 100644 src/retriv/index-embedding-identity.ts create mode 100644 src/retriv/transformers-embeddings.ts delete mode 100644 test/unit/embedding-cache-identity.test.ts create mode 100644 test/unit/index-embedding-identity.test.ts create mode 100644 test/unit/transformers-embeddings.test.ts diff --git a/README.md b/README.md index b69911c5..576ed5ae 100644 --- a/README.md +++ b/README.md @@ -241,7 +241,7 @@ The large default context can exceed memory for big models on constrained hardwa ### Embedding Model -`skilld search` is powered by a local embedding model. It runs offline through transformers.js — no API key, and no network traffic after the first download. Pick one under **Embedding model** in `skilld config`: +`skilld search` uses a local embedding model. It runs offline through transformers.js. It needs no API key or network after the first download. Pick one under **Embedding model** in `skilld config`: | Model | Dimensions | Notes | |-------|-----------:|-------| @@ -249,13 +249,14 @@ The large default context can exceed memory for big models on constrained hardwa | `bge-base-en-v1.5` | 768 | Balanced accuracy and speed. | | `bge-m3` | 1024 | Multilingual, 8192-token context. | -Larger models retrieve more accurately but cost more time and memory when indexing. Set `SKILLD_EMBED_MODEL` to override the saved setting for a single run: +Larger models retrieve more accurately but cost more time and memory when indexing. Set `SKILLD_EMBED_MODEL` to override the saved setting: ```bash -SKILLD_EMBED_MODEL=bge-m3 skilld add npm:vue +export SKILLD_EMBED_MODEL=bge-m3 +skilld update --force ``` -Search indexes store fixed-width vectors, so changing to a model with different dimensions strands existing indexes. Rebuild them after switching: +Each search index belongs to one model and device. Keep environment overrides set for both indexing and querying. Rebuild indexes after either setting changes: ```bash skilld update --force @@ -267,7 +268,7 @@ The embedding model runs on the CPU by default. **Embedding device** in `skilld | Device | Notes | |--------|-------| -| `auto` | Default. Lets transformers.js choose — CPU under Node. | +| `auto` | Default. Lets transformers.js choose, CPU under Node. | | `cpu` | Always available, predictable. | | `webgpu` | Fastest on Apple Silicon in testing. | | `coreml` | Apple Neural Engine. Measured slower than CPU for these models. | @@ -280,15 +281,16 @@ Measured on an Apple M5 Max, 120 documents, best of 3 after warm-up (docs/sec): | `bge-base-en-v1.5` | 198 | 68 | **580** | | `Xenova/bge-large-en-v1.5` | 71 | 9 | **201** | -WebGPU was 2.6-2.9x faster than CPU at every size, which means `bge-large` on WebGPU indexes faster than `bge-base` does on CPU — better retrieval for less wall-clock. CoreML was consistently slower. +WebGPU was 2.6 to 2.9 times faster than CPU at every size. `bge-large` on WebGPU indexed faster than `bge-base` on CPU. CoreML was consistently slower. -The ranking is hardware-specific, so benchmark before trusting a device on other machines. Override for a single run with `SKILLD_EMBED_DEVICE`: +The ranking is hardware-specific, so benchmark before trusting a device on other machines. Set `SKILLD_EMBED_DEVICE` to override the saved setting: ```bash -SKILLD_EMBED_DEVICE=cpu skilld update --force +export SKILLD_EMBED_DEVICE=cpu +skilld update --force ``` -If a backend is unavailable, indexing fails to start — switch back to `auto`. +If a backend is unavailable, indexing fails to start. Switch back to `auto`. ### Eject diff --git a/src/commands/config.ts b/src/commands/config.ts index bb7b8e43..3f1ac873 100644 --- a/src/commands/config.ts +++ b/src/commands/config.ts @@ -11,7 +11,7 @@ import { guard, menuLoop } from '../cli/menu.ts' import { NO_MODELS_MESSAGE, OAUTH_NOTE, pickModel } from '../cli/model-picker.ts' import { defaultFeatures, readConfig, updateConfig } from '../core/config.ts' import { getProjectState } from '../core/skills.ts' -import { DEFAULT_EMBED_DEVICE, DEFAULT_EMBED_MODEL, EMBED_DEVICES, EMBED_MODELS, getEmbedModelInfo, resolveEmbedModel } from '../retriv/models.ts' +import { DEFAULT_EMBED_DEVICE, DEFAULT_EMBED_MODEL, EMBED_DEVICES, EMBED_MODELS, resolveEmbedModel } from '../retriv/models.ts' export async function configCommand(): Promise { const initConfig = readConfig() @@ -248,7 +248,7 @@ async function configureModel(): Promise { } } -// ── Embedding model selection ──────────────────────────────────────── +// Embedding model selection async function configureEmbedModel(): Promise { const config = readConfig() @@ -260,7 +260,7 @@ async function configureEmbedModel(): Promise { } const choice = guard(await p.select({ - message: 'Embedding model — indexes and queries docs for skilld search', + message: 'Embedding model for indexing and querying skilld search', options: EMBED_MODELS.map(m => ({ label: m.label, value: m.id, @@ -274,22 +274,12 @@ async function configureEmbedModel(): Promise { return } - const previous = getEmbedModelInfo(current) - const next = getEmbedModelInfo(choice as string) updateConfig({ embedModel: choice === DEFAULT_EMBED_MODEL ? undefined : choice as string }) p.log.success(`Embedding model set to ${choice}`) - - // sqlite-vec columns are fixed-width, so a dimension change strands existing - // indexes: they stay queryable at the old width but new docs cannot join them. - if (previous && next && previous.dimensions !== next.dimensions) { - p.log.warn( - `Vector width changed ${previous.dimensions}d → ${next.dimensions}d. ` - + 'Existing search indexes must be rebuilt: skilld update --force', - ) - } + p.log.warn('Embedding model changed. Rebuild existing search indexes: skilld update --force') } -// ── Embedding device selection ─────────────────────────────────────── +// Embedding device selection async function configureEmbedDevice(): Promise { const config = readConfig() @@ -301,23 +291,28 @@ async function configureEmbedDevice(): Promise { p.note( 'The fastest backend depends on your hardware. On an Apple M5 Max, WebGPU\n' - + 'ran 2.6-2.9x faster than CPU across every model size, while CoreML ran\n' - + '3-8x slower. Benchmark before trusting a device on other machines.', + + 'ran 2.6 to 2.9 times faster than CPU. CoreML ran 3 to 8 times slower.\n' + + 'Benchmark before trusting a device on other machines.', 'Choosing a device', ) const choice = guard(await p.select({ - message: 'Embedding device — where the model runs', + message: 'Embedding device where the model runs', options: EMBED_DEVICES.map(d => ({ label: d.label, value: d.id, hint: d.hint })), initialValue: current, })) + if (choice === current) { + p.log.info(`Embedding device unchanged (${choice})`) + return + } + updateConfig({ embedDevice: choice === DEFAULT_EMBED_DEVICE ? undefined : choice as string }) p.log.success(`Embedding device set to ${choice}`) + p.log.warn('Embedding device changed. Rebuild existing search indexes: skilld update --force') - if (choice !== DEFAULT_EMBED_DEVICE && choice !== 'cpu') { - p.log.info('If indexing fails to start, the backend is unavailable on this machine — switch back to Auto.') - } + if (choice !== DEFAULT_EMBED_DEVICE && choice !== 'cpu') + p.log.info('If indexing fails to start, switch back to Auto. The backend may be unavailable on this machine.') } export const configCommandDef = defineCommand({ diff --git a/src/retriv/embedding-cache.ts b/src/retriv/embedding-cache.ts index 126d3dde..0b7efc88 100644 --- a/src/retriv/embedding-cache.ts +++ b/src/retriv/embedding-cache.ts @@ -55,7 +55,7 @@ function createSqliteStorage(db: DatabaseSync) { * * `model` identifies which embedder produced the cached vectors. Entries are * keyed by text hash alone, so vectors from a different model would be served - * for the same text — two models of equal width (bge-large and + * for the same text. Two models of equal width (bge-large and * qwen3-embedding:0.6b are both 1024d) would silently mix embedding spaces and * destroy ranking. Dimensions alone cannot catch that; the model id can. */ diff --git a/src/retriv/index-embedding-identity.ts b/src/retriv/index-embedding-identity.ts new file mode 100644 index 00000000..dbea2f43 --- /dev/null +++ b/src/retriv/index-embedding-identity.ts @@ -0,0 +1,57 @@ +import { DatabaseSync } from 'node:sqlite' +import { existsSync } from 'node:fs' +import { DEFAULT_EMBEDDING_IDENTITY } from './models.ts' + +const META_TABLE = 'skilld_meta' +const IDENTITY_KEY = 'embedding_identity' + +export type IndexEmbeddingIdentityState + = | { _tag: 'Current' } + | { _tag: 'Missing' } + | { _tag: 'Mismatch', current: string, stored: string } + +function tableExists(db: DatabaseSync, name: string): boolean { + return db.prepare("SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?").get(name) !== undefined +} + +function hasIndexedDocuments(db: DatabaseSync): boolean { + if (!tableExists(db, 'documents_meta')) + return false + const row = db.prepare('SELECT EXISTS(SELECT 1 FROM documents_meta) AS found').get() as { found: number } + return row.found === 1 +} + +export function checkIndexEmbeddingIdentity(dbPath: string, current: string): IndexEmbeddingIdentityState { + if (dbPath === ':memory:' || !existsSync(dbPath)) + return { _tag: 'Missing' } + + const db = new DatabaseSync(dbPath, { open: true, readOnly: true }) + try { + const row = tableExists(db, META_TABLE) + ? db.prepare(`SELECT value FROM ${META_TABLE} WHERE key = ?`).get(IDENTITY_KEY) as { value: string } | undefined + : undefined + const stored = row?.value ?? (hasIndexedDocuments(db) ? DEFAULT_EMBEDDING_IDENTITY : undefined) + + if (!stored) + return { _tag: 'Missing' } + if (stored !== current) + return { _tag: 'Mismatch', current, stored } + return { _tag: 'Current' } + } + finally { + db.close() + } +} + +export function recordIndexEmbeddingIdentity(dbPath: string, identity: string): void { + if (dbPath === ':memory:') + return + const db = new DatabaseSync(dbPath) + try { + db.exec(`CREATE TABLE IF NOT EXISTS ${META_TABLE} (key TEXT PRIMARY KEY, value TEXT NOT NULL)`) + db.prepare(`INSERT OR REPLACE INTO ${META_TABLE} (key, value) VALUES (?, ?)`).run(IDENTITY_KEY, identity) + } + finally { + db.close() + } +} diff --git a/src/retriv/index.ts b/src/retriv/index.ts index c33941e4..b5ce532a 100644 --- a/src/retriv/index.ts +++ b/src/retriv/index.ts @@ -1,7 +1,9 @@ import type { ChunkEntity, Document, IndexConfig, IndexPhase, IndexProgress, SearchFilter, SearchOptions, SearchResult, SearchSnippet } from './types.ts' import { readConfig } from '../core/config.ts' import { stripFrontmatter } from '../core/markdown.ts' -import { resolveEmbedDevice, resolveEmbedModel } from './models.ts' +import { checkIndexEmbeddingIdentity, recordIndexEmbeddingIdentity } from './index-embedding-identity.ts' +import { getEmbeddingIdentity, resolveEmbedDevice, resolveEmbedModel } from './models.ts' +import { transformersEmbeddings } from './transformers-embeddings.ts' export type { ChunkEntity, Document, IndexConfig, IndexPhase, IndexProgress, SearchFilter, SearchOptions, SearchResult, SearchSnippet } @@ -15,6 +17,14 @@ export class SearchDepsUnavailableError extends Error { } } +export class EmbeddingIndexMismatchError extends Error { + constructor(dbPath: string, stored: string, current: string) { + super(`Search index uses ${stored}, but embedding settings resolve to ${current}. Rebuild indexes with: skilld update --force`) + this.name = 'EmbeddingIndexMismatchError' + this.cause = { dbPath, stored, current } + } +} + let _fts5Available: boolean | null = null /** @@ -49,21 +59,27 @@ export async function getDb(config: Pick) { if (!checkFts5()) throw new SearchDepsUnavailableError(new Error('FTS5 module not available'), 'SQLite FTS5 module not available. Search indexing skipped. On Windows, run from WSL where FTS5 is included.') - let createRetriv, autoChunker, sqliteMod, sqliteVec, transformersJs, cachedEmbeddings + const userConfig = readConfig() + const embedModel = resolveEmbedModel(userConfig.embedModel) + const device = resolveEmbedDevice(userConfig.embedDevice) + const embeddingIdentity = getEmbeddingIdentity(embedModel, device) + const identityState = checkIndexEmbeddingIdentity(config.dbPath, embeddingIdentity) + if (identityState._tag === 'Mismatch') + throw new EmbeddingIndexMismatchError(config.dbPath, identityState.stored, identityState.current) + + let createRetriv, autoChunker, sqliteMod, sqliteVec, cachedEmbeddings try { ;([ { createRetriv }, { autoChunker }, sqliteMod, sqliteVec, - { transformersJs }, { cachedEmbeddings }, ] = await Promise.all([ import('retriv'), import('retriv/chunkers/auto'), import('retriv/db/sqlite'), import('sqlite-vec'), - import('retriv/embeddings/transformers-js'), import('./embedding-cache.ts'), ])) } @@ -72,20 +88,14 @@ export async function getDb(config: Pick) { throw new SearchDepsUnavailableError(err) throw err } - const userConfig = readConfig() - const embedModel = resolveEmbedModel(userConfig.embedModel) - const device = resolveEmbedDevice(userConfig.embedDevice) - // Cache identity pairs model with device: cached vectors are only valid for - // the embedder that produced them, and backends can differ numerically. const embeddings = await cachedEmbeddings( - transformersJs({ + transformersEmbeddings({ model: embedModel, - // Omitted when `auto` so transformers.js keeps its own device resolution. ...(device ? { device } : {}), }), - `${embedModel}@${device ?? 'auto'}`, + embeddingIdentity, ) - return createRetriv({ + const db = await createRetriv({ driver: sqliteMod.default({ path: config.dbPath, embeddings, @@ -93,6 +103,14 @@ export async function getDb(config: Pick) { }), chunking: autoChunker(), }) + try { + recordIndexEmbeddingIdentity(config.dbPath, embeddingIdentity) + } + catch (error) { + await db.close?.() + throw error + } + return db } /** diff --git a/src/retriv/models.ts b/src/retriv/models.ts index ee2d7eda..b667c4a5 100644 --- a/src/retriv/models.ts +++ b/src/retriv/models.ts @@ -1,19 +1,21 @@ +import { resolveModelForPreset } from 'retriv/embeddings/model-info' + /** * Local embedding models available to the search index. * - * Every model here runs offline through transformers.js — no API key, no + * Every model here runs offline through transformers.js. It needs no API key or * network after the initial download. Larger models retrieve more accurately * but cost more time and memory to index with. * * Dimensions are fixed per model and sqlite-vec columns are fixed-width, so - * switching model invalidates existing indexes. Rebuild with + * switching a model or device invalidates existing indexes. Rebuild with * `skilld update --force`. */ export interface EmbedModelInfo { /** Model id passed to retriv (resolved to a Hugging Face repo internally) */ id: string label: string - /** Vector width — determines index layout */ + /** Vector width determines index layout */ dimensions: number hint: string } @@ -57,8 +59,8 @@ export function getEmbedModelInfo(id: string): EmbedModelInfo | undefined { /** * Resolve the embedding model to index and query with. * - * `SKILLD_EMBED_MODEL` wins so a single run can be overridden without touching - * saved config; otherwise the configured value, otherwise the default. + * `SKILLD_EMBED_MODEL` overrides saved config. The configured value overrides + * the default. */ export function resolveEmbedModel(configured?: string): string { const fromEnv = process.env.SKILLD_EMBED_MODEL?.trim() @@ -72,9 +74,9 @@ export function resolveEmbedModel(configured?: string): string { * * `auto` means "let transformers.js decide", which resolves to CPU under Node. * Everything else is opt-in because the fastest backend is hardware-specific: - * on an Apple M5 Max `webgpu` measured 2.6-2.9x faster than CPU across every - * bge size, while `coreml` measured 3-8x slower (it falls back to CPU for - * unsupported ops and pays for graph partitioning). + * on an Apple M5 Max `webgpu` measured 2.6 to 2.9 times faster than CPU. + * `coreml` measured 3 to 8 times slower because it falls back to CPU for + * unsupported ops and pays for graph partitioning. */ export interface EmbedDeviceInfo { id: string @@ -84,11 +86,11 @@ export interface EmbedDeviceInfo { export const DEFAULT_EMBED_DEVICE = 'auto' -export const EMBED_DEVICES: readonly EmbedDeviceInfo[] = [ +export const EMBED_DEVICES = [ { id: 'auto', label: 'Auto', - hint: 'let transformers.js choose — CPU under Node', + hint: 'let transformers.js choose, CPU under Node', }, { id: 'cpu', @@ -98,14 +100,17 @@ export const EMBED_DEVICES: readonly EmbedDeviceInfo[] = [ { id: 'webgpu', label: 'GPU (WebGPU)', - hint: 'fastest on Apple Silicon in testing — verify on your hardware', + hint: 'fastest on Apple Silicon in testing, verify on your hardware', }, { id: 'coreml', label: 'CoreML', - hint: 'Apple Neural Engine — measured slower than CPU for these models', + hint: 'Apple Neural Engine, measured slower than CPU for these models', }, -] +] as const satisfies readonly EmbedDeviceInfo[] + +export type EmbedDevice = typeof EMBED_DEVICES[number]['id'] +export type RuntimeEmbedDevice = Exclude export function getEmbedDeviceInfo(id: string): EmbedDeviceInfo | undefined { return EMBED_DEVICES.find(d => d.id === id) @@ -115,8 +120,18 @@ export function getEmbedDeviceInfo(id: string): EmbedDeviceInfo | undefined { * Resolve the execution device. Returns `undefined` for `auto` so the option * is omitted entirely and transformers.js keeps its own default resolution. */ -export function resolveEmbedDevice(configured?: string): string | undefined { +export function resolveEmbedDevice(configured?: string): RuntimeEmbedDevice | undefined { const fromEnv = process.env.SKILLD_EMBED_DEVICE?.trim() const value = fromEnv || configured || DEFAULT_EMBED_DEVICE - return value === DEFAULT_EMBED_DEVICE ? undefined : value + if (value === DEFAULT_EMBED_DEVICE) + return undefined + if (!getEmbedDeviceInfo(value)) + throw new Error(`Unsupported embedding device: ${value}`) + return value as RuntimeEmbedDevice } + +export function getEmbeddingIdentity(model: string, device?: string): string { + return `${resolveModelForPreset(model, 'transformers.js')}@${device ?? DEFAULT_EMBED_DEVICE}` +} + +export const DEFAULT_EMBEDDING_IDENTITY = getEmbeddingIdentity(DEFAULT_EMBED_MODEL) diff --git a/src/retriv/pool.ts b/src/retriv/pool.ts index de432799..b54036ba 100644 --- a/src/retriv/pool.ts +++ b/src/retriv/pool.ts @@ -4,11 +4,16 @@ import { existsSync } from 'node:fs' import { fileURLToPath } from 'node:url' import { Worker } from 'node:worker_threads' import { dirname, join } from 'pathe' -import { SearchDepsUnavailableError } from './index.ts' +import { EmbeddingIndexMismatchError, SearchDepsUnavailableError } from './index.ts' function reconstructError(message: string, name?: string): Error { if (name === 'SearchDepsUnavailableError') return new SearchDepsUnavailableError(undefined, message) + if (name === 'EmbeddingIndexMismatchError') { + const error = new Error(message) + error.name = EmbeddingIndexMismatchError.name + return error + } return new Error(message) } diff --git a/src/retriv/transformers-embeddings.ts b/src/retriv/transformers-embeddings.ts new file mode 100644 index 00000000..33809b40 --- /dev/null +++ b/src/retriv/transformers-embeddings.ts @@ -0,0 +1,63 @@ +import type { EmbeddingConfig } from 'retriv' +import { rm } from 'node:fs/promises' +import { resolve } from 'pathe' +import { getModelDimensions, getModelMaxTokens, resolveModelForPreset } from 'retriv/embeddings/model-info' +import type { RuntimeEmbedDevice } from './models.ts' + +export interface TransformersEmbeddingOptions { + model: string + device?: RuntimeEmbedDevice +} + +/** Transformers.js provider with explicit device support. */ +export function transformersEmbeddings(options: TransformersEmbeddingOptions): EmbeddingConfig { + const model = resolveModelForPreset(options.model, 'transformers.js') + let cached: Awaited> | undefined + + return { + async resolve() { + if (cached) + return cached + + // Search is optional, so keep the model runtime outside the main CLI bundle. + const { env, pipeline } = await import('@huggingface/transformers') + const load = () => pipeline('feature-extraction', model, { + dtype: 'fp32', + ...(options.device ? { device: options.device } : {}), + }) + const extractor = await load().catch(async (error) => { + const corrupted = error instanceof Error + && (error.message.includes('Protobuf parsing failed') || String(error.cause).includes('Protobuf parsing failed')) + if (!corrupted || !env.cacheDir) + throw error + const cacheRoot = resolve(env.cacheDir) + const modelCache = resolve(cacheRoot, model) + if (!modelCache.startsWith(`${cacheRoot}/`)) + throw error + await rm(modelCache, { recursive: true, force: true }) + console.warn(`[skilld] Cleared corrupted model cache for ${model}, retrying...`) + return load() + }) + + const dimensions = getModelDimensions(model) + if (!dimensions) + throw new Error(`Unknown dimensions for model ${model}.`) + + const embedder = async (texts: string[]) => { + const output = await extractor(texts, { pooling: 'mean', normalize: true }) + const data = output.data as Float32Array + return Array.from( + { length: texts.length }, + (_, index) => data.slice(index * dimensions, (index + 1) * dimensions), + ) + } + + cached = { + embedder, + dimensions, + maxTokens: getModelMaxTokens(model), + } + return cached + }, + } +} diff --git a/test/unit/embed-models.test.ts b/test/unit/embed-models.test.ts index 88a64786..f3c698c3 100644 --- a/test/unit/embed-models.test.ts +++ b/test/unit/embed-models.test.ts @@ -1,6 +1,6 @@ import { getModelDimensions, resolveModelForPreset } from 'retriv/embeddings/model-info' import { afterEach, beforeEach, describe, expect, it } from 'vitest' -import { DEFAULT_EMBED_DEVICE, DEFAULT_EMBED_MODEL, EMBED_DEVICES, EMBED_MODELS, getEmbedDeviceInfo, getEmbedModelInfo, resolveEmbedDevice, resolveEmbedModel } from '../../src/retriv/models.ts' +import { DEFAULT_EMBED_DEVICE, DEFAULT_EMBED_MODEL, EMBED_DEVICES, EMBED_MODELS, getEmbedDeviceInfo, getEmbedModelInfo, getEmbeddingIdentity, resolveEmbedDevice, resolveEmbedModel } from '../../src/retriv/models.ts' describe('resolveEmbedModel', () => { let original: string | undefined @@ -115,6 +115,17 @@ describe('resolveEmbedDevice', () => { process.env.SKILLD_EMBED_DEVICE = 'auto' expect(resolveEmbedDevice('webgpu')).toBeUndefined() }) + + it('rejects an unsupported device', () => { + expect(() => resolveEmbedDevice('invalid')).toThrow('Unsupported embedding device: invalid') + }) +}) + +describe('getEmbeddingIdentity', () => { + it('uses the resolved model repo and device', () => { + expect(getEmbeddingIdentity('bge-small-en-v1.5', 'webgpu')) + .toBe('Xenova/bge-small-en-v1.5@webgpu') + }) }) describe('embed device registry', () => { diff --git a/test/unit/embedding-cache-identity.test.ts b/test/unit/embedding-cache-identity.test.ts deleted file mode 100644 index 00b6c569..00000000 --- a/test/unit/embedding-cache-identity.test.ts +++ /dev/null @@ -1,103 +0,0 @@ -import { DatabaseSync } from 'node:sqlite' -import { describe, expect, it } from 'vitest' - -/** - * Guards the cache-invalidation rule in `src/retriv/embedding-cache.ts`. - * - * Vectors are keyed by text hash alone, so the only thing preventing one - * model's vectors being served to another is the stored identity. Dimensions - * are not enough: `Xenova/bge-large-en-v1.5` and `ollama:qwen3-embedding:0.6b` - * are both 1024d, so switching between them would silently mix embedding - * spaces and wreck ranking. - * - * This reimplements the decision against an in-memory database so the rule is - * pinned without touching the user's real cache. - */ -function applyIdentity(db: DatabaseSync, dimensions: number, model?: string): void { - const get = db.prepare('SELECT value FROM meta WHERE key = ?') - const set = db.prepare('INSERT OR REPLACE INTO meta (key, value) VALUES (?, ?)') - - const storedDims = get.get('dimensions') as { value: string } | undefined - const storedModel = get.get('model') as { value: string } | undefined - const dimsChanged = storedDims && Number(storedDims.value) !== dimensions - const modelChanged = model !== undefined && storedModel?.value !== model - - if (dimsChanged || modelChanged) - db.exec('DELETE FROM embeddings') - - set.run('dimensions', String(dimensions)) - if (model !== undefined) - set.run('model', model) -} - -function makeDb(): DatabaseSync { - const db = new DatabaseSync(':memory:') - db.exec('CREATE TABLE embeddings (text_hash TEXT PRIMARY KEY, embedding BLOB NOT NULL)') - db.exec('CREATE TABLE meta (key TEXT PRIMARY KEY, value TEXT NOT NULL)') - return db -} - -function seed(db: DatabaseSync, n = 3): void { - const stmt = db.prepare('INSERT OR IGNORE INTO embeddings (text_hash, embedding) VALUES (?, ?)') - for (let i = 0; i < n; i++) - stmt.run(`hash-${i}`, Buffer.from(new Float32Array([i, i, i]).buffer)) -} - -function count(db: DatabaseSync): number { - return (db.prepare('SELECT COUNT(*) c FROM embeddings').get() as { c: number }).c -} - -describe('embedding cache identity', () => { - it('keeps cached vectors when model and dimensions are unchanged', () => { - const db = makeDb() - applyIdentity(db, 1024, 'model-a') - seed(db) - applyIdentity(db, 1024, 'model-a') - expect(count(db)).toBe(3) - db.close() - }) - - // The regression: equal width, different model. - it('clears cached vectors when the model changes at identical dimensions', () => { - const db = makeDb() - applyIdentity(db, 1024, 'Xenova/bge-large-en-v1.5@webgpu') - seed(db) - expect(count(db)).toBe(3) - - applyIdentity(db, 1024, 'ollama:qwen3-embedding:0.6b') - expect(count(db)).toBe(0) - db.close() - }) - - it('clears cached vectors when dimensions change', () => { - const db = makeDb() - applyIdentity(db, 384, 'model-a') - seed(db) - applyIdentity(db, 1024, 'model-a') - expect(count(db)).toBe(0) - db.close() - }) - - // Same model on a different backend: numeric output can differ, so vectors - // are only interchangeable within a device. - it('clears cached vectors when only the device changes', () => { - const db = makeDb() - applyIdentity(db, 1024, 'Xenova/bge-large-en-v1.5@cpu') - seed(db) - applyIdentity(db, 1024, 'Xenova/bge-large-en-v1.5@webgpu') - expect(count(db)).toBe(0) - db.close() - }) - - // A cache written before the model key existed has unknown provenance. - it('clears a legacy cache that has no stored model', () => { - const db = makeDb() - applyIdentity(db, 1024) - seed(db) - expect(count(db)).toBe(3) - - applyIdentity(db, 1024, 'model-a') - expect(count(db)).toBe(0) - db.close() - }) -}) diff --git a/test/unit/embedding-cache.test.ts b/test/unit/embedding-cache.test.ts index fa613264..99215ecd 100644 --- a/test/unit/embedding-cache.test.ts +++ b/test/unit/embedding-cache.test.ts @@ -120,6 +120,56 @@ describe('embedding-cache', () => { expect((result[0] as Float32Array).length).toBe(8) }) + it('wipes cache when the model changes at the same dimension', async () => { + const { config: first } = fakeEmbeddingConfig(4, async texts => + texts.map(() => new Float32Array(4).fill(1))) + const firstWrapped = await cachedEmbeddings(first, 'model-a@cpu') + const { embedder: firstEmbedder } = await firstWrapped.resolve() + await firstEmbedder(['hello']) + + const { config: second, calls } = fakeEmbeddingConfig(4, async (texts) => { + calls.push(texts) + return texts.map(() => new Float32Array(4).fill(2)) + }) + const secondWrapped = await cachedEmbeddings(second, 'model-b@cpu') + const { embedder: secondEmbedder } = await secondWrapped.resolve() + + const result = await secondEmbedder(['hello']) + expect(calls).toEqual([['hello']]) + expect([...result[0] as Float32Array]).toEqual([2, 2, 2, 2]) + }) + + it('keeps cache when model and device are unchanged', async () => { + const { config: first } = fakeEmbeddingConfig(4, async texts => + texts.map(() => new Float32Array(4).fill(1))) + const firstWrapped = await cachedEmbeddings(first, 'model-a@cpu') + const { embedder: firstEmbedder } = await firstWrapped.resolve() + await firstEmbedder(['hello']) + + const { config: second, calls } = fakeEmbeddingConfig(4) + const secondWrapped = await cachedEmbeddings(second, 'model-a@cpu') + const { embedder: secondEmbedder } = await secondWrapped.resolve() + + const result = await secondEmbedder(['hello']) + expect(calls).toHaveLength(0) + expect([...result[0] as Float32Array]).toEqual([1, 1, 1, 1]) + }) + + it('wipes a cache with unknown legacy identity', async () => { + const { config: legacy } = fakeEmbeddingConfig(4, async texts => + texts.map(() => new Float32Array(4).fill(1))) + const legacyWrapped = await cachedEmbeddings(legacy) + const { embedder: legacyEmbedder } = await legacyWrapped.resolve() + await legacyEmbedder(['hello']) + + const { config: current, calls } = fakeEmbeddingConfig(4) + const currentWrapped = await cachedEmbeddings(current, 'model-a@cpu') + const { embedder: currentEmbedder } = await currentWrapped.resolve() + + await currentEmbedder(['hello']) + expect(calls).toEqual([['hello']]) + }) + it('clearEmbeddingCache removes the db file', async () => { const { config } = fakeEmbeddingConfig() const wrapped = await cachedEmbeddings(config) diff --git a/test/unit/index-embedding-identity.test.ts b/test/unit/index-embedding-identity.test.ts new file mode 100644 index 00000000..98e9a728 --- /dev/null +++ b/test/unit/index-embedding-identity.test.ts @@ -0,0 +1,79 @@ +import { DatabaseSync } from 'node:sqlite' +import { mkdtempSync, rmSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'pathe' +import { afterEach, describe, expect, it } from 'vitest' +import { EmbeddingIndexMismatchError, getDb } from '../../src/retriv/index.ts' +import { checkIndexEmbeddingIdentity, recordIndexEmbeddingIdentity } from '../../src/retriv/index-embedding-identity.ts' + +const dirs: string[] = [] + +function dbPath(): string { + const dir = mkdtempSync(join(tmpdir(), 'skilld-index-identity-')) + dirs.push(dir) + return join(dir, 'search.db') +} + +afterEach(() => { + for (const dir of dirs.splice(0)) + rmSync(dir, { recursive: true, force: true }) +}) + +describe('index embedding identity', () => { + it('rejects an equal-width model change', () => { + const path = dbPath() + recordIndexEmbeddingIdentity(path, 'Xenova/bge-large-en-v1.5@auto') + + expect(checkIndexEmbeddingIdentity(path, 'Xenova/bge-m3@auto')).toEqual({ + _tag: 'Mismatch', + current: 'Xenova/bge-m3@auto', + stored: 'Xenova/bge-large-en-v1.5@auto', + }) + }) + + it('blocks opening an index with a different model', async () => { + const path = dbPath() + recordIndexEmbeddingIdentity(path, 'Xenova/bge-large-en-v1.5@auto') + const originalModel = process.env.SKILLD_EMBED_MODEL + const originalDevice = process.env.SKILLD_EMBED_DEVICE + process.env.SKILLD_EMBED_MODEL = 'bge-m3' + process.env.SKILLD_EMBED_DEVICE = 'auto' + + try { + await expect(getDb({ dbPath: path })).rejects.toBeInstanceOf(EmbeddingIndexMismatchError) + } + finally { + if (originalModel === undefined) + delete process.env.SKILLD_EMBED_MODEL + else + process.env.SKILLD_EMBED_MODEL = originalModel + if (originalDevice === undefined) + delete process.env.SKILLD_EMBED_DEVICE + else + process.env.SKILLD_EMBED_DEVICE = originalDevice + } + }) + + it('treats existing pre-identity indexes as the old default', () => { + const path = dbPath() + const db = new DatabaseSync(path) + db.exec('CREATE TABLE documents_meta (id TEXT PRIMARY KEY)') + db.prepare('INSERT INTO documents_meta (id) VALUES (?)').run('doc') + db.close() + + expect(checkIndexEmbeddingIdentity(path, 'Xenova/bge-m3@auto')).toEqual({ + _tag: 'Mismatch', + current: 'Xenova/bge-m3@auto', + stored: 'Xenova/bge-small-en-v1.5@auto', + }) + }) + + it('accepts the identity recorded for a new index', () => { + const path = dbPath() + recordIndexEmbeddingIdentity(path, 'Xenova/bge-m3@webgpu') + + expect(checkIndexEmbeddingIdentity(path, 'Xenova/bge-m3@webgpu')).toEqual({ + _tag: 'Current', + }) + }) +}) diff --git a/test/unit/transformers-embeddings.test.ts b/test/unit/transformers-embeddings.test.ts new file mode 100644 index 00000000..fe4c0652 --- /dev/null +++ b/test/unit/transformers-embeddings.test.ts @@ -0,0 +1,34 @@ +import { describe, expect, it, vi } from 'vitest' + +const { pipeline } = vi.hoisted(() => ({ + pipeline: vi.fn(async () => async (texts: string[]) => ({ + data: new Float32Array(texts.length * 384), + })), +})) + +vi.mock('@huggingface/transformers', () => ({ + env: {}, + pipeline, +})) + +import { transformersEmbeddings } from '../../src/retriv/transformers-embeddings.ts' + +describe('transformersEmbeddings', () => { + it('runs the model on the selected device', async () => { + const embeddings = transformersEmbeddings({ + model: 'bge-small-en-v1.5', + device: 'webgpu', + }) + + const { embedder } = await embeddings.resolve() + const result = await embedder(['one', 'two']) + + expect(pipeline).toHaveBeenCalledWith( + 'feature-extraction', + 'Xenova/bge-small-en-v1.5', + { device: 'webgpu', dtype: 'fp32' }, + ) + expect(result).toHaveLength(2) + expect(result[0]).toHaveLength(384) + }) +}) From 97d67b1bb5e6352853340a3873bccc37e0e33cb3 Mon Sep 17 00:00:00 2001 From: Harlan Wilton Date: Wed, 12 Aug 2026 16:01:26 +1000 Subject: [PATCH 3/3] chore(search): use retriv device support Upgrade retriv to 0.15.0 and remove the temporary local Transformers.js provider now that device forwarding is available upstream. --- pnpm-lock.yaml | 91 +++++------------------ pnpm-workspace.yaml | 3 +- src/retriv/index.ts | 7 +- src/retriv/transformers-embeddings.ts | 63 ---------------- test/unit/transformers-embeddings.test.ts | 34 --------- 5 files changed, 23 insertions(+), 175 deletions(-) delete mode 100644 src/retriv/transformers-embeddings.ts delete mode 100644 test/unit/transformers-embeddings.test.ts diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index b9e7e140..cb923ee9 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -43,8 +43,8 @@ catalogs: specifier: ^0.3.23 version: 0.3.23 retriv: - specifier: ^0.14.7 - version: 0.14.7 + specifier: ^0.15.0 + version: 0.15.0 std-env: specifier: ^4.2.0 version: 4.2.0 @@ -167,7 +167,7 @@ importers: version: 2.0.3 retriv: specifier: 'catalog:' - version: 0.14.7(@huggingface/transformers@4.2.0)(sqlite-vec@0.1.9)(typescript@7.0.2) + version: 0.15.0(@huggingface/transformers@4.2.0)(sqlite-vec@0.1.9)(typescript@7.0.2) skilld-protocol: specifier: workspace:* version: link:packages/protocol @@ -434,27 +434,14 @@ packages: resolution: {integrity: sha512-sl4Bm6yiMNYrZKkqqDFWN0UfnWhlS8ivKxrYl+6t0gCLrqr8y3B2IqZZbFRkfaVVp7C/baApyh71P+LeE1A2sQ==} engines: {node: '>=18.0.0'} - '@babel/helper-string-parser@7.27.1': - resolution: {integrity: sha512-qMlSxKbpRlAridDExk92nSobyDdpPijUq2DW6oDnUqd0iOGxmQjyqhMIihI9+zv4LPyZdRje2cavWPbCbWm3eA==} - engines: {node: '>=6.9.0'} - '@babel/helper-string-parser@7.29.7': resolution: {integrity: sha512-Pb5ijPrZ89GDH8223L4UP8i6QApWxs04RbPQJTeWDV0/keR2E36MeKnyr6LYmUUvqRRI+Iv87SuF1W6ErINzYw==} engines: {node: '>=6.9.0'} - '@babel/helper-validator-identifier@7.28.5': - resolution: {integrity: sha512-qSs4ifwzKJSV39ucNjsvc6WVHs6b7S03sOh2OcHF9UHfVPqWWALUsNUVzhSBiItjRZoLHx7nIarVjqKVusUZ1Q==} - engines: {node: '>=6.9.0'} - '@babel/helper-validator-identifier@7.29.7': resolution: {integrity: sha512-qehxGkRj55h/ff8EMaJ+cYhyaKlHIxqYDn682wQD7RNp9UujOQsHog2uS0r2vzr4pW+sXf90NeeayjcNaX3fFg==} engines: {node: '>=6.9.0'} - '@babel/parser@7.29.3': - resolution: {integrity: sha512-b3ctpQwp+PROvU/cttc4OYl4MzfJUWy6FZg+PMXfzmt/+39iHVF0sDfqay8TQM3JA2EUOyKcFZt75jWriQijsA==} - engines: {node: '>=6.0.0'} - hasBin: true - '@babel/parser@7.29.8': resolution: {integrity: sha512-E8lTAYNB1KW+FH+VGJuZM1ioAx2E6oVlvQFRrf5P8ZZmsiJXYAD9vTFV7yyEURNzgh1dFqMZuO6tUwcARbqFCA==} engines: {node: '>=6.0.0'} @@ -464,10 +451,6 @@ packages: resolution: {integrity: sha512-Nq8OhGWiZIZGV6hLHoyAKLLcJihP/xFeBMGJoUrxTX2psI8dCifzLhZISFb+VWS3wFMRDmCGw5R+dOySCqPLhw==} engines: {node: '>=6.9.0'} - '@babel/types@7.29.0': - resolution: {integrity: sha512-LwdZHpScM4Qz8Xw2iKSzS+cfglZzJGvofQICy7W7v4caru4EaAmyUuO6BGrbyQ2mYV11W0U8j5mBhd14dd3B0A==} - engines: {node: '>=6.9.0'} - '@babel/types@7.29.8': resolution: {integrity: sha512-Vj1jF3cPfxg7OAfoI7QnVKLoILlm2JF9pnVHrX8qx7AHMiYWT+NDAA7jChlNgRS4WTLc/fD1lXLmPixluj+3Gg==} engines: {node: '>=6.9.0'} @@ -846,12 +829,6 @@ packages: peerDependencies: eslint: ^6.0.0 || ^7.0.0 || >=8.0.0 - '@eslint-community/eslint-utils@4.9.1': - resolution: {integrity: sha512-phrYmNiYppR7znFEdqgfWHXR6NCkZEK7hwWDHZUjit/2/U0r6XvkDl0SYnoM51Hq7FhCGdLDT6zxCCOY1hexsQ==} - engines: {node: ^12.22.0 || ^14.17.0 || >=16.0.0} - peerDependencies: - eslint: ^6.0.0 || ^7.0.0 || >=8.0.0 - '@eslint-community/regexpp@4.12.2': resolution: {integrity: sha512-EriSTlt5OC9/7SXkRSCAhfSxxoSUgBm33OH+IkwbdpgoqsSsUg7y3uh+IICI/Qg4BBWr3U2i39RpmycbxMq4ew==} engines: {node: ^12.0.0 || ^14.0.0 || >=16.0.0} @@ -2266,10 +2243,6 @@ packages: bowser@2.14.1: resolution: {integrity: sha512-tzPjzCxygAKWFOJP011oxFHs57HzIhOEracIgAePE4pqB3LikALKnSzUyU4MGs9/iCEUuHlAJTjTc5M+u7YEGg==} - brace-expansion@5.0.6: - resolution: {integrity: sha512-kLpxurY4Z4r9sgMsyG0Z9uzsBlgiU/EFKhj/h91/8yHu0edo7XuixOIH3VcJ8kkxs6/jPzoI6U9Vj3WqbMQ94g==} - engines: {node: 18 || 20 || >=22} - brace-expansion@5.0.9: resolution: {integrity: sha512-ScQ4IuvIEF1TMlP7Zt+vjJ//9zlPb2SDcxWxM3bk8s6t6GGdJ7KO1dCcTidOPJKePW30LE/2cT7wCyPho9/Wxg==} engines: {node: 20 || >=22} @@ -3272,10 +3245,6 @@ packages: resolution: {integrity: sha512-VP79XUPxV2CigYP3jWwAUFSku2aKqBH7uTAapFWCBqutsbmDo96KY5o8uh6U+/YSIn5OxJnXp73beVkpqMIGhA==} engines: {node: '>=18'} - minimatch@10.2.5: - resolution: {integrity: sha512-MULkVLfKGYDFYejP07QOurDLLQpcjk7Fw+7jXS2R2czRQzR56yHRveU5NDJEOviH+hETZKSkIk5c+T23GjFUMg==} - engines: {node: 18 || 20 || >=22} - minimatch@10.2.6: resolution: {integrity: sha512-vpLQEs+VLCr1nU0BXS07maYoFwlDAH0gngQuuttxIwutDFEMHq2blX+8vpgxDdK3J1PwjCJiep77OitTZ4Ll1A==} engines: {node: 18 || 20 || >=22} @@ -3553,18 +3522,18 @@ packages: resolution: {integrity: sha512-oMA2dcrw6u0YfxJQXm342bFKX/E4sG9rbTzO9ptUcR/e8A33cHuvStiYOwH7fszkZlZ1z/ta9AAoPk2F4qIOHA==} engines: {node: '>=18'} - retriv@0.14.7: - resolution: {integrity: sha512-j06MPUvbwLBp5XZHzYdZkBZG/IUmf1NH9M+DDLIVbZls4Q4wcOxQBSpPlCczAUzDDPg27IVo95gEAeaoPJSquw==} + retriv@0.15.0: + resolution: {integrity: sha512-Ya5hmjM1g8PJ3fCXDyCS44WBqOekMiv4Yr4Dspt29wqWUoY0j6LBPumRbqIOmYw8/u8z25C8KtubcfgBWGNVdg==} peerDependencies: - '@ai-sdk/cohere': ^3.0.0 - '@ai-sdk/google': ^3.0.0 - '@ai-sdk/mistral': ^3.0.0 - '@ai-sdk/openai': ^3.0.0 - '@huggingface/transformers': ^3.0.0 + '@ai-sdk/cohere': ^4.0.0 + '@ai-sdk/google': ^4.0.0 + '@ai-sdk/mistral': ^4.0.0 + '@ai-sdk/openai': ^4.0.0 + '@huggingface/transformers': ^3.0.0 || ^4.0.0 '@libsql/client': ^0.14.0 || ^0.15.0 || ^0.16.0 || ^0.17.0 '@upstash/vector': ^1.0.0 - ai: ^4.0.0 || ^5.0.0 || ^6.0.0 - ollama-ai-provider-v2: ^1.0.0 || ^2.0.0 || ^3.0.0 + ai: ^7.0.0 + ollama-ai-provider-v2: ^4.0.0 pg: ^8.0.0 sqlite-vec: ^0.1.0-alpha.0 typescript: ^5.0.0 || ^6.0.0-0 @@ -4453,29 +4422,16 @@ snapshots: '@aws/lambda-invoke-store@0.3.0': {} - '@babel/helper-string-parser@7.27.1': {} - '@babel/helper-string-parser@7.29.7': {} - '@babel/helper-validator-identifier@7.28.5': {} - '@babel/helper-validator-identifier@7.29.7': {} - '@babel/parser@7.29.3': - dependencies: - '@babel/types': 7.29.0 - '@babel/parser@7.29.8': dependencies: '@babel/types': 7.29.8 '@babel/runtime@7.29.7': {} - '@babel/types@7.29.0': - dependencies: - '@babel/helper-string-parser': 7.27.1 - '@babel/helper-validator-identifier': 7.28.5 - '@babel/types@7.29.8': dependencies: '@babel/helper-string-parser': 7.29.7 @@ -4719,11 +4675,6 @@ snapshots: eslint: 10.8.0(jiti@2.7.0)(supports-color@7.2.0) eslint-visitor-keys: 3.4.3 - '@eslint-community/eslint-utils@4.9.1(eslint@10.8.0(jiti@2.7.0)(supports-color@7.2.0))': - dependencies: - eslint: 10.8.0(jiti@2.7.0)(supports-color@7.2.0) - eslint-visitor-keys: 3.4.3 - '@eslint-community/regexpp@4.12.2': {} '@eslint/compat@2.1.0(eslint@10.8.0(jiti@2.7.0)(supports-color@7.2.0))': @@ -5441,7 +5392,7 @@ snapshots: '@typescript-eslint/types': 8.66.0 '@typescript-eslint/visitor-keys': 8.66.0 debug: 4.4.3(supports-color@7.2.0) - minimatch: 10.2.5 + minimatch: 10.2.6 semver: 7.8.5 tinyglobby: 0.2.17 ts-api-utils: 2.5.0(typescript@7.0.2) @@ -5451,7 +5402,7 @@ snapshots: '@typescript-eslint/utils@8.66.0(eslint@10.8.0(jiti@2.7.0)(supports-color@7.2.0))(supports-color@7.2.0)(typescript@7.0.2)': dependencies: - '@eslint-community/eslint-utils': 4.9.1(eslint@10.8.0(jiti@2.7.0)(supports-color@7.2.0)) + '@eslint-community/eslint-utils': 4.10.1(eslint@10.8.0(jiti@2.7.0)(supports-color@7.2.0)) '@typescript-eslint/scope-manager': 8.66.0 '@typescript-eslint/types': 8.66.0 '@typescript-eslint/typescript-estree': 8.66.0(supports-color@7.2.0)(typescript@7.0.2) @@ -5594,7 +5545,7 @@ snapshots: '@vue/compiler-core@3.5.27': dependencies: - '@babel/parser': 7.29.3 + '@babel/parser': 7.29.8 '@vue/shared': 3.5.27 entities: 7.0.1 estree-walker: 2.0.2 @@ -5607,7 +5558,7 @@ snapshots: '@vue/compiler-sfc@3.5.27': dependencies: - '@babel/parser': 7.29.3 + '@babel/parser': 7.29.8 '@vue/compiler-core': 3.5.27 '@vue/compiler-dom': 3.5.27 '@vue/compiler-ssr': 3.5.27 @@ -5757,10 +5708,6 @@ snapshots: bowser@2.14.1: {} - brace-expansion@5.0.6: - dependencies: - balanced-match: 4.0.4 - brace-expansion@5.0.9: dependencies: balanced-match: 4.0.4 @@ -7027,10 +6974,6 @@ snapshots: mimic-function@5.0.1: {} - minimatch@10.2.5: - dependencies: - brace-expansion: 5.0.6 - minimatch@10.2.6: dependencies: brace-expansion: 5.0.9 @@ -7337,7 +7280,7 @@ snapshots: onetime: 7.0.0 signal-exit: 4.1.0 - retriv@0.14.7(@huggingface/transformers@4.2.0)(sqlite-vec@0.1.9)(typescript@7.0.2): + retriv@0.15.0(@huggingface/transformers@4.2.0)(sqlite-vec@0.1.9)(typescript@7.0.2): optionalDependencies: '@huggingface/transformers': 4.2.0 sqlite-vec: 0.1.9 diff --git a/pnpm-workspace.yaml b/pnpm-workspace.yaml index 698b7894..98298b4e 100644 --- a/pnpm-workspace.yaml +++ b/pnpm-workspace.yaml @@ -2,6 +2,7 @@ catalogMode: prefer minimumReleaseAgeExclude: - verkit@0.2.0 - '@mdream/rust-wasm32-wasi@1.5.12' + - retriv@0.15.0 shellEmulator: true trustPolicy: no-downgrade @@ -23,7 +24,7 @@ catalog: ofetch: ^1.5.1 pathe: ^2.0.3 publint: ^0.3.23 - retriv: ^0.14.7 + retriv: ^0.15.0 std-env: ^4.2.0 tsx: ^4.23.8 typebox: ^1.3.10 diff --git a/src/retriv/index.ts b/src/retriv/index.ts index b5ce532a..1e12c931 100644 --- a/src/retriv/index.ts +++ b/src/retriv/index.ts @@ -3,7 +3,6 @@ import { readConfig } from '../core/config.ts' import { stripFrontmatter } from '../core/markdown.ts' import { checkIndexEmbeddingIdentity, recordIndexEmbeddingIdentity } from './index-embedding-identity.ts' import { getEmbeddingIdentity, resolveEmbedDevice, resolveEmbedModel } from './models.ts' -import { transformersEmbeddings } from './transformers-embeddings.ts' export type { ChunkEntity, Document, IndexConfig, IndexPhase, IndexProgress, SearchFilter, SearchOptions, SearchResult, SearchSnippet } @@ -67,19 +66,21 @@ export async function getDb(config: Pick) { if (identityState._tag === 'Mismatch') throw new EmbeddingIndexMismatchError(config.dbPath, identityState.stored, identityState.current) - let createRetriv, autoChunker, sqliteMod, sqliteVec, cachedEmbeddings + let createRetriv, autoChunker, sqliteMod, sqliteVec, transformersJs, cachedEmbeddings try { ;([ { createRetriv }, { autoChunker }, sqliteMod, sqliteVec, + { transformersJs }, { cachedEmbeddings }, ] = await Promise.all([ import('retriv'), import('retriv/chunkers/auto'), import('retriv/db/sqlite'), import('sqlite-vec'), + import('retriv/embeddings/transformers-js'), import('./embedding-cache.ts'), ])) } @@ -89,7 +90,7 @@ export async function getDb(config: Pick) { throw err } const embeddings = await cachedEmbeddings( - transformersEmbeddings({ + transformersJs({ model: embedModel, ...(device ? { device } : {}), }), diff --git a/src/retriv/transformers-embeddings.ts b/src/retriv/transformers-embeddings.ts deleted file mode 100644 index 33809b40..00000000 --- a/src/retriv/transformers-embeddings.ts +++ /dev/null @@ -1,63 +0,0 @@ -import type { EmbeddingConfig } from 'retriv' -import { rm } from 'node:fs/promises' -import { resolve } from 'pathe' -import { getModelDimensions, getModelMaxTokens, resolveModelForPreset } from 'retriv/embeddings/model-info' -import type { RuntimeEmbedDevice } from './models.ts' - -export interface TransformersEmbeddingOptions { - model: string - device?: RuntimeEmbedDevice -} - -/** Transformers.js provider with explicit device support. */ -export function transformersEmbeddings(options: TransformersEmbeddingOptions): EmbeddingConfig { - const model = resolveModelForPreset(options.model, 'transformers.js') - let cached: Awaited> | undefined - - return { - async resolve() { - if (cached) - return cached - - // Search is optional, so keep the model runtime outside the main CLI bundle. - const { env, pipeline } = await import('@huggingface/transformers') - const load = () => pipeline('feature-extraction', model, { - dtype: 'fp32', - ...(options.device ? { device: options.device } : {}), - }) - const extractor = await load().catch(async (error) => { - const corrupted = error instanceof Error - && (error.message.includes('Protobuf parsing failed') || String(error.cause).includes('Protobuf parsing failed')) - if (!corrupted || !env.cacheDir) - throw error - const cacheRoot = resolve(env.cacheDir) - const modelCache = resolve(cacheRoot, model) - if (!modelCache.startsWith(`${cacheRoot}/`)) - throw error - await rm(modelCache, { recursive: true, force: true }) - console.warn(`[skilld] Cleared corrupted model cache for ${model}, retrying...`) - return load() - }) - - const dimensions = getModelDimensions(model) - if (!dimensions) - throw new Error(`Unknown dimensions for model ${model}.`) - - const embedder = async (texts: string[]) => { - const output = await extractor(texts, { pooling: 'mean', normalize: true }) - const data = output.data as Float32Array - return Array.from( - { length: texts.length }, - (_, index) => data.slice(index * dimensions, (index + 1) * dimensions), - ) - } - - cached = { - embedder, - dimensions, - maxTokens: getModelMaxTokens(model), - } - return cached - }, - } -} diff --git a/test/unit/transformers-embeddings.test.ts b/test/unit/transformers-embeddings.test.ts deleted file mode 100644 index fe4c0652..00000000 --- a/test/unit/transformers-embeddings.test.ts +++ /dev/null @@ -1,34 +0,0 @@ -import { describe, expect, it, vi } from 'vitest' - -const { pipeline } = vi.hoisted(() => ({ - pipeline: vi.fn(async () => async (texts: string[]) => ({ - data: new Float32Array(texts.length * 384), - })), -})) - -vi.mock('@huggingface/transformers', () => ({ - env: {}, - pipeline, -})) - -import { transformersEmbeddings } from '../../src/retriv/transformers-embeddings.ts' - -describe('transformersEmbeddings', () => { - it('runs the model on the selected device', async () => { - const embeddings = transformersEmbeddings({ - model: 'bge-small-en-v1.5', - device: 'webgpu', - }) - - const { embedder } = await embeddings.resolve() - const result = await embedder(['one', 'two']) - - expect(pipeline).toHaveBeenCalledWith( - 'feature-extraction', - 'Xenova/bge-small-en-v1.5', - { device: 'webgpu', dtype: 'fp32' }, - ) - expect(result).toHaveLength(2) - expect(result[0]).toHaveLength(384) - }) -})