@optave/codegraph 3.1.2 → 3.1.4

package/src/embeddings/models.js

@@ -0,0 +1,218 @@
+import { execFileSync } from 'node:child_process';
+import { createInterface } from 'node:readline';
+import { ConfigError, EngineError } from '../errors.js';
+import { info } from '../logger.js';
+
+// Lazy-load transformers (heavy, optional module)
+let pipeline = null;
+let extractor = null;
+let activeModel = null;
+
+export const MODELS = {
+  minilm: {
+    name: 'Xenova/all-MiniLM-L6-v2',
+    dim: 384,
+    contextWindow: 256,
+    desc: 'Smallest, fastest (~23MB). General text.',
+    quantized: true,
+  },
+  'jina-small': {
+    name: 'Xenova/jina-embeddings-v2-small-en',
+    dim: 512,
+    contextWindow: 8192,
+    desc: 'Small, good quality (~33MB). General text.',
+    quantized: false,
+  },
+  'jina-base': {
+    name: 'Xenova/jina-embeddings-v2-base-en',
+    dim: 768,
+    contextWindow: 8192,
+    desc: 'Good quality (~137MB). General text, 8192 token context.',
+    quantized: false,
+  },
+  'jina-code': {
+    name: 'Xenova/jina-embeddings-v2-base-code',
+    dim: 768,
+    contextWindow: 8192,
+    desc: 'Code-aware (~137MB). Trained on code+text, best for code search.',
+    quantized: false,
+  },
+  nomic: {
+    name: 'Xenova/nomic-embed-text-v1',
+    dim: 768,
+    contextWindow: 8192,
+    desc: 'Good local quality (~137MB). 8192 context.',
+    quantized: false,
+  },
+  'nomic-v1.5': {
+    name: 'nomic-ai/nomic-embed-text-v1.5',
+    dim: 768,
+    contextWindow: 8192,
+    desc: 'Improved nomic (~137MB). Matryoshka dimensions, 8192 context.',
+    quantized: false,
+  },
+  'bge-large': {
+    name: 'Xenova/bge-large-en-v1.5',
+    dim: 1024,
+    contextWindow: 512,
+    desc: 'Best general retrieval (~335MB). Top MTEB scores.',
+    quantized: false,
+  },
+};
+
+export const EMBEDDING_STRATEGIES = ['structured', 'source'];
+
+export const DEFAULT_MODEL = 'nomic-v1.5';
+const BATCH_SIZE_MAP = {
+  minilm: 32,
+  'jina-small': 16,
+  'jina-base': 8,
+  'jina-code': 8,
+  nomic: 8,
+  'nomic-v1.5': 8,
+  'bge-large': 4,
+};
+const DEFAULT_BATCH_SIZE = 32;
+
+/** @internal Used by generator.js — not part of the public barrel. */
+export function getModelConfig(modelKey) {
+  const key = modelKey || DEFAULT_MODEL;
+  const config = MODELS[key];
+  if (!config) {
+    throw new ConfigError(`Unknown model: ${key}. Available: ${Object.keys(MODELS).join(', ')}`);
+  }
+  return config;
+}
+
+/**
+ * Prompt the user to install a missing package interactively.
+ * Returns true if the package was installed, false otherwise.
+ * Skips the prompt entirely in non-TTY environments (CI, piped stdin).
+ * @internal Not part of the public barrel.
+ */
+export function promptInstall(packageName) {
+  if (!process.stdin.isTTY) return Promise.resolve(false);
+
+  return new Promise((resolve) => {
+    const rl = createInterface({ input: process.stdin, output: process.stderr });
+    rl.question(`Semantic search requires ${packageName}. Install it now? [y/N] `, (answer) => {
+      rl.close();
+      if (answer.trim().toLowerCase() !== 'y') return resolve(false);
+      try {
+        execFileSync('npm', ['install', packageName], {
+          stdio: 'inherit',
+          timeout: 300_000,
+        });
+        resolve(true);
+      } catch {
+        resolve(false);
+      }
+    });
+  });
+}
+
+/**
+ * Lazy-load @huggingface/transformers.
+ * If the package is missing, prompts the user to install it interactively.
+ * In non-TTY environments, prints an error and exits.
+ * @internal Not part of the public barrel.
+ */
+export async function loadTransformers() {
+  try {
+    return await import('@huggingface/transformers');
+  } catch {
+    const pkg = '@huggingface/transformers';
+    const installed = await promptInstall(pkg);
+    if (installed) {
+      try {
+        return await import(pkg);
+      } catch (loadErr) {
+        throw new EngineError(
+          `${pkg} was installed but failed to load. Please check your environment.`,
+          { cause: loadErr },
+        );
+      }
+    }
+    throw new EngineError(`Semantic search requires ${pkg}.\nInstall it with: npm install ${pkg}`);
+  }
+}
+
+/**
+ * Dispose the current ONNX session and free memory.
+ * Safe to call when no model is loaded (no-op).
+ */
+export async function disposeModel() {
+  if (extractor) {
+    await extractor.dispose();
+    extractor = null;
+  }
+  activeModel = null;
+}
+
+async function loadModel(modelKey) {
+  const config = getModelConfig(modelKey);
+
+  if (extractor && activeModel === config.name) return { extractor, config };
+
+  // Dispose previous model before loading a different one
+  await disposeModel();
+
+  const transformers = await loadTransformers();
+  pipeline = transformers.pipeline;
+
+  info(`Loading embedding model: ${config.name} (${config.dim}d)...`);
+  const pipelineOpts = config.quantized ? { quantized: true } : {};
+  try {
+    extractor = await pipeline('feature-extraction', config.name, pipelineOpts);
+  } catch (err) {
+    const msg = err.message || String(err);
+    if (msg.includes('Unauthorized') || msg.includes('401') || msg.includes('gated')) {
+      throw new EngineError(
+        `Model "${config.name}" requires authentication.\n` +
+          `This model is gated on HuggingFace and needs an access token.\n\n` +
+          `Options:\n` +
+          `  1. Set HF_TOKEN env var: export HF_TOKEN=hf_...\n` +
+          `  2. Use a public model instead: codegraph embed --model minilm`,
+        { cause: err },
+      );
+    }
+    throw new EngineError(
+      `Failed to load model "${config.name}": ${msg}\n` +
+        `Try a different model: codegraph embed --model minilm`,
+      { cause: err },
+    );
+  }
+  activeModel = config.name;
+  info('Model loaded.');
+  return { extractor, config };
+}
+
+/**
+ * Generate embeddings for an array of texts.
+ */
+export async function embed(texts, modelKey) {
+  const { extractor: ext, config } = await loadModel(modelKey);
+  const dim = config.dim;
+  const results = [];
+  const batchSize = BATCH_SIZE_MAP[modelKey || DEFAULT_MODEL] || DEFAULT_BATCH_SIZE;
+
+  for (let i = 0; i < texts.length; i += batchSize) {
+    const batch = texts.slice(i, i + batchSize);
+    const output = await ext(batch, { pooling: 'mean', normalize: true });
+
+    for (let j = 0; j < batch.length; j++) {
+      const start = j * dim;
+      const vec = new Float32Array(dim);
+      for (let k = 0; k < dim; k++) {
+        vec[k] = output.data[start + k];
+      }
+      results.push(vec);
+    }
+
+    if (texts.length > batchSize) {
+      process.stdout.write(`  Embedded ${Math.min(i + batchSize, texts.length)}/${texts.length}\r`);
+    }
+  }
+
+  return { vectors: results, dim };
+}

