sweet-search 0.0.1 → 2.3.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (161) hide show
  1. package/LICENSE +190 -0
  2. package/NOTICE +23 -0
  3. package/core/cli.js +51 -0
  4. package/core/config.js +27 -0
  5. package/core/embedding/embedding-cache.js +467 -0
  6. package/core/embedding/embedding-local-model.js +845 -0
  7. package/core/embedding/embedding-remote.js +492 -0
  8. package/core/embedding/embedding-service.js +712 -0
  9. package/core/embedding/embedding-telemetry.js +219 -0
  10. package/core/embedding/index.js +40 -0
  11. package/core/graph/community-detector.js +294 -0
  12. package/core/graph/graph-expansion.js +839 -0
  13. package/core/graph/graph-extractor.js +2304 -0
  14. package/core/graph/graph-search.js +2148 -0
  15. package/core/graph/hcgs-generator.js +666 -0
  16. package/core/graph/index.js +16 -0
  17. package/core/graph/leiden-algorithm.js +547 -0
  18. package/core/graph/relationship-resolver.js +366 -0
  19. package/core/graph/repo-map.js +408 -0
  20. package/core/graph/summary-manager.js +549 -0
  21. package/core/indexing/artifact-builder.js +1054 -0
  22. package/core/indexing/ast-chunker.js +709 -0
  23. package/core/indexing/chunking/chunk-builder.js +170 -0
  24. package/core/indexing/chunking/markdown-chunker.js +503 -0
  25. package/core/indexing/chunking/plaintext-chunker.js +104 -0
  26. package/core/indexing/dedup/dedup-phase.js +159 -0
  27. package/core/indexing/dedup/exemplar-selector.js +65 -0
  28. package/core/indexing/document-chunker.js +56 -0
  29. package/core/indexing/incremental-parser.js +390 -0
  30. package/core/indexing/incremental-tracker.js +761 -0
  31. package/core/indexing/index-codebase-v21.js +472 -0
  32. package/core/indexing/index-maintainer.mjs +1674 -0
  33. package/core/indexing/index.js +90 -0
  34. package/core/indexing/indexer-ann.js +1077 -0
  35. package/core/indexing/indexer-build.js +742 -0
  36. package/core/indexing/indexer-phases.js +800 -0
  37. package/core/indexing/indexer-pool.js +764 -0
  38. package/core/indexing/indexer-sparse-gram.js +98 -0
  39. package/core/indexing/indexer-utils.js +536 -0
  40. package/core/indexing/indexer-worker.js +148 -0
  41. package/core/indexing/li-skip-policy.js +225 -0
  42. package/core/indexing/merkle-tracker.js +244 -0
  43. package/core/indexing/model-pool.js +166 -0
  44. package/core/infrastructure/code-graph-repository.js +120 -0
  45. package/core/infrastructure/codebase-repository.js +131 -0
  46. package/core/infrastructure/config/dedup.js +54 -0
  47. package/core/infrastructure/config/embedding.js +298 -0
  48. package/core/infrastructure/config/graph.js +80 -0
  49. package/core/infrastructure/config/index.js +82 -0
  50. package/core/infrastructure/config/indexing.js +8 -0
  51. package/core/infrastructure/config/platform.js +254 -0
  52. package/core/infrastructure/config/ranking.js +221 -0
  53. package/core/infrastructure/config/search.js +396 -0
  54. package/core/infrastructure/config/translation.js +89 -0
  55. package/core/infrastructure/config/vector-store.js +114 -0
  56. package/core/infrastructure/constants.js +86 -0
  57. package/core/infrastructure/coreml-cascade.js +909 -0
  58. package/core/infrastructure/coreml-cascade.json +46 -0
  59. package/core/infrastructure/coreml-provider.js +81 -0
  60. package/core/infrastructure/db-utils.js +69 -0
  61. package/core/infrastructure/dedup-hashing.js +83 -0
  62. package/core/infrastructure/hardware-capability.js +332 -0
  63. package/core/infrastructure/index.js +104 -0
  64. package/core/infrastructure/language-patterns/maps.js +121 -0
  65. package/core/infrastructure/language-patterns/registry-core.js +323 -0
  66. package/core/infrastructure/language-patterns/registry-data-query.js +155 -0
  67. package/core/infrastructure/language-patterns/registry-object-oriented.js +285 -0
  68. package/core/infrastructure/language-patterns/registry-tooling.js +240 -0
  69. package/core/infrastructure/language-patterns/registry-web-style.js +143 -0
  70. package/core/infrastructure/language-patterns/registry.js +19 -0
  71. package/core/infrastructure/language-patterns.js +141 -0
  72. package/core/infrastructure/llm-provider.js +733 -0
  73. package/core/infrastructure/manifest.json +46 -0
  74. package/core/infrastructure/maxsim.wasm +0 -0
  75. package/core/infrastructure/model-fetcher.js +423 -0
  76. package/core/infrastructure/model-registry.js +214 -0
  77. package/core/infrastructure/native-inference.js +587 -0
  78. package/core/infrastructure/native-resolver.js +187 -0
  79. package/core/infrastructure/native-sparse-gram.js +257 -0
  80. package/core/infrastructure/native-tokenizer.js +160 -0
  81. package/core/infrastructure/onnx-mutex.js +45 -0
  82. package/core/infrastructure/onnx-session-utils.js +261 -0
  83. package/core/infrastructure/ort-pipeline.js +111 -0
  84. package/core/infrastructure/project-detector.js +102 -0
  85. package/core/infrastructure/quantization.js +410 -0
  86. package/core/infrastructure/simd-distance.js +502 -0
  87. package/core/infrastructure/simd-distance.wasm +0 -0
  88. package/core/infrastructure/tree-sitter-provider.js +665 -0
  89. package/core/infrastructure/webgpu-maxsim.js +222 -0
  90. package/core/query/index.js +35 -0
  91. package/core/query/intent-detector.js +201 -0
  92. package/core/query/intent-router.js +156 -0
  93. package/core/query/query-router-catboost.js +222 -0
  94. package/core/query/query-router-ml.js +266 -0
  95. package/core/query/query-router.js +213 -0
  96. package/core/ranking/cascaded-scorer.js +379 -0
  97. package/core/ranking/flashrank.js +810 -0
  98. package/core/ranking/index.js +49 -0
  99. package/core/ranking/late-interaction-index.js +2383 -0
  100. package/core/ranking/late-interaction-model.js +812 -0
  101. package/core/ranking/local-reranker.js +374 -0
  102. package/core/ranking/mmr.js +379 -0
  103. package/core/ranking/quality-scorer.js +363 -0
  104. package/core/search/context-expander.js +1167 -0
  105. package/core/search/dedup/sibling-expander.js +327 -0
  106. package/core/search/index.js +16 -0
  107. package/core/search/search-boost.js +259 -0
  108. package/core/search/search-cli.js +544 -0
  109. package/core/search/search-format.js +282 -0
  110. package/core/search/search-fusion.js +327 -0
  111. package/core/search/search-hybrid.js +204 -0
  112. package/core/search/search-pattern-chunks.js +337 -0
  113. package/core/search/search-pattern-planner.js +439 -0
  114. package/core/search/search-pattern-prefilter.js +412 -0
  115. package/core/search/search-pattern-ripgrep.js +663 -0
  116. package/core/search/search-pattern.js +463 -0
  117. package/core/search/search-postprocess.js +452 -0
  118. package/core/search/search-semantic.js +706 -0
  119. package/core/search/search-server.js +554 -0
  120. package/core/search/session-daemon-prewarm.mjs +164 -0
  121. package/core/search/session-warmup.js +595 -0
  122. package/core/search/sweet-search.js +632 -0
  123. package/core/search/warmup-metrics.js +532 -0
  124. package/core/start-server.js +6 -0
  125. package/core/training/query-router/features/extractor.js +762 -0
  126. package/core/training/query-router/features/multilingual-patterns.js +431 -0
  127. package/core/training/query-router/features/text-segmenter.js +303 -0
  128. package/core/training/query-router/features/unicode-utils.js +383 -0
  129. package/core/training/query-router/output/v45_router_d4.js +11521 -0
  130. package/core/training/query-router/output/v46_router_d4.js +11498 -0
  131. package/core/vector-store/binary-heap.js +227 -0
  132. package/core/vector-store/binary-hnsw-index.js +1004 -0
  133. package/core/vector-store/float-vector-store.js +234 -0
  134. package/core/vector-store/hnsw-index.js +580 -0
  135. package/core/vector-store/index.js +39 -0
  136. package/core/vector-store/seismic-index.js +498 -0
  137. package/core/vocabulary/index.js +84 -0
  138. package/core/vocabulary/vocab-constants.js +20 -0
  139. package/core/vocabulary/vocab-miner-extractors.js +375 -0
  140. package/core/vocabulary/vocab-miner-nl.js +404 -0
  141. package/core/vocabulary/vocab-miner-utils.js +146 -0
  142. package/core/vocabulary/vocab-miner.js +574 -0
  143. package/core/vocabulary/vocab-prewarm-cli.js +110 -0
  144. package/core/vocabulary/vocab-ranker.js +492 -0
  145. package/core/vocabulary/vocab-warmer.js +523 -0
  146. package/core/vocabulary/vocab-warmup-orchestrator.js +425 -0
  147. package/core/vocabulary/vocabulary-utils.js +704 -0
  148. package/crates/wasm-router/pkg/package.json +13 -0
  149. package/crates/wasm-router/pkg/query_router_wasm.d.ts +36 -0
  150. package/crates/wasm-router/pkg/query_router_wasm.js +271 -0
  151. package/crates/wasm-router/pkg/query_router_wasm_bg.wasm +0 -0
  152. package/crates/wasm-router/pkg/query_router_wasm_bg.wasm.d.ts +19 -0
  153. package/mcp/config-gen.js +121 -0
  154. package/mcp/server.js +335 -0
  155. package/mcp/tool-handlers.js +476 -0
  156. package/package.json +131 -9
  157. package/scripts/benchmark-harness.js +794 -0
  158. package/scripts/init.js +1058 -0
  159. package/scripts/smoke-test.js +435 -0
  160. package/scripts/uninstall.js +478 -0
  161. package/scripts/verify-runtime.js +176 -0
