diff --git a/README.md b/README.md index f93ebc15..576ed5ae 100644 --- a/README.md +++ b/README.md @@ -239,6 +239,59 @@ 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` 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 | +|-------|-----------:|-------| +| `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: + +```bash +export SKILLD_EMBED_MODEL=bge-m3 +skilld update --force +``` + +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 +``` + +### 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 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. Set `SKILLD_EMBED_DEVICE` to override the saved setting: + +```bash +export 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/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/commands/config.ts b/src/commands/config.ts index 20d76282..3f1ac873 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, 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,73 @@ 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 for indexing and querying 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 + } + + updateConfig({ embedModel: choice === DEFAULT_EMBED_MODEL ? undefined : choice as string }) + p.log.success(`Embedding model set to ${choice}`) + p.log.warn('Embedding model changed. Rebuild existing search indexes: 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 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', + 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, switch back to Auto. The backend may be unavailable on this machine.') +} + 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..0b7efc88 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) { 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.') + 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, transformersJs, cachedEmbeddings try { ;([ @@ -70,8 +89,14 @@ export async function getDb(config: Pick) { throw new SearchDepsUnavailableError(err) throw err } - const embeddings = await cachedEmbeddings(transformersJs()) - return createRetriv({ + const embeddings = await cachedEmbeddings( + transformersJs({ + model: embedModel, + ...(device ? { device } : {}), + }), + embeddingIdentity, + ) + const db = await createRetriv({ driver: sqliteMod.default({ path: config.dbPath, embeddings, @@ -79,6 +104,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 new file mode 100644 index 00000000..b667c4a5 --- /dev/null +++ b/src/retriv/models.ts @@ -0,0 +1,137 @@ +import { resolveModelForPreset } from 'retriv/embeddings/model-info' + +/** + * Local embedding models available to the search index. + * + * 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 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 */ + 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` overrides saved config. The configured value overrides + * 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 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 + label: string + hint: string +} + +export const DEFAULT_EMBED_DEVICE = 'auto' + +export const EMBED_DEVICES = [ + { + 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', + }, +] 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) +} + +/** + * 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): RuntimeEmbedDevice | undefined { + const fromEnv = process.env.SKILLD_EMBED_DEVICE?.trim() + const value = fromEnv || configured || DEFAULT_EMBED_DEVICE + 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/test/unit/embed-models.test.ts b/test/unit/embed-models.test.ts new file mode 100644 index 00000000..f3c698c3 --- /dev/null +++ b/test/unit/embed-models.test.ts @@ -0,0 +1,145 @@ +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, getEmbeddingIdentity, 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() + }) + + 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', () => { + 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.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', + }) + }) +})