package/src/embeddings/search/cli-formatter.js

@@ -0,0 +1,151 @@
+import { warn } from '../../logger.js';
+import { hybridSearchData } from './hybrid.js';
+import { ftsSearchData } from './keyword.js';
+import { multiSearchData, searchData } from './semantic.js';
+
+/**
+ * Search with mode support — CLI wrapper with multi-query detection.
+ * Modes: 'hybrid' (default), 'semantic', 'keyword'
+ */
+export async function search(query, customDbPath, opts = {}) {
+  const mode = opts.mode || 'hybrid';
+
+  // Split by semicolons, trim, filter empties
+  const queries = query
+    .split(';')
+    .map((q) => q.trim())
+    .filter((q) => q.length > 0);
+
+  const kindIcon = (kind) => (kind === 'function' ? 'f' : kind === 'class' ? '*' : 'o');
+
+  // ─── Keyword-only mode ──────────────────────────────────────────────
+  if (mode === 'keyword') {
+    const singleQuery = queries.length === 1 ? queries[0] : query;
+    const data = ftsSearchData(singleQuery, customDbPath, opts);
+    if (!data) {
+      console.log('No FTS5 index found. Run `codegraph embed` to build the keyword index.');
+      return;
+    }
+
+    if (opts.json) {
+      console.log(JSON.stringify(data, null, 2));
+      return;
+    }
+
+    console.log(`\nKeyword search: "${singleQuery}" (BM25)\n`);
+    if (data.results.length === 0) {
+      console.log('  No results found.');
+    } else {
+      for (const r of data.results) {
+        console.log(
+          `  BM25 ${r.bm25Score.toFixed(2)}  ${kindIcon(r.kind)} ${r.name} -- ${r.file}:${r.line}`,
+        );
+      }
+    }
+    console.log(`\n  ${data.results.length} results shown\n`);
+    return;
+  }
+
+  // ─── Semantic-only mode ─────────────────────────────────────────────
+  if (mode === 'semantic') {
+    if (queries.length <= 1) {
+      const singleQuery = queries[0] || query;
+      const data = await searchData(singleQuery, customDbPath, opts);
+      if (!data) return;
+
+      if (opts.json) {
+        console.log(JSON.stringify(data, null, 2));
+        return;
+      }
+
+      console.log(`\nSemantic search: "${singleQuery}"\n`);
+      if (data.results.length === 0) {
+        console.log('  No results above threshold.');
+      } else {
+        for (const r of data.results) {
+          const bar = '#'.repeat(Math.round(r.similarity * 20));
+          console.log(`  ${(r.similarity * 100).toFixed(1)}% ${bar}`);
+          console.log(`    ${kindIcon(r.kind)} ${r.name} -- ${r.file}:${r.line}`);
+        }
+      }
+      console.log(`\n  ${data.results.length} results shown\n`);
+    } else {
+      const data = await multiSearchData(queries, customDbPath, opts);
+      if (!data) return;
+
+      if (opts.json) {
+        console.log(JSON.stringify(data, null, 2));
+        return;
+      }
+
+      console.log(`\nMulti-query semantic search (RRF, k=${opts.rrfK || 60}):`);
+      for (let i = 0; i < queries.length; i++) console.log(`  [${i + 1}] "${queries[i]}"`);
+      console.log();
+      if (data.results.length === 0) {
+        console.log('  No results above threshold.');
+      } else {
+        for (const r of data.results) {
+          console.log(
+            `  RRF ${r.rrf.toFixed(4)}  ${kindIcon(r.kind)} ${r.name} -- ${r.file}:${r.line}`,
+          );
+          for (const qs of r.queryScores) {
+            const bar = '#'.repeat(Math.round(qs.similarity * 20));
+            console.log(
+              `    [${queries.indexOf(qs.query) + 1}] ${(qs.similarity * 100).toFixed(1)}% ${bar} (rank ${qs.rank})`,
+            );
+          }
+        }
+      }
+      console.log(`\n  ${data.results.length} results shown\n`);
+    }
+    return;
+  }
+
+  // ─── Hybrid mode (default) ──────────────────────────────────────────
+  const data = await hybridSearchData(query, customDbPath, opts);
+
+  if (!data) {
+    // No FTS5 index — fall back to semantic-only
+    warn(
+      'FTS5 index not found — using semantic search only. Re-run `codegraph embed` to enable hybrid mode.',
+    );
+    return search(query, customDbPath, { ...opts, mode: 'semantic' });
+  }
+
+  if (opts.json) {
+    console.log(JSON.stringify(data, null, 2));
+    return;
+  }
+
+  const rrfK = opts.rrfK || 60;
+  if (queries.length <= 1) {
+    const singleQuery = queries[0] || query;
+    console.log(`\nHybrid search: "${singleQuery}" (BM25 + semantic, RRF k=${rrfK})\n`);
+  } else {
+    console.log(`\nHybrid multi-query search (BM25 + semantic, RRF k=${rrfK}):`);
+    for (let i = 0; i < queries.length; i++) console.log(`  [${i + 1}] "${queries[i]}"`);
+    console.log();
+  }
+
+  if (data.results.length === 0) {
+    console.log('  No results found.');
+  } else {
+    for (const r of data.results) {
+      console.log(
+        `  RRF ${r.rrf.toFixed(4)}  ${kindIcon(r.kind)} ${r.name} -- ${r.file}:${r.line}`,
+      );
+      const parts = [];
+      if (r.bm25Rank != null) {
+        parts.push(`BM25: rank ${r.bm25Rank} (score ${r.bm25Score.toFixed(2)})`);
+      }
+      if (r.semanticRank != null) {
+        parts.push(`Semantic: rank ${r.semanticRank} (${(r.similarity * 100).toFixed(1)}%)`);
+      }
+      if (parts.length > 0) {
+        console.log(`    ${parts.join('  |  ')}`);
+      }
+    }
+  }
+
+  console.log(`\n  ${data.results.length} results shown\n`);
+}

