@optave/codegraph 3.1.3 → 3.1.4

package/src/db/repository/sqlite-repository.js

@@ -0,0 +1,219 @@
+import { Repository } from './base.js';
+import { hasCfgTables } from './cfg.js';
+import { getComplexityForNode } from './complexity.js';
+import { hasDataflowTable } from './dataflow.js';
+import {
+  countCrossFileCallers,
+  findAllIncomingEdges,
+  findAllOutgoingEdges,
+  findCalleeNames,
+  findCallees,
+  findCallerNames,
+  findCallers,
+  findCrossFileCallTargets,
+  findDistinctCallers,
+  findImportDependents,
+  findImportSources,
+  findImportTargets,
+  findIntraFileCallEdges,
+  getClassHierarchy,
+} from './edges.js';
+import { hasEmbeddings } from './embeddings.js';
+import { getCallableNodes, getCallEdges, getFileNodesAll, getImportEdges } from './graph-read.js';
+import {
+  bulkNodeIdsByFile,
+  countEdges,
+  countFiles,
+  countNodes,
+  findFileNodes,
+  findNodeById,
+  findNodeByQualifiedName,
+  findNodeChildren,
+  findNodesByFile,
+  findNodesByScope,
+  findNodesForTriage,
+  findNodesWithFanIn,
+  getFunctionNodeId,
+  getNodeId,
+  iterateFunctionNodes,
+  listFunctionNodes,
+} from './nodes.js';
+
+/**
+ * SqliteRepository — wraps existing `fn(db, ...)` repository functions
+ * behind the Repository interface so callers can use `repo.method(...)`.
+ */
+export class SqliteRepository extends Repository {
+  #db;
+
+  /** @param {object} db - better-sqlite3 Database instance */
+  constructor(db) {
+    super();
+    this.#db = db;
+  }
+
+  /** Expose the underlying db for code that still needs raw access. */
+  get db() {
+    return this.#db;
+  }
+
+  // ── Node lookups ──────────────────────────────────────────────────
+
+  findNodeById(id) {
+    return findNodeById(this.#db, id);
+  }
+
+  findNodesByFile(file) {
+    return findNodesByFile(this.#db, file);
+  }
+
+  findFileNodes(fileLike) {
+    return findFileNodes(this.#db, fileLike);
+  }
+
+  findNodesWithFanIn(namePattern, opts) {
+    return findNodesWithFanIn(this.#db, namePattern, opts);
+  }
+
+  countNodes() {
+    return countNodes(this.#db);
+  }
+
+  countEdges() {
+    return countEdges(this.#db);
+  }
+
+  countFiles() {
+    return countFiles(this.#db);
+  }
+
+  getNodeId(name, kind, file, line) {
+    return getNodeId(this.#db, name, kind, file, line);
+  }
+
+  getFunctionNodeId(name, file, line) {
+    return getFunctionNodeId(this.#db, name, file, line);
+  }
+
+  bulkNodeIdsByFile(file) {
+    return bulkNodeIdsByFile(this.#db, file);
+  }
+
+  findNodeChildren(parentId) {
+    return findNodeChildren(this.#db, parentId);
+  }
+
+  findNodesByScope(scopeName, opts) {
+    return findNodesByScope(this.#db, scopeName, opts);
+  }
+
+  findNodeByQualifiedName(qualifiedName, opts) {
+    return findNodeByQualifiedName(this.#db, qualifiedName, opts);
+  }
+
+  listFunctionNodes(opts) {
+    return listFunctionNodes(this.#db, opts);
+  }
+
+  iterateFunctionNodes(opts) {
+    return iterateFunctionNodes(this.#db, opts);
+  }
+
+  findNodesForTriage(opts) {
+    return findNodesForTriage(this.#db, opts);
+  }
+
+  // ── Edge queries ──────────────────────────────────────────────────
+
+  findCallees(nodeId) {
+    return findCallees(this.#db, nodeId);
+  }
+
+  findCallers(nodeId) {
+    return findCallers(this.#db, nodeId);
+  }
+
+  findDistinctCallers(nodeId) {
+    return findDistinctCallers(this.#db, nodeId);
+  }
+
+  findAllOutgoingEdges(nodeId) {
+    return findAllOutgoingEdges(this.#db, nodeId);
+  }
+
+  findAllIncomingEdges(nodeId) {
+    return findAllIncomingEdges(this.#db, nodeId);
+  }
+
+  findCalleeNames(nodeId) {
+    return findCalleeNames(this.#db, nodeId);
+  }
+
+  findCallerNames(nodeId) {
+    return findCallerNames(this.#db, nodeId);
+  }
+
+  findImportTargets(nodeId) {
+    return findImportTargets(this.#db, nodeId);
+  }
+
+  findImportSources(nodeId) {
+    return findImportSources(this.#db, nodeId);
+  }
+
+  findImportDependents(nodeId) {
+    return findImportDependents(this.#db, nodeId);
+  }
+
+  findCrossFileCallTargets(file) {
+    return findCrossFileCallTargets(this.#db, file);
+  }
+
+  countCrossFileCallers(nodeId, file) {
+    return countCrossFileCallers(this.#db, nodeId, file);
+  }
+
+  getClassHierarchy(classNodeId) {
+    return getClassHierarchy(this.#db, classNodeId);
+  }
+
+  findIntraFileCallEdges(file) {
+    return findIntraFileCallEdges(this.#db, file);
+  }
+
+  // ── Graph-read queries ────────────────────────────────────────────
+
+  getCallableNodes() {
+    return getCallableNodes(this.#db);
+  }
+
+  getCallEdges() {
+    return getCallEdges(this.#db);
+  }
+
+  getFileNodesAll() {
+    return getFileNodesAll(this.#db);
+  }
+
+  getImportEdges() {
+    return getImportEdges(this.#db);
+  }
+
+  // ── Optional table checks ─────────────────────────────────────────
+
+  hasCfgTables() {
+    return hasCfgTables(this.#db);
+  }
+
+  hasEmbeddings() {
+    return hasEmbeddings(this.#db);
+  }
+
+  hasDataflowTable() {
+    return hasDataflowTable(this.#db);
+  }
+
+  getComplexityForNode(nodeId) {
+    return getComplexityForNode(this.#db, nodeId);
+  }
+}

package/src/db.js

@@ -29,8 +29,10 @@ export {
   findImportTargets,
   findIntraFileCallEdges,
   findNodeById,
+  findNodeByQualifiedName,
   findNodeChildren,
   findNodesByFile,
+  findNodesByScope,
   findNodesForTriage,
   findNodesWithFanIn,
   getCallableNodes,
@@ -50,9 +52,12 @@ export {
   hasCoChanges,
   hasDataflowTable,
   hasEmbeddings,
+  InMemoryRepository,
   iterateFunctionNodes,
   listFunctionNodes,
   purgeFileData,
   purgeFilesData,
+  Repository,
+  SqliteRepository,
   upsertCoChangeMeta,
 } from './db/repository/index.js';

package/src/embeddings/generator.js

@@ -0,0 +1,163 @@
+import fs from 'node:fs';
+import path from 'node:path';
+import { closeDb, findDbPath, openDb } from '../db.js';
+import { DbError } from '../errors.js';
+import { warn } from '../logger.js';
+import { embed, getModelConfig } from './models.js';
+import { buildSourceText } from './strategies/source.js';
+import { buildStructuredText } from './strategies/structured.js';
+
+/**
+ * Rough token estimate (~4 chars per token for code/English).
+ * Conservative — avoids adding a tokenizer dependency.
+ */
+export function estimateTokens(text) {
+  return Math.ceil(text.length / 4);
+}
+
+function initEmbeddingsSchema(db) {
+  db.exec(`
+    CREATE TABLE IF NOT EXISTS embeddings (
+      node_id INTEGER PRIMARY KEY,
+      vector BLOB NOT NULL,
+      text_preview TEXT,
+      FOREIGN KEY(node_id) REFERENCES nodes(id)
+    );
+    CREATE TABLE IF NOT EXISTS embedding_meta (
+      key TEXT PRIMARY KEY,
+      value TEXT
+    );
+  `);
+
+  // Add full_text column (idempotent — ignore if already exists)
+  try {
+    db.exec('ALTER TABLE embeddings ADD COLUMN full_text TEXT');
+  } catch {
+    /* column already exists */
+  }
+
+  // FTS5 virtual table for BM25 keyword search
+  db.exec(`
+    CREATE VIRTUAL TABLE IF NOT EXISTS fts_index USING fts5(
+      name,
+      content,
+      tokenize='unicode61'
+    );
+  `);
+}
+
+/**
+ * Build embeddings for all functions/methods/classes in the graph.
+ * @param {string} rootDir - Project root directory
+ * @param {string} modelKey - Model identifier from MODELS registry
+ * @param {string} [customDbPath] - Override path to graph.db
+ * @param {object} [options] - Embedding options
+ * @param {string} [options.strategy='structured'] - 'structured' (graph-enriched) or 'source' (raw code)
+ */
+export async function buildEmbeddings(rootDir, modelKey, customDbPath, options = {}) {
+  const strategy = options.strategy || 'structured';
+  const dbPath = customDbPath || findDbPath(null);
+
+  if (!fs.existsSync(dbPath)) {
+    throw new DbError(
+      `No codegraph database found at ${dbPath}.\nRun "codegraph build" first to analyze your codebase.`,
+      { file: dbPath },
+    );
+  }
+
+  const db = openDb(dbPath);
+  initEmbeddingsSchema(db);
+
+  db.exec('DELETE FROM embeddings');
+  db.exec('DELETE FROM embedding_meta');
+  db.exec('DELETE FROM fts_index');
+
+  const nodes = db
+    .prepare(
+      `SELECT * FROM nodes WHERE kind IN ('function', 'method', 'class') ORDER BY file, line`,
+    )
+    .all();
+
+  console.log(`Building embeddings for ${nodes.length} symbols (strategy: ${strategy})...`);
+
+  const byFile = new Map();
+  for (const node of nodes) {
+    if (!byFile.has(node.file)) byFile.set(node.file, []);
+    byFile.get(node.file).push(node);
+  }
+
+  const texts = [];
+  const nodeIds = [];
+  const nodeNames = [];
+  const previews = [];
+  const config = getModelConfig(modelKey);
+  const contextWindow = config.contextWindow;
+  let overflowCount = 0;
+
+  for (const [file, fileNodes] of byFile) {
+    const fullPath = path.join(rootDir, file);
+    let lines;
+    try {
+      lines = fs.readFileSync(fullPath, 'utf-8').split('\n');
+    } catch (err) {
+      warn(`Cannot read ${file} for embeddings: ${err.message}`);
+      continue;
+    }
+
+    for (const node of fileNodes) {
+      let text =
+        strategy === 'structured'
+          ? buildStructuredText(node, file, lines, db)
+          : buildSourceText(node, file, lines);
+
+      // Detect and handle context window overflow
+      const tokens = estimateTokens(text);
+      if (tokens > contextWindow) {
+        overflowCount++;
+        const maxChars = contextWindow * 4;
+        text = text.slice(0, maxChars);
+      }
+
+      texts.push(text);
+      nodeIds.push(node.id);
+      nodeNames.push(node.name);
+      previews.push(`${node.name} (${node.kind}) -- ${file}:${node.line}`);
+    }
+  }
+
+  if (overflowCount > 0) {
+    warn(
+      `${overflowCount} symbol(s) exceeded model context window (${contextWindow} tokens) and were truncated`,
+    );
+  }
+
+  console.log(`Embedding ${texts.length} symbols...`);
+  const { vectors, dim } = await embed(texts, modelKey);
+
+  const insert = db.prepare(
+    'INSERT OR REPLACE INTO embeddings (node_id, vector, text_preview, full_text) VALUES (?, ?, ?, ?)',
+  );
+  const insertFts = db.prepare('INSERT INTO fts_index(rowid, name, content) VALUES (?, ?, ?)');
+  const insertMeta = db.prepare('INSERT OR REPLACE INTO embedding_meta (key, value) VALUES (?, ?)');
+  const insertAll = db.transaction(() => {
+    for (let i = 0; i < vectors.length; i++) {
+      insert.run(nodeIds[i], Buffer.from(vectors[i].buffer), previews[i], texts[i]);
+      insertFts.run(nodeIds[i], nodeNames[i], texts[i]);
+    }
+    insertMeta.run('model', config.name);
+    insertMeta.run('dim', String(dim));
+    insertMeta.run('count', String(vectors.length));
+    insertMeta.run('fts_count', String(vectors.length));
+    insertMeta.run('strategy', strategy);
+    insertMeta.run('built_at', new Date().toISOString());
+    if (overflowCount > 0) {
+      insertMeta.run('truncated_count', String(overflowCount));
+    }
+  });
+  insertAll();
+
+  console.log(
+    `\nStored ${vectors.length} embeddings (${dim}d, ${config.name}, strategy: ${strategy}) in graph.db`,
+  );
+  closeDb(db);
+}

package/src/embeddings/index.js

@@ -0,0 +1,13 @@
+/**
+ * Embeddings subsystem — public API barrel.
+ *
+ * Re-exports everything consumers previously imported from `../embedder.js`.
+ */
+
+export { buildEmbeddings, estimateTokens } from './generator.js';
+export { DEFAULT_MODEL, disposeModel, EMBEDDING_STRATEGIES, embed, MODELS } from './models.js';
+export { search } from './search/cli-formatter.js';
+export { hybridSearchData } from './search/hybrid.js';
+export { ftsSearchData } from './search/keyword.js';
+export { multiSearchData, searchData } from './search/semantic.js';
+export { cosineSim } from './stores/sqlite-blob.js';

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 };
+}