@lincoln504/pi-research 0.1.12 → 1.0.0

package/src/knowledge/chunker.ts

@@ -0,0 +1,157 @@
+export interface Chunk {
+  text: string;
+  actual_overlap: number;
+}
+
+export interface ChunkerOptions {
+  targetSize: number;
+  overlap: number;
+}
+
+import { metrics } from '../utils/metrics.ts';
+
+export class Chunker {
+  private targetSize: number;
+  private overlap: number;
+
+  constructor(options: ChunkerOptions) {
+    if (options.overlap >= options.targetSize) {
+      throw new Error(
+        `Chunker: overlap (${options.overlap}) must be less than targetSize (${options.targetSize}). ` +
+        `This prevents infinite loops during chunking.`
+      );
+    }
+    this.targetSize = options.targetSize;
+    this.overlap = options.overlap;
+  }
+
+  chunk(text: string): Chunk[] {
+    if (!text) {
+      metrics.increment('chunker_operations_total', 1, { status: 'empty' });
+      return [];
+    }
+
+    const startTime = Date.now();
+    const chunks: Chunk[] = [];
+    let start = 0;
+    let prevEnd = 0;
+    let codeBlockExtensions = 0;
+    let headingExtensions = 0;
+    let sentenceExtensions = 0;
+    let newlineExtensions = 0;
+    let spaceExtensions = 0;
+
+    while (start < text.length) {
+      let end = start + this.targetSize;
+      
+      if (end < text.length) {
+        let slice = text.slice(start, end);
+        let extendedForCodeBlock = false;
+        
+        // 0. Code blocks atomicity
+        const textBefore = text.slice(0, start);
+        const codeBlockMatchesBefore = textBefore.match(/```/g);
+        const startsInCodeBlock = codeBlockMatchesBefore && codeBlockMatchesBefore.length % 2 !== 0;
+
+        // Hard cap: never extend a chunk beyond 4x targetSize (min 2000 chars) for code blocks.
+        // A very large code block (e.g. a package compatibility table) can otherwise
+        // produce chunks of 12000+ tokens that OOM the embedding model at inference time.
+        const MAX_CHUNK_CHARS = Math.max(this.targetSize * 4, 2000);
+
+        if (startsInCodeBlock) {
+          const nextEnd = text.indexOf('```', start);
+          if (nextEnd !== -1 && nextEnd + 3 - start <= MAX_CHUNK_CHARS) {
+            end = nextEnd + 3;
+            slice = text.slice(start, end);
+            extendedForCodeBlock = true;
+            codeBlockExtensions++;
+          }
+        } else {
+          const codeBlockMatchesInSlice = slice.match(/```/g);
+          if (codeBlockMatchesInSlice && codeBlockMatchesInSlice.length % 2 !== 0) {
+            const nextEnd = text.indexOf('```', start + slice.lastIndexOf('```') + 3);
+            if (nextEnd !== -1 && nextEnd + 3 - start <= MAX_CHUNK_CHARS) {
+              end = nextEnd + 3;
+              slice = text.slice(start, end);
+              extendedForCodeBlock = true;
+              codeBlockExtensions++;
+            }
+          }
+        }
+
+        if (end < text.length && !extendedForCodeBlock) {
+          const lastHeading = slice.lastIndexOf('\n#');
+          if (lastHeading !== -1 && lastHeading > this.targetSize * 0.4) {
+            end = start + lastHeading + 1;
+            headingExtensions++;
+          } else {
+            const sentenceMatches = [...slice.matchAll(/[.!?](?=\s|\n)/g)];
+            let lastSentencePos = -1;
+            if (sentenceMatches.length > 0) {
+              const lastMatch = sentenceMatches[sentenceMatches.length - 1];
+              if (lastMatch && lastMatch.index !== undefined) {
+                lastSentencePos = lastMatch.index + 1;
+              }
+            }
+
+            if (lastSentencePos !== -1 && lastSentencePos > this.targetSize * 0.6) {
+              end = start + lastSentencePos;
+              sentenceExtensions++;
+            } else {
+              const lastNL = slice.lastIndexOf('\n');
+              if (lastNL !== -1 && lastNL > this.targetSize * 0.7) {
+                end = start + lastNL + 1;
+                newlineExtensions++;
+              } else {
+                const lastSpace = slice.lastIndexOf(' ');
+                if (lastSpace !== -1 && lastSpace > this.targetSize * 0.8) {
+                  end = start + lastSpace + 1;
+                  spaceExtensions++;
+                }
+              }
+            }
+          }
+        }
+      } else {
+        end = text.length;
+      }
+
+      if (end <= start) end = Math.min(start + this.targetSize, text.length);
+
+      const chunkText = text.slice(start, end);
+      const actual_overlap = chunks.length === 0 ? 0 : Math.max(0, Math.min(chunkText.length, prevEnd - start));
+      
+      chunks.push({ text: chunkText, actual_overlap });
+      
+      if (end >= text.length) break;
+      
+      prevEnd = Math.max(prevEnd, end);
+      start = end - this.overlap;
+      if (start <= (end - chunkText.length)) {
+        start = end - Math.min(this.overlap, chunkText.length - 1);
+      }
+      if (start < 0) start = 0;
+      if (start >= end) start = end - 1; 
+    }
+
+    const duration = Date.now() - startTime;
+    metrics.observe('chunker_duration_ms', duration);
+    metrics.increment('chunker_operations_total', 1, { status: 'success' });
+    metrics.increment('chunker_chunks_generated_total', chunks.length);
+    metrics.increment('chunker_extensions_total', codeBlockExtensions, { type: 'code_block' });
+    metrics.increment('chunker_extensions_total', headingExtensions, { type: 'heading' });
+    metrics.increment('chunker_extensions_total', sentenceExtensions, { type: 'sentence' });
+    metrics.increment('chunker_extensions_total', newlineExtensions, { type: 'newline' });
+    metrics.increment('chunker_extensions_total', spaceExtensions, { type: 'space' });
+
+    // Track chunk size distribution
+    for (const chunk of chunks) {
+      metrics.observe('chunker_chunk_size_chars', chunk.text.length);
+      if (chunk.actual_overlap > 0) {
+        metrics.observe('chunker_overlap_chars', chunk.actual_overlap);
+      }
+    }
+
+    return chunks;
+  }
+}

package/src/knowledge/embedder-init.ts

@@ -0,0 +1,413 @@
+/**
+ * Embedder Initialization Logic
+ *
+ * Handles initialization and warmup logic for the embedder
+ */
+
+import type { FeatureExtractionPipeline } from '@huggingface/transformers';
+import { logger } from '../logger.ts';
+import type { IStateManager } from '../core/service-interfaces.ts';
+import { DisposablePipeline } from '../core/interfaces/knowledge-interfaces.ts';
+import { withTimeout, markWebGpuFallback } from './embedder-utils.ts';
+import { getHFEnv } from './onnx-env.ts';
+import { getTransformers } from './transformers-loader.ts';
+
+/**
+ * Load pipeline with timeout
+ */
+export async function loadPipelineWithTimeout(
+  model: string,
+  device: string,
+  timeoutMs: number,
+  useCache: boolean
+): Promise<{ pipeline: FeatureExtractionPipeline; errorMessage: string }> {
+  const errorMessage = `Model load timed out after ${timeoutMs}ms. Check network connection or try a smaller model.`;
+
+  // pipeline() must be created INSIDE runCapturingStderr so that the FD-level
+  // stderr redirect is in place before the native ONNX/Dawn C++ code runs.
+  // Creating the Promise outside means the background thread could write its
+  // Dawn limit-clamping warnings before captureStdio redirects FD 2.
+  const loadedPipeline = await logger.runCapturingStderr(async () => {
+    const { pipeline } = await getTransformers();
+    const pipelinePromise = pipeline('feature-extraction', model, {
+      device: device as 'webgpu' | 'cpu' | 'auto' | 'gpu' | 'wasm' | 'webnn' | 'webnn-npu' | 'webnn-gpu' | 'webnn-cpu',
+      ...(useCache === false ? { use_cache: false } : {}),
+      // Clamp ONNX intra-op thread pool to 2 threads per session.
+      // Default (0) = one thread per physical CPU core. With multiple concurrent
+      // processes each loading their own pipeline, the default spawns N_cores * N_procs
+      // threads that all busy-spin simultaneously, saturating the CPU.
+      // 2 threads gives adequate within-op parallelism without thrashing.
+      session_options: {
+        intraOpNumThreads: 2,
+        interOpNumThreads: 1,
+      },
+    });
+    return await withTimeout(pipelinePromise, timeoutMs, errorMessage);
+  });
+
+  return { pipeline: loadedPipeline, errorMessage };
+}
+
+/**
+ * Warmup the pipeline
+ */
+export async function warmupPipeline(
+  pipeline: FeatureExtractionPipeline,
+  poolingMode: 'mean' | 'cls' | 'last_token',
+  useCache: boolean
+): Promise<{ dummy: any; success: boolean; error?: Error }> {
+  try {
+    const dummy = await withTimeout(
+      pipeline('warmup', { 
+        pooling: poolingMode as any, 
+        normalize: false,
+        ...(useCache === false ? { use_cache: false } : {}),
+      }),
+      20_000,
+      'Model warmup timed out after 20000ms.'
+    );
+    return { dummy, success: true };
+  } catch (warmupErr) {
+    return { dummy: null, success: false, error: warmupErr as Error };
+  }
+}
+
+/**
+ * Check if error is a WebGPU device error
+ */
+export function isWebGpuDeviceError(err: unknown): boolean {
+  if (!err) return false;
+
+  // Type-safe string extraction from error objects
+  let msg: string;
+  let stack: string;
+
+  if (err instanceof Error) {
+    msg = err.message ?? '';
+    stack = err.stack ?? '';
+  } else if (typeof err === 'object') {
+    msg = String((err as Record<string, unknown>)['message'] ?? '');
+    stack = String((err as Record<string, unknown>)['stack'] ?? '');
+    if (!msg && !stack) {
+      try { msg = JSON.stringify(err); } catch { msg = String(err); }
+    }
+  } else {
+    msg = String(err);
+    stack = '';
+  }
+
+  const combined = (msg + ' ' + stack).toLowerCase();
+  
+  return (
+    combined.includes('webgpu') ||
+    combined.includes('out_of_device_memory') ||
+    combined.includes('out of memory') ||
+    combined.includes('vk_error_out_of_device_memory') ||
+    combined.includes('vkallocatememory') ||
+    combined.includes('device lost') ||
+    combined.includes('devicelost') ||
+    combined.includes('validation failed') ||
+    combined.includes('invalid buffer') ||
+    combined.includes('bindgroup') ||
+    combined.includes('minbindingsize') ||
+    combined.includes('past_key_values') ||
+    combined.includes('unlabeled') ||
+    combined.includes('bufferbindingtype') ||
+    combined.includes('createdevice') ||
+    combined.includes('createbindgroup') ||
+    combined.includes('out-of-memory')
+  );
+}
+
+/**
+ * Check if a model is FULLY cached on disk.
+ *
+ * Verifying only that `model.onnx` exists is not enough on two axes:
+ *
+ *  1. WEIGHTS: large models keep their weights in a sibling external-data file
+ *     (`model.onnx_data`). An interrupted download can leave the small graph file
+ *     (`model.onnx`) complete while the weights file is missing, zero-length, or a
+ *     leftover partial-download artifact — which passes a presence-only check and then
+ *     crashes the ONNX loader ("Deserialize tensor … file_length … out of bounds").
+ *  2. ROOT METADATA: transformers.js also needs `config.json` and a tokenizer at the
+ *     MODEL ROOT (not in `onnx/`). A truncated/garbage `config.json` (external disk
+ *     error, partial delete) otherwise passes a weights-only check and then fails to
+ *     parse with an error isCorruptModelError does NOT match — a permanent, un-self-
+ *     healing init failure on every startup. Requiring them here re-fetches instead.
+ *
+ * Returning false routes the caller back through a fresh download, which heals the gap.
+ * A leftover `.tmp` artifact alongside COMPLETE weights is just trash from a prior
+ * interrupted attempt that later succeeded — it is swept (best-effort) and does not
+ * defeat the cached fast-path (which would otherwise re-download every startup).
+ */
+export async function isModelCached(model: string): Promise<boolean> {
+  try {
+    const env = await getHFEnv();
+    const cacheDir = env.cacheDir;
+    if (!cacheDir) return false;
+    const { access, readdir, readFile, stat, rm } = await import('node:fs/promises');
+    const path = await import('node:path');
+    const modelDir = path.default.join(cacheDir, model);
+    const onnxDir = path.default.join(modelDir, 'onnx');
+
+    // Graph file must exist.
+    await access(path.default.join(onnxDir, 'model.onnx'));
+
+    // Root config.json must exist AND parse — catches a missing/empty/truncated/garbage
+    // config (the un-self-healing corruption case). It is tiny, so the parse is cheap.
+    try {
+      const cfgRaw = await readFile(path.default.join(modelDir, 'config.json'), 'utf-8');
+      JSON.parse(cfgRaw);
+    } catch {
+      return false;
+    }
+
+    // A tokenizer must be present and non-empty. Accept any recognized form so this never
+    // false-negatives across the supported models (fast tokenizer.json) or a custom model
+    // (sentencepiece tokenizer.model / WordPiece vocab.txt).
+    const tokenizerFiles = ['tokenizer.json', 'tokenizer.model', 'vocab.txt'];
+    let hasTokenizer = false;
+    for (const f of tokenizerFiles) {
+      const st = await stat(path.default.join(modelDir, f)).catch(() => null);
+      if (st && st.size > 0) { hasTokenizer = true; break; }
+    }
+    if (!hasTokenizer) return false;
+
+    const entries = await readdir(onnxDir).catch(() => [] as string[]);
+
+    // If an external weights file is listed, it must be non-empty — an interrupted
+    // download can leave a zero-length model.onnx_data that crashes the ONNX loader.
+    const hasExternalWeights = entries.includes('model.onnx_data');
+    if (hasExternalWeights) {
+      const dataStat = await stat(path.default.join(onnxDir, 'model.onnx_data')).catch(() => null);
+      if (!dataStat || dataStat.size === 0) return false;
+    }
+
+    const partialArtifacts = entries.filter(
+      (f) => /\.(incomplete|downloading|part|tmp)$/i.test(f) || f.includes('.tmp.')
+    );
+    if (partialArtifacts.length > 0) {
+      // A leftover partial-download artifact. Treat it as trash to sweep ONLY when the
+      // real weights are confirmed present and complete (model.onnx_data non-empty,
+      // checked above) — a prior interrupted attempt that later succeeded. With no
+      // confirmed external-weights file the partial IS an in-flight/interrupted weights
+      // download, so the cache is genuinely incomplete → re-download.
+      if (hasExternalWeights) {
+        for (const f of partialArtifacts) {
+          await rm(path.default.join(onnxDir, f), { force: true }).catch(() => {});
+        }
+      } else {
+        return false;
+      }
+    }
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+/**
+ * Whether a model-load failure indicates a corrupt/truncated on-disk cache (as opposed to a
+ * network, timeout, OOM, opset-mismatch, or GPU error). These are unrecoverable by retrying the
+ * same files — the only fix is to purge the cache and re-download.
+ *
+ * Match ONLY corruption-specific signals. ONNX Runtime wraps essentially every session-creation
+ * failure in two GENERIC strings — `Load model from <uri> failed: <ex>` and `Failed to load
+ * model with error: <ex>` — and also emits broad `out of bounds` / `out_of_range` for non-cache
+ * conditions (CPU OOM during weight load, unregistered-op/opset drift, shape inference, WebGPU
+ * buffer-offset checks). Matching those would purge + re-download a perfectly good multi-hundred-MB
+ * cache on a transient error, and PERMANENTLY destroy it if the host is offline. The genuine
+ * truncated-weights error ("Deserialize tensor … file_length … out of bounds") is still caught
+ * here via 'deserialize'/'file_length', so narrowing loses no real corruption coverage.
+ */
+export function isCorruptModelError(err: unknown): boolean {
+  const msg = (err instanceof Error ? (err.message + ' ' + (err.stack ?? '')) : String(err)).toLowerCase();
+  return (
+    msg.includes('deserialize') ||
+    msg.includes('file_length') ||
+    msg.includes('protobuf parsing failed') ||
+    msg.includes('invalid protobuf')
+  );
+}
+
+/**
+ * Delete a model's on-disk cache directory so the next load re-downloads it from scratch.
+ * Best-effort: a deletion failure is logged and swallowed (the subsequent re-download attempt
+ * will surface any real, persistent problem).
+ */
+export async function purgeModelCache(model: string): Promise<void> {
+  try {
+    const env = await getHFEnv();
+    const cacheDir = env.cacheDir;
+    if (!cacheDir) return;
+    const { rm } = await import('node:fs/promises');
+    const path = await import('node:path');
+    // Containment guard: the consequence here is a recursive delete, so never let an empty,
+    // absolute, or traversal model id collapse `modelDir` to the cache root (which would wipe
+    // every model) or escape it. `model` is infrastructure config (EMBEDDING_MODEL), not
+    // attacker-derived, but validate in depth regardless.
+    const resolvedCache = path.default.resolve(cacheDir);
+    const modelDir = path.default.resolve(resolvedCache, model);
+    if (!model || path.default.isAbsolute(model) || modelDir === resolvedCache || !modelDir.startsWith(resolvedCache + path.default.sep)) {
+      logger.warn(`[embedder] Refusing to purge unsafe model-cache path for model id ${JSON.stringify(model)}`);
+      return;
+    }
+    await rm(modelDir, { recursive: true, force: true });
+    logger.warn(`[embedder] Purged corrupt model cache at ${modelDir}`);
+  } catch (err) {
+    logger.warn('[embedder] Failed to purge model cache:', err);
+  }
+}
+
+/**
+ * Load model on CPU (fallback from WebGPU)
+ */
+export async function loadModelOnCPU(
+  model: string,
+  initializationTimeoutMs: number,
+  useCache: boolean
+): Promise<FeatureExtractionPipeline> {
+  logger.info(`[embedder] Loading model on CPU after WebGPU error...`);
+  
+  const loadedPipeline = await logger.runCapturingStderr(async () => {
+    const { pipeline } = await getTransformers();
+    const pipelinePromise = pipeline('feature-extraction', model, {
+      device: 'cpu',
+      ...(useCache === false ? { use_cache: false } : {}),
+      session_options: {
+        intraOpNumThreads: 2,
+        interOpNumThreads: 1,
+      },
+    });
+    return await withTimeout(pipelinePromise, initializationTimeoutMs, 'CPU fallback model load timed out');
+  });
+  
+  return loadedPipeline;
+}
+
+/**
+ * Release GPU lock if held
+ */
+export async function releaseGpuLock(stateManager: IStateManager | null, gpuLockHeld: boolean): Promise<void> {
+  if (gpuLockHeld && stateManager) {
+    await stateManager.releaseGpuLock().catch(err => {
+      logger.warn('[embedder] Failed to release GPU lock:', err);
+    });
+  }
+}
+
+/**
+ * Acquire GPU lock for initialization
+ */
+export async function acquireGpuLock(
+  stateManager: IStateManager | null
+): Promise<{ acquired: boolean; shouldFallback: boolean }> {
+  if (!stateManager) {
+    return { acquired: false, shouldFallback: false };
+  }
+
+  const gpuLockHeld = await stateManager.acquireGpuLock(undefined, 30_000);
+  if (!gpuLockHeld) {
+    logger.warn('[embedder] Failed to acquire GPU init lock within 30s — falling back to CPU');
+    return { acquired: false, shouldFallback: true };
+  }
+  
+  logger.debug('[embedder] Acquired GPU init lock');
+  return { acquired: true, shouldFallback: false };
+}
+
+/**
+ * Handle WebGPU error during load
+ */
+export async function handleWebGPULoadError(
+  loadErr: unknown,
+  pipeline: FeatureExtractionPipeline | null,
+  stateManager: IStateManager | null,
+  gpuLockHeld: boolean,
+  model: string,
+  initializationTimeoutMs: number,
+  useCache: boolean
+): Promise<{ success: boolean; pipeline?: FeatureExtractionPipeline; error?: Error; dummy?: any }> {
+  const errorMsg = loadErr instanceof Error ? loadErr.message : String(loadErr);
+  const isValidationError = errorMsg.toLowerCase().includes('validation');
+  
+  markWebGpuFallback();
+  
+  if (isValidationError) {
+    logger.warn('[embedder] WebGPU validation error during pipeline loading — falling back to CPU');
+    logger.debug('[embedder] Validation error details:', errorMsg);
+  } else {
+    logger.warn('[embedder] WebGPU OOM during pipeline loading — falling back to CPU');
+  }
+  
+  // Clean up
+  if (pipeline) {
+    try { if (typeof (pipeline as any).dispose === 'function') await (pipeline as DisposablePipeline).dispose(); } catch (err) { logger.warn('[embedder] Error disposing pipeline:', err); }
+  }
+  await releaseGpuLock(stateManager, gpuLockHeld);
+
+  // Load on CPU instead
+  try {
+    const cpuPipeline = await loadModelOnCPU(model, initializationTimeoutMs, useCache);
+    logger.info(`[embedder] Pipeline loaded (device: cpu)`);
+    return { success: true, pipeline: cpuPipeline };
+  } catch (cpuLoadErr) {
+    logger.error('[embedder] CPU fallback pipeline load failed:', cpuLoadErr);
+    const error = new Error(`WebGPU initialization failed and CPU fallback load also failed: ${cpuLoadErr instanceof Error ? cpuLoadErr.message : String(cpuLoadErr)}`, { cause: cpuLoadErr });
+    return { success: false, error };
+  }
+}
+
+/**
+ * Handle WebGPU error during warmup
+ */
+export async function handleWebGPUWarmupError(
+  warmupErr: Error,
+  pipeline: FeatureExtractionPipeline | null,
+  stateManager: IStateManager | null,
+  gpuLockHeld: boolean,
+  model: string,
+  initializationTimeoutMs: number,
+  useCache: boolean
+): Promise<{ success: boolean; pipeline?: FeatureExtractionPipeline; error?: Error; dummy?: any }> {
+  const isValidationError = warmupErr.message.toLowerCase().includes('validation');
+  
+  markWebGpuFallback();
+  
+  if (isValidationError) {
+    logger.warn('[embedder] WebGPU validation error during warmup — falling back to CPU');
+    logger.debug('[embedder] Validation error details:', warmupErr.message);
+  } else {
+    logger.warn('[embedder] WebGPU OOM during warmup — falling back to CPU');
+  }
+  
+  // Clean up
+  if (pipeline) {
+    try { if (typeof (pipeline as any).dispose === 'function') await (pipeline as DisposablePipeline).dispose(); } catch (err) { logger.warn('[embedder] Error disposing pipeline:', err); }
+  }
+  await releaseGpuLock(stateManager, gpuLockHeld);
+
+  // Load on CPU
+  try {
+    const cpuPipeline = await loadModelOnCPU(model, initializationTimeoutMs, useCache);
+    logger.info('[embedder] CPU pipeline loaded successfully');
+    
+    // Warmup CPU pipeline
+    const { dummy, success, error: cpuWarmupErr } = await warmupPipeline(cpuPipeline, 'mean', useCache);
+    if (!success || cpuWarmupErr) {
+      logger.error('[embedder] CPU fallback pipeline warmup failed:', cpuWarmupErr);
+      if (cpuPipeline) {
+        try { if (typeof (cpuPipeline as any).dispose === 'function') await (cpuPipeline as DisposablePipeline).dispose(); } catch (err) { logger.warn('[embedder] Error disposing fallback CPU pipeline:', err); }
+      }
+      const error = new Error(`WebGPU initialization failed and CPU fallback warmup also failed: ${cpuWarmupErr?.message || 'Unknown error'}`, { cause: cpuWarmupErr });
+      return { success: false, error };
+    }
+    
+    logger.info(`[embedder] CPU pipeline warmup successful. Ready with device: cpu`);
+    return { success: true, pipeline: cpuPipeline, dummy };
+  } catch (cpuLoadErr) {
+    logger.error('[embedder] CPU pipeline load failed:', cpuLoadErr);
+    const error = new Error(`WebGPU initialization failed and CPU fallback load also failed: ${cpuLoadErr instanceof Error ? cpuLoadErr.message : String(cpuLoadErr)}`, { cause: cpuLoadErr });
+    return { success: false, error };
+  }
+}