package/src/embeddings/search/filters.js

@@ -0,0 +1,46 @@
+/**
+ * Match a file path against a glob pattern.
+ * Supports *, **, and ? wildcards. Zero dependencies.
+ */
+export function globMatch(filePath, pattern) {
+  // Normalize separators to forward slashes
+  const normalized = filePath.replace(/\\/g, '/');
+  // Escape regex specials except glob chars
+  let regex = pattern.replace(/\\/g, '/').replace(/[.+^${}()|[\]\\]/g, '\\$&');
+  // Replace ** first (matches any path segment), then * and ?
+  regex = regex.replace(/\*\*/g, '\0');
+  regex = regex.replace(/\*/g, '[^/]*');
+  regex = regex.replace(/\0/g, '.*');
+  regex = regex.replace(/\?/g, '[^/]');
+  try {
+    return new RegExp(`^${regex}$`).test(normalized);
+  } catch {
+    // Malformed pattern — fall back to substring match
+    return normalized.includes(pattern);
+  }
+}
+
+const TEST_PATTERN = /\.(test|spec)\.|__test__|__tests__|\.stories\./;
+
+/**
+ * Apply post-query filters (glob pattern, noTests) to a set of rows.
+ * Mutates nothing — returns a new filtered array.
+ * @param {Array} rows - Rows with at least a `file` property
+ * @param {object} opts
+ * @param {string} [opts.filePattern] - Glob pattern (only applied if it contains glob chars)
+ * @param {boolean} [opts.noTests] - Exclude test/spec files
+ * @param {boolean} [opts.isGlob] - Pre-computed: does filePattern contain glob chars?
+ * @returns {Array}
+ */
+export function applyFilters(rows, opts = {}) {
+  let filtered = rows;
+  const isGlob =
+    opts.isGlob !== undefined ? opts.isGlob : opts.filePattern && /[*?[\]]/.test(opts.filePattern);
+  if (isGlob) {
+    filtered = filtered.filter((row) => globMatch(row.file, opts.filePattern));
+  }
+  if (opts.noTests) {
+    filtered = filtered.filter((row) => !TEST_PATTERN.test(row.file));
+  }
+  return filtered;
+}