@@ -0,0 +1,587 @@
1
+ /**
2
+ * Native Inference — candle-based FP32 embedding + late interaction inference.
3
+ *
4
+ * Provides native model wrappers that serve as drop-in replacements for the
5
+ * ORT inference paths. Uses the napi-rs addon with Metal GPU acceleration on
6
+ * Apple Silicon, CPU fallback elsewhere.
7
+ *
8
+ * On M3+ Apple Silicon, the addon also loads a CoreML variant cascade
9
+ * alongside the candle backbone. This module is the single point where
10
+ * the cascade dirs resolved by `coreml-cascade.js` are handed to the
11
+ * Rust `NativeEmbeddingModel::load` / `NativeLateInteractionModel::load`
12
+ * constructors. Lower hardware, no cache, or
13
+ * `SWEET_SEARCH_COREML_CASCADE=0` collapses to the candle-only path
14
+ * transparently.
15
+ *
16
+ * Environment:
17
+ * SWEET_SEARCH_NATIVE_INFERENCE=0|1 — force disable/enable
18
+ * SWEET_SEARCH_COREML_CASCADE=0 — force-disable the cascade
19
+ * even when hardware +
20
+ * cache are both eligible
21
+ * (diagnostic / benchmarking)
22
+ * SWEET_SEARCH_COREML_STATS=1 — dump per-variant
23
+ * dispatch report on
24
+ * addon shutdown (see
25
+ * coreml_embedding.rs)
26
+ * SWEET_SEARCH_CUDA=0 — force-disable CUDA even
27
+ * on an eligible host
28
+ * (diagnostic)
29
+ * SWEET_SEARCH_CUDA_COMPUTE_CAP=<cc> — set by this module
30
+ * before addon loads so
31
+ * the Rust dtype policy
32
+ * picks BF16/F16/F32 by
33
+ * compute capability.
34
+ * See mod.rs::optimal_dtype
35
+ * CANDLE_METAL_COMPUTE_PER_BUFFER=<N> — candle default 50 (tuned)
36
+ * CANDLE_METAL_COMMAND_POOL_SIZE=<N> — candle default 5 (tuned)
37
+ */
38
+
39
+ import { existsSync } from 'fs';
40
+ import { join } from 'path';
41
+ import { createRequire } from 'module';
42
+ import { resolveNativeAddon } from './native-resolver.js';
43
+ import { createTokenizer } from './native-tokenizer.js';
44
+ import { getModelCacheDir, fetchModel } from './model-fetcher.js';
45
+ import { getModelEntry } from './model-registry.js';
46
+ import { getCoremlCascadeResolvedDirs } from './coreml-cascade.js';
47
+ import { detectHardwareCapability } from './hardware-capability.js';
48
+
49
+ const require = createRequire(import.meta.url);
50
+
51
+ // ─── State ───
52
+
53
+ let _addon = null;
54
+ let _embeddingModel = null;
55
+ let _embeddingModelLoadPromise = null; // race-gate for concurrent first calls
56
+ let _liModel = null;
57
+ let _liModelLoadPromise = null;
58
+ let _embTokenizer = null;
59
+ let _embTokenizerLoadPromise = null;
60
+ let _liTokenizer = null;
61
+ let _liTokenizerLoadPromise = null;
62
+ let _available = null;
63
+ let _coremlCascadeLogged = false;
64
+
65
+ // ─── Cascade-dir selection (pure helper, unit-tested) ───
66
+
67
+ /**
68
+ * Decide whether the Rust addon should load a CoreML cascade for this
69
+ * deviceKind. Only `'metal'` triggers cascade resolution; `'cuda'`, `'cpu'`,
70
+ * and `'auto'` skip cascade (CUDA has no cascade, CPU has no GPU dispatch).
71
+ *
72
+ * Pure function — takes a resolver for the cascade dirs and the override.
73
+ * Returns `undefined` (explicit "no cascade") or a string path.
74
+ *
75
+ * @param {'cpu'|'metal'|'cuda'|'auto'} deviceKind
76
+ * @param {string|undefined} cascadeDirOverride - explicit caller override (wins)
77
+ * @param {() => string|null|undefined} resolveCascadeDir - called only when needed
78
+ */
79
+ export function pickCascadeDirForDevice(deviceKind, cascadeDirOverride, resolveCascadeDir) {
80
+ if (cascadeDirOverride !== undefined) return cascadeDirOverride;
81
+ if (deviceKind !== 'metal') return undefined;
82
+ return resolveCascadeDir() || undefined;
83
+ }
84
+
85
+ // ─── CUDA compute-capability env propagation ───
86
+
87
+ /**
88
+ * Ensure `SWEET_SEARCH_CUDA_COMPUTE_CAP` is set for the current process
89
+ * before the addon loads a CUDA model. The Rust `optimal_dtype` reads
90
+ * this env var to pick BF16 on Ampere+ and F16/F32 on older GPUs.
91
+ *
92
+ * Idempotent: honors an already-set value (useful for forcing a dtype
93
+ * tier in benchmarks) and silently no-ops when there is no NVIDIA GPU.
94
+ */
95
+ function propagateCudaComputeCapToAddonEnv() {
96
+ if (process.env.SWEET_SEARCH_CUDA_COMPUTE_CAP) return;
97
+ try {
98
+ const hw = detectHardwareCapability();
99
+ if (hw.nvidiaGpu?.computeCapability) {
100
+ process.env.SWEET_SEARCH_CUDA_COMPUTE_CAP = hw.nvidiaGpu.computeCapability;
101
+ }
102
+ } catch {
103
+ // Hardware detection is advisory. If it fails, the Rust side falls
104
+ // back to F32 which is always correct, if slower.
105
+ }
106
+ }
107
+
108
+ // ─── CoreML cascade resolution ───
109
+
110
+ /**
111
+ * Resolve which CoreML cascade dirs the Rust addon should try to load.
112
+ * Logged exactly once per process so a mis-configured cascade surfaces
113
+ * at startup instead of silently falling through on every call.
114
+ *
115
+ * Always returns an object — never throws. The returned dirs can be
116
+ * `null`, which the Rust addon treats as "CoreML path disabled" and
117
+ * falls back to candle unconditionally.
118
+ */
119
+ function resolveCoremlCascadeForAddon() {
120
+ const resolved = getCoremlCascadeResolvedDirs();
121
+ if (!_coremlCascadeLogged) {
122
+ _coremlCascadeLogged = true;
123
+ const hw = detectHardwareCapability();
124
+ // Log a single line so every subsequent `[NativeInference]` message
125
+ // can be correlated with the cascade decision made here.
126
+ if (resolved.embedDir || resolved.liDir) {
127
+ process.stderr.write(
128
+ `[NativeInference] CoreML cascade: ${resolved.status}` +
129
+ ` (embed=${resolved.embedDir ? 'yes' : 'no'}, li=${resolved.liDir ? 'yes' : 'no'},` +
130
+ ` chip=${hw.brandString || 'unknown'})\n`
131
+ );
132
+ } else if (hw.coremlCascadeEligible) {
133
+ process.stderr.write(
134
+ `[NativeInference] CoreML cascade: ${resolved.status} —` +
135
+ ` hardware eligible but cache missing. Run \`node scripts/build-coreml-cascade.js\`` +
136
+ ` to enable the 18% speedup on ${hw.brandString}\n`
137
+ );
138
+ }
139
+ // Ineligible hardware: silent. The user has no action to take and
140
+ // we already log the candle device via the existing "Loaded"
141
+ // line below.
142
+ }
143
+ return resolved;
144
+ }
145
+
146
+ // ─── Addon Loading ───
147
+
148
+ function loadAddon() {
149
+ if (_addon) return _addon;
150
+ const addonPath = resolveNativeAddon();
151
+ if (!addonPath) return null;
152
+ try {
153
+ _addon = require(addonPath);
154
+ return _addon;
155
+ } catch {
156
+ return null;
157
+ }
158
+ }
159
+
160
+ // ─── Detection ───
161
+
162
+ /**
163
+ * Check if native inference is available and not disabled via env var.
164
+ * Caches result after first check.
165
+ */
166
+ export function isNativeInferenceAvailable() {
167
+ if (_available !== null) return _available;
168
+
169
+ const envFlag = (process.env.SWEET_SEARCH_NATIVE_INFERENCE ?? '').trim().toLowerCase();
170
+ if (envFlag === '0' || envFlag === 'false' || envFlag === 'off') {
171
+ _available = false;
172
+ return false;
173
+ }
174
+
175
+ const addon = loadAddon();
176
+ _available = typeof addon?.NativeEmbeddingModel?.load === 'function';
177
+ return _available;
178
+ }
179
+
180
+ // ─── Embedding Model ───
181
+
182
+ /**
183
+ * Load the native embedding model (CodeRankEmbed FP32 safetensors).
184
+ * Returns the model instance or null if unavailable.
185
+ *
186
+ * Race-gated: concurrent first calls share a single load promise so the
187
+ * underlying napi `addon.NativeEmbeddingModel.load(...)` runs exactly once.
188
+ * Without this gate, multiple parallel queries (e.g. eval runner with
189
+ * --concurrency=N) all see _embeddingModel == null and each load a fresh
190
+ * model copy, wasting Metal memory and printing N "Loaded" lines.
191
+ */
192
+ export async function getNativeEmbeddingModel() {
193
+ if (_embeddingModel) return _embeddingModel;
194
+ if (_embeddingModelLoadPromise) return _embeddingModelLoadPromise;
195
+
196
+ _embeddingModelLoadPromise = (async () => {
197
+ const addon = loadAddon();
198
+ if (!addon?.NativeEmbeddingModel) return null;
199
+
200
+ await fetchModel('coderankembed-fp32');
201
+
202
+ const entry = getModelEntry('coderankembed-fp32');
203
+ const modelDir = getModelCacheDir(entry.hfId);
204
+ const safetensorsPath = join(modelDir, 'model.safetensors');
205
+ const configPath = join(modelDir, 'config.json');
206
+
207
+ if (!existsSync(safetensorsPath) || !existsSync(configPath)) return null;
208
+
209
+ // Resolve the CoreML cascade dir for NomicBERT embeddings. Returns
210
+ // null on ineligible hardware, empty cache, or explicit opt-out
211
+ // (SWEET_SEARCH_COREML_CASCADE=0). The Rust addon's constructor
212
+ // runs a parity check against the smallest variant before arming
213
+ // the dispatch path; failure there falls back to candle.
214
+ const cascade = resolveCoremlCascadeForAddon();
215
+
216
+ const t0 = Date.now();
217
+ _embeddingModel = addon.NativeEmbeddingModel.load(
218
+ safetensorsPath,
219
+ configPath,
220
+ cascade.embedDir || undefined,
221
+ );
222
+ console.log(`[NativeInference] Embedding model loaded in ${Date.now() - t0}ms (dim: ${_embeddingModel.dim}, device: ${addon.nativeInferenceDevice()})`);
223
+
224
+ return _embeddingModel;
225
+ })();
226
+
227
+ try {
228
+ return await _embeddingModelLoadPromise;
229
+ } finally {
230
+ // Clear the promise on resolve so a future explicit unload can re-load.
231
+ // Keep it set on success: subsequent calls hit the _embeddingModel cache
232
+ // first and never reach the promise check.
233
+ }
234
+ }
235
+
236
+ /**
237
+ * Get or create the embedding tokenizer. Race-gated like getNativeEmbeddingModel.
238
+ * Uses the INT8 model's tokenizer (same vocab as FP32).
239
+ */
240
+ async function getEmbTokenizer() {
241
+ if (_embTokenizer) return _embTokenizer;
242
+ if (_embTokenizerLoadPromise) return _embTokenizerLoadPromise;
243
+ _embTokenizerLoadPromise = (async () => {
244
+ const entry = getModelEntry('coderankembed-int8');
245
+ const tokenizerPath = join(getModelCacheDir(entry.hfId), 'tokenizer.json');
246
+ _embTokenizer = await createTokenizer(tokenizerPath);
247
+ return _embTokenizer;
248
+ })();
249
+ return _embTokenizerLoadPromise;
250
+ }
251
+
252
+ /**
253
+ * Tokenize texts for the native model, returning napi-compatible arrays.
254
+ */
255
+ function tokenizedToNapi(tokenized, batchSize, seqLen) {
256
+ const inputIds = new Array(batchSize);
257
+ const attentionMask = new Array(batchSize);
258
+ for (let b = 0; b < batchSize; b++) {
259
+ const ids = new Array(seqLen);
260
+ const mask = new Array(seqLen);
261
+ const base = b * seqLen;
262
+ for (let s = 0; s < seqLen; s++) {
263
+ ids[s] = Number(tokenized.input_ids.data[base + s]);
264
+ mask[s] = Number(tokenized.attention_mask.data[base + s]);
265
+ }
266
+ inputIds[b] = ids;
267
+ attentionMask[b] = mask;
268
+ }
269
+ return { inputIds, attentionMask };
270
+ }
271
+
272
+ /**
273
+ * Native embedding inference — drop-in replacement for callLocalModel().
274
+ *
275
+ * @param {string[]} texts - Array of texts to embed
276
+ * @param {object} [options]
277
+ * @param {number} [options.maxLength=512] - Max sequence length
278
+ * @returns {Float32Array[]} Array of L2-normalized embedding vectors (768d)
279
+ */
280
+ export async function nativeEmbed(texts, options = {}) {
281
+ const { maxLength = 512 } = options;
282
+ const model = await getNativeEmbeddingModel();
283
+ if (!model) throw new Error('[NativeInference] Embedding model not loaded');
284
+
285
+ const tokenizer = await getEmbTokenizer();
286
+ const tokenized = tokenizer(texts, { padding: true, truncation: true, max_length: maxLength });
287
+ const batchSize = texts.length;
288
+ const seqLen = tokenized.input_ids.dims[1];
289
+
290
+ const { inputIds, attentionMask } = tokenizedToNapi(tokenized, batchSize, seqLen);
291
+ // The native addon returns a flat `Float32Array` of length batch * dim —
292
+ // see the comment on embed_batch in embedding_model.rs. Use `.subarray()`
293
+ // to return zero-copy views over the shared backing buffer (2026-04-15 perf
294
+ // fix — previous `.slice()` was a per-row memcpy for no reason).
295
+ //
296
+ // Callers MUST treat each view as read-only. Mutation audit (2026-04-15):
297
+ // - callLocalModelGpu at core/embedding/embedding-local-model.js:547
298
+ // returns the array unchanged to its caller
299
+ // - downstream embedding-service paths call truncateForHNSW, which
300
+ // allocates a fresh Float32Array per row (quantization.js)
301
+ // - SQLite persistence paths take a BLOB copy via better-sqlite3
302
+ // - no grep hit for `embeddings[i][j] = ` or `.set(` on the returned rows
303
+ // Any future consumer that writes through a returned view will silently
304
+ // corrupt neighbouring rows in the shared buffer. If you need to mutate,
305
+ // copy first with `new Float32Array(view)` or `view.slice()`.
306
+ const flat = await model.embedBatch(inputIds, attentionMask);
307
+ const dim = model.dim;
308
+
309
+ const embeddings = new Array(batchSize);
310
+ for (let i = 0; i < batchSize; i++) {
311
+ embeddings[i] = flat.subarray(i * dim, (i + 1) * dim);
312
+ }
313
+ return embeddings;
314
+ }
315
+
316
+ // ─── Late Interaction Model ───
317
+
318
+ /**
319
+ * Load the native LI model (LateOn-Code FP32 safetensors + projection).
320
+ * Returns the model instance or null if unavailable. Race-gated.
321
+ */
322
+ export async function getNativeLiModel() {
323
+ if (_liModel) return _liModel;
324
+ if (_liModelLoadPromise) return _liModelLoadPromise;
325
+ _liModelLoadPromise = (async () => {
326
+ const addon = loadAddon();
327
+ if (!addon?.NativeLateInteractionModel) return null;
328
+
329
+ await fetchModel('lateon-code-fp32');
330
+
331
+ const entry = getModelEntry('lateon-code-fp32');
332
+ const modelDir = getModelCacheDir(entry.hfId);
333
+ const backbonePath = join(modelDir, 'model.safetensors');
334
+ const projPath = join(modelDir, '1_Dense', 'model.safetensors');
335
+ const configPath = join(modelDir, 'config.json');
336
+
337
+ if (!existsSync(backbonePath) || !existsSync(projPath) || !existsSync(configPath)) return null;
338
+
339
+ // Resolve the CoreML cascade dir for ModernBERT LI. Same contract
340
+ // as the embedding model above — see that comment.
341
+ const cascade = resolveCoremlCascadeForAddon();
342
+
343
+ const t0 = Date.now();
344
+ _liModel = addon.NativeLateInteractionModel.load(
345
+ backbonePath,
346
+ projPath,
347
+ configPath,
348
+ cascade.liDir || undefined,
349
+ );
350
+ console.log(`[NativeInference] LI model loaded in ${Date.now() - t0}ms (dim: ${_liModel.dim}, device: ${addon.nativeInferenceDevice()})`);
351
+
352
+ return _liModel;
353
+ })();
354
+ return _liModelLoadPromise;
355
+ }
356
+
357
+ /**
358
+ * Get or create the LI tokenizer. Race-gated.
359
+ */
360
+ async function getLiTokenizer() {
361
+ if (_liTokenizer) return _liTokenizer;
362
+ if (_liTokenizerLoadPromise) return _liTokenizerLoadPromise;
363
+ _liTokenizerLoadPromise = (async () => {
364
+ const entry = getModelEntry('lateon-code');
365
+ const tokenizerPath = join(getModelCacheDir(entry.hfId), 'tokenizer.json');
366
+ _liTokenizer = await createTokenizer(tokenizerPath);
367
+ return _liTokenizer;
368
+ })();
369
+ return _liTokenizerLoadPromise;
370
+ }
371
+
372
+ /**
373
+ * Native LI encoding from a pre-tokenized batch. Avoids re-tokenizing when
374
+ * the caller already has the tokenizer output (e.g. because it needs
375
+ * input_ids for skiplist filtering). Accepts the object returned by the
376
+ * native tokenizer: `{ input_ids: { data, dims }, attention_mask: { data, dims } }`.
377
+ *
378
+ * @param {object} tokenized - Pre-tokenized batch from native tokenizer
379
+ * @returns {Float32Array[][]} Per-document arrays of Float32Array token vectors
380
+ */
381
+ export async function nativeLiEncodeTokenized(tokenized) {
382
+ const model = await getNativeLiModel();
383
+ if (!model) throw new Error('[NativeInference] LI model not loaded');
384
+
385
+ const [batchSize, seqLen] = tokenized.input_ids.dims;
386
+ const { inputIds, attentionMask } = tokenizedToNapi(tokenized, batchSize, seqLen);
387
+ const result = await model.encodeBatch(inputIds, attentionMask);
388
+ const dim = model.dim;
389
+
390
+ // The native addon returns `vectors` as a flat `Float32Array` (zero-copy
391
+ // from the Rust Vec<f32> via napi typed array) — see encode_batch in
392
+ // li_model.rs. Use `.subarray()` to return zero-copy views over the shared
393
+ // backing buffer (2026-04-15 perf fix — previous `.slice()` was a per-token
394
+ // memcpy for no reason).
395
+ //
396
+ // LI token views are READ-ONLY. Mutation audit (2026-04-15):
397
+ // - nativeLiEncodeTokenized is called from late-interaction-model.js:387
398
+ // (encodeDocumentsGpu) which pushes views into a JS array and hands
399
+ // them to indexer-ann.js::finalizeBatchResults
400
+ // - finalizeBatchResults calls liIndex.add which at
401
+ // late-interaction-index.js:429-430 does
402
+ // `tokens.slice(...).map(emb => emb.slice(0, tokenDim))` — the inner
403
+ // `.slice()` creates a fresh Float32Array per token (copy)
404
+ // - poolTokens at late-interaction-model.js:619 does
405
+ // `new Float32Array(clusterInput[i])` (copy via constructor)
406
+ // - persistence goes through _writeSegmentFile which copies bytes into
407
+ // a pre-allocated Buffer
408
+ // Any future consumer that writes through a returned view will corrupt
409
+ // the underlying batch buffer and neighbouring tokens. Copy before mutating.
410
+ const flat = result.vectors;
411
+
412
+ const allVectors = new Array(batchSize);
413
+ let offset = 0;
414
+ for (let b = 0; b < batchSize; b++) {
415
+ const count = result.tokenCounts[b];
416
+ const docVectors = new Array(count);
417
+ for (let t = 0; t < count; t++) {
418
+ docVectors[t] = flat.subarray(offset + t * dim, offset + (t + 1) * dim);
419
+ }
420
+ allVectors[b] = docVectors;
421
+ offset += count * dim;
422
+ }
423
+
424
+ return allVectors;
425
+ }
426
+
427
+ /**
428
+ * Native LI encoding — returns per-token L2-normalized 128d vectors.
429
+ *
430
+ * @param {string[]} texts - Texts to encode (already prefixed with [Q]/[D])
431
+ * @param {object} [options]
432
+ * @param {number} [options.maxLength=2048] - Max sequence length
433
+ * @returns {Float32Array[][]} Per-document arrays of Float32Array token vectors
434
+ */
435
+ export async function nativeLiEncode(texts, options = {}) {
436
+ const { maxLength = 2048 } = options;
437
+ const tokenizer = await getLiTokenizer();
438
+ const tokenized = tokenizer(texts, { padding: true, truncation: true, max_length: maxLength });
439
+ return nativeLiEncodeTokenized(tokenized);
440
+ }
441
+
442
+ // ─── Model state queries ───
443
+
444
+ export function isNativeEmbeddingModelLoaded() {
445
+ return _embeddingModel != null;
446
+ }
447
+
448
+ export function isNativeLiModelLoaded() {
449
+ return _liModel != null;
450
+ }
451
+
452
+ // ─── Device-explicit loading ───
453
+
454
+ /**
455
+ * Load the native embedding model on a specific device ("cpu" | "metal" | "auto").
456
+ * Used by model-pool.js to arm GPU models for indexing. Sets the module-scope
457
+ * singleton so subsequent nativeEmbed() calls use this instance.
458
+ */
459
+ export async function loadNativeEmbeddingModelWithDevice(deviceKind, cascadeDirOverride) {
460
+ if (_embeddingModel) return _embeddingModel;
461
+ if (_embeddingModelLoadPromise) return _embeddingModelLoadPromise;
462
+
463
+ _embeddingModelLoadPromise = (async () => {
464
+ const addon = loadAddon();
465
+ if (!addon?.NativeEmbeddingModel?.loadWithDevice) return null;
466
+
467
+ // CUDA needs the compute-capability env set BEFORE the Rust load
468
+ // path picks a dtype. No-op on non-CUDA kinds and already-set envs.
469
+ if (deviceKind === 'cuda') propagateCudaComputeCapToAddonEnv();
470
+
471
+ await fetchModel('coderankembed-fp32');
472
+
473
+ const entry = getModelEntry('coderankembed-fp32');
474
+ const modelDir = getModelCacheDir(entry.hfId);
475
+ const safetensorsPath = join(modelDir, 'model.safetensors');
476
+ const configPath = join(modelDir, 'config.json');
477
+
478
+ if (!existsSync(safetensorsPath) || !existsSync(configPath)) return null;
479
+
480
+ // Cascade is a CoreML-specific concept. CUDA has no cascade — the Rust
481
+ // addon JITs CUDA kernels per forward pass; there is no `.mlpackage`
482
+ // equivalent to prefetch. Skip cascade resolution entirely on CUDA so
483
+ // we don't log "cascade missing" warnings on a Linux+NVIDIA host.
484
+ const cascadeDir = pickCascadeDirForDevice(
485
+ deviceKind,
486
+ cascadeDirOverride,
487
+ () => resolveCoremlCascadeForAddon().embedDir,
488
+ );
489
+
490
+ const t0 = Date.now();
491
+ _embeddingModel = addon.NativeEmbeddingModel.loadWithDevice(
492
+ safetensorsPath,
493
+ configPath,
494
+ cascadeDir,
495
+ deviceKind,
496
+ );
497
+ console.log(`[NativeInference] Embedding model loaded in ${Date.now() - t0}ms (dim: ${_embeddingModel.dim}, device: ${deviceKind})`);
498
+
499
+ return _embeddingModel;
500
+ })();
501
+
502
+ try {
503
+ return await _embeddingModelLoadPromise;
504
+ } finally {
505
+ // Keep promise set on success; clear only on re-load.
506
+ }
507
+ }
508
+
509
+ /**
510
+ * Load the native LI model on a specific device.
511
+ */
512
+ export async function loadNativeLiModelWithDevice(deviceKind, cascadeDirOverride) {
513
+ if (_liModel) return _liModel;
514
+ if (_liModelLoadPromise) return _liModelLoadPromise;
515
+
516
+ _liModelLoadPromise = (async () => {
517
+ const addon = loadAddon();
518
+ if (!addon?.NativeLateInteractionModel?.loadWithDevice) return null;
519
+
520
+ // See loadNativeEmbeddingModelWithDevice for why this is CUDA-only.
521
+ if (deviceKind === 'cuda') propagateCudaComputeCapToAddonEnv();
522
+
523
+ await fetchModel('lateon-code-fp32');
524
+
525
+ const entry = getModelEntry('lateon-code-fp32');
526
+ const modelDir = getModelCacheDir(entry.hfId);
527
+ const backbonePath = join(modelDir, 'model.safetensors');
528
+ const projPath = join(modelDir, '1_Dense', 'model.safetensors');
529
+ const configPath = join(modelDir, 'config.json');
530
+
531
+ if (!existsSync(backbonePath) || !existsSync(projPath) || !existsSync(configPath)) return null;
532
+
533
+ // CUDA has no cascade — see the matching comment in
534
+ // loadNativeEmbeddingModelWithDevice.
535
+ const cascadeDir = pickCascadeDirForDevice(
536
+ deviceKind,
537
+ cascadeDirOverride,
538
+ () => resolveCoremlCascadeForAddon().liDir,
539
+ );
540
+
541
+ const t0 = Date.now();
542
+ _liModel = addon.NativeLateInteractionModel.loadWithDevice(
543
+ backbonePath,
544
+ projPath,
545
+ configPath,
546
+ cascadeDir,
547
+ deviceKind,
548
+ );
549
+ console.log(`[NativeInference] LI model loaded in ${Date.now() - t0}ms (dim: ${_liModel.dim}, device: ${deviceKind})`);
550
+
551
+ return _liModel;
552
+ })();
553
+
554
+ return _liModelLoadPromise;
555
+ }
556
+
557
+ // ─── Warmup primitives ───
558
+
559
+ export async function warmupNativeEmbeddingModel() {
560
+ if (!_embeddingModel?.warmupForward) return;
561
+ const t0 = Date.now();
562
+ await _embeddingModel.warmupForward();
563
+ console.log(`[NativeInference] Embedding warmup forward in ${Date.now() - t0}ms`);
564
+ }
565
+
566
+ export async function warmupNativeLiModel() {
567
+ if (!_liModel?.warmupForward) return;
568
+ const t0 = Date.now();
569
+ await _liModel.warmupForward();
570
+ console.log(`[NativeInference] LI warmup forward in ${Date.now() - t0}ms`);
571
+ }
572
+
573
+ // ─── Cleanup ───
574
+
575
+ export function unloadNativeModels() {
576
+ _embeddingModel = null;
577
+ _embeddingModelLoadPromise = null;
578
+ _liModel = null;
579
+ _liModelLoadPromise = null;
580
+ _embTokenizer = null;
581
+ _embTokenizerLoadPromise = null;
582
+ _liTokenizer = null;
583
+ _liTokenizerLoadPromise = null;
584
+ _addon = null;
585
+ _available = null;
586
+ _coremlCascadeLogged = false;
587
+ }