package/src/embeddings/search/hybrid.js

@@ -0,0 +1,121 @@
+import { openReadonlyOrFail } from '../../db.js';
+import { hasFtsIndex } from '../stores/fts5.js';
+import { ftsSearchData } from './keyword.js';
+import { searchData } from './semantic.js';
+
+/**
+ * Hybrid BM25 + semantic search with RRF fusion.
+ * Returns { results: [{ name, kind, file, line, rrf, bm25Score, bm25Rank, similarity, semanticRank }] }
+ * or null if no FTS5 index (caller should fall back to semantic-only).
+ */
+export async function hybridSearchData(query, customDbPath, opts = {}) {
+  const limit = opts.limit || 15;
+  const k = opts.rrfK || 60;
+  const topK = (opts.limit || 15) * 5;
+
+  // Split semicolons for multi-query support
+  const queries =
+    typeof query === 'string'
+      ? query
+          .split(';')
+          .map((q) => q.trim())
+          .filter((q) => q.length > 0)
+      : [query];
+
+  // Check FTS5 availability first (sync, cheap)
+  const checkDb = openReadonlyOrFail(customDbPath);
+  const ftsAvailable = hasFtsIndex(checkDb);
+  checkDb.close();
+  if (!ftsAvailable) return null;
+
+  // Collect ranked lists: for each query, one BM25 list + one semantic list
+  const rankedLists = [];
+
+  for (const q of queries) {
+    // BM25 ranked list (sync)
+    const bm25Data = ftsSearchData(q, customDbPath, { ...opts, limit: topK });
+    if (bm25Data?.results) {
+      rankedLists.push(
+        bm25Data.results.map((r, idx) => ({
+          key: `${r.name}:${r.file}:${r.line}`,
+          rank: idx + 1,
+          source: 'bm25',
+          ...r,
+        })),
+      );
+    }
+
+    // Semantic ranked list (async)
+    const semData = await searchData(q, customDbPath, {
+      ...opts,
+      limit: topK,
+      minScore: opts.minScore || 0.2,
+    });
+    if (semData?.results) {
+      rankedLists.push(
+        semData.results.map((r, idx) => ({
+          key: `${r.name}:${r.file}:${r.line}`,
+          rank: idx + 1,
+          source: 'semantic',
+          ...r,
+        })),
+      );
+    }
+  }
+
+  // RRF fusion across all ranked lists
+  const fusionMap = new Map();
+  for (const list of rankedLists) {
+    for (const item of list) {
+      if (!fusionMap.has(item.key)) {
+        fusionMap.set(item.key, {
+          name: item.name,
+          kind: item.kind,
+          file: item.file,
+          line: item.line,
+          endLine: item.endLine ?? null,
+          role: item.role ?? null,
+          fileHash: item.fileHash ?? null,
+          rrfScore: 0,
+          bm25Score: null,
+          bm25Rank: null,
+          similarity: null,
+          semanticRank: null,
+        });
+      }
+      const entry = fusionMap.get(item.key);
+      entry.rrfScore += 1 / (k + item.rank);
+      if (item.source === 'bm25') {
+        if (entry.bm25Rank === null || item.rank < entry.bm25Rank) {
+          entry.bm25Score = item.bm25Score;
+          entry.bm25Rank = item.rank;
+        }
+      } else {
+        if (entry.semanticRank === null || item.rank < entry.semanticRank) {
+          entry.similarity = item.similarity;
+          entry.semanticRank = item.rank;
+        }
+      }
+    }
+  }
+
+  const results = [...fusionMap.values()]
+    .sort((a, b) => b.rrfScore - a.rrfScore)
+    .slice(0, limit)
+    .map((e) => ({
+      name: e.name,
+      kind: e.kind,
+      file: e.file,
+      line: e.line,
+      endLine: e.endLine,
+      role: e.role,
+      fileHash: e.fileHash,
+      rrf: e.rrfScore,
+      bm25Score: e.bm25Score,
+      bm25Rank: e.bm25Rank,
+      similarity: e.similarity,
+      semanticRank: e.semanticRank,
+    }));
+
+  return { results };
+}

package/src/embeddings/search/keyword.js

@@ -0,0 +1,68 @@
+import { openReadonlyOrFail } from '../../db.js';
+import { normalizeSymbol } from '../../queries.js';
+import { hasFtsIndex, sanitizeFtsQuery } from '../stores/fts5.js';
+import { applyFilters } from './filters.js';
+
+/**
+ * BM25 keyword search via FTS5.
+ * Returns { results: [{ name, kind, file, line, bm25Score }] } or null if no FTS5 index.
+ */
+export function ftsSearchData(query, customDbPath, opts = {}) {
+  const limit = opts.limit || 15;
+
+  const db = openReadonlyOrFail(customDbPath);
+
+  try {
+    if (!hasFtsIndex(db)) {
+      return null;
+    }
+
+    const ftsQuery = sanitizeFtsQuery(query);
+    if (!ftsQuery) {
+      return { results: [] };
+    }
+
+    let sql = `
+      SELECT f.rowid AS node_id, rank AS bm25_score,
+             n.name, n.kind, n.file, n.line, n.end_line, n.role
+      FROM fts_index f
+      JOIN nodes n ON f.rowid = n.id
+      WHERE fts_index MATCH ?
+    `;
+    const params = [ftsQuery];
+
+    if (opts.kind) {
+      sql += ' AND n.kind = ?';
+      params.push(opts.kind);
+    }
+
+    const isGlob = opts.filePattern && /[*?[\]]/.test(opts.filePattern);
+    if (opts.filePattern && !isGlob) {
+      sql += ' AND n.file LIKE ?';
+      params.push(`%${opts.filePattern}%`);
+    }
+
+    sql += ' ORDER BY rank LIMIT ?';
+    params.push(limit * 5); // fetch generous set for post-filtering
+
+    let rows;
+    try {
+      rows = db.prepare(sql).all(...params);
+    } catch {
+      // Invalid FTS5 query syntax — return empty
+      return { results: [] };
+    }
+
+    rows = applyFilters(rows, { ...opts, isGlob });
+
+    const hc = new Map();
+    const results = rows.slice(0, limit).map((row) => ({
+      ...normalizeSymbol(row, db, hc),
+      bm25Score: -row.bm25_score, // FTS5 rank is negative; negate for display
+    }));
+
+    return { results };
+  } finally {
+    db.close();
+  }
+}

package/src/embeddings/search/prepare.js

@@ -0,0 +1,66 @@
+import { getEmbeddingCount, getEmbeddingMeta } from '../../db/repository/embeddings.js';
+import { openReadonlyOrFail } from '../../db.js';
+import { MODELS } from '../models.js';
+import { applyFilters } from './filters.js';
+
+/**
+ * Shared setup for search functions: opens DB, validates embeddings/model, loads rows.
+ * Returns { db, rows, modelKey, storedDim } or null on failure (prints error).
+ * On null return, the DB is closed. On exception, the DB is also closed
+ * (callers only need to close DB from the returned object on the happy path).
+ */
+export function prepareSearch(customDbPath, opts = {}) {
+  const db = openReadonlyOrFail(customDbPath);
+
+  try {
+    const count = getEmbeddingCount(db);
+    if (count === 0) {
+      console.log('No embeddings found. Run `codegraph embed` first.');
+      db.close();
+      return null;
+    }
+
+    const storedModel = getEmbeddingMeta(db, 'model') || null;
+    const dimStr = getEmbeddingMeta(db, 'dim');
+    const storedDim = dimStr ? parseInt(dimStr, 10) : null;
+
+    let modelKey = opts.model || null;
+    if (!modelKey && storedModel) {
+      for (const [key, config] of Object.entries(MODELS)) {
+        if (config.name === storedModel) {
+          modelKey = key;
+          break;
+        }
+      }
+    }
+
+    // Pre-filter: allow filtering by kind or file pattern to reduce search space
+    const isGlob = opts.filePattern && /[*?[\]]/.test(opts.filePattern);
+    let sql = `
+    SELECT e.node_id, e.vector, e.text_preview, n.name, n.kind, n.file, n.line, n.end_line, n.role
+    FROM embeddings e
+    JOIN nodes n ON e.node_id = n.id
+  `;
+    const params = [];
+    const conditions = [];
+    if (opts.kind) {
+      conditions.push('n.kind = ?');
+      params.push(opts.kind);
+    }
+    if (opts.filePattern && !isGlob) {
+      conditions.push('n.file LIKE ?');
+      params.push(`%${opts.filePattern}%`);
+    }
+    if (conditions.length > 0) {
+      sql += ` WHERE ${conditions.join(' AND ')}`;
+    }
+
+    let rows = db.prepare(sql).all(...params);
+    rows = applyFilters(rows, { ...opts, isGlob });
+
+    return { db, rows, modelKey, storedDim };
+  } catch (err) {
+    db.close();
+    throw err;
+  }
+}