codebase-contextualizer-cli 1.0.0

package/src/indexer.js

@@ -0,0 +1,290 @@
+const fs = require("fs/promises");
+const path = require("path");
+const { createAbortError, throwIfAborted } = require("./abort");
+const { loadCache, saveCache } = require("./cache");
+const { mapLimit } = require("./concurrency");
+const { databaseExists, writeEmbeddingResults } = require("./database");
+const { hashFile } = require("./file-hash");
+const { WorkerPool, getDefaultWorkerCount } = require("./worker-pool");
+const { walkSourceFiles } = require("./walker");
+
+async function resolveTargetRoot(target) {
+  const root = path.resolve(target);
+  const stat = await fs.stat(root);
+
+  if (!stat.isDirectory()) {
+    throw new Error(`Target is not a directory: ${root}`);
+  }
+
+  return root;
+}
+
+function createCounts(changes) {
+  return {
+    scanned: changes.new.length + changes.modified.length + changes.unchanged.length,
+    new: changes.new.length,
+    modified: changes.modified.length,
+    unchanged: changes.unchanged.length,
+    removed: changes.removed.length,
+  };
+}
+
+async function collectState(target, options = {}) {
+  throwIfAborted(options.signal);
+
+  const root = await resolveTargetRoot(target);
+  throwIfAborted(options.signal);
+
+  const cacheState = await loadCache(root);
+  throwIfAborted(options.signal);
+
+  const walkResult = await walkSourceFiles(root, {
+    signal: options.signal,
+  });
+  throwIfAborted(options.signal);
+
+  const hashConcurrency = options.hashConcurrency || 16;
+  const hashResults = await mapLimit(walkResult.files, hashConcurrency, async (file) => {
+    throwIfAborted(options.signal);
+
+    try {
+      const hash = await hashFile(file.absolutePath);
+      throwIfAborted(options.signal);
+
+      return {
+        file: {
+          ...file,
+          hash,
+        },
+      };
+    } catch (error) {
+      if (options.signal && options.signal.aborted) {
+        throw createAbortError(options.signal);
+      }
+
+      return {
+        error: {
+          path: file.absolutePath,
+          message: error.message,
+        },
+      };
+    }
+  });
+  const hashedFiles = hashResults.filter((result) => result.file).map((result) => result.file);
+  const errors = [
+    ...walkResult.errors,
+    ...hashResults.filter((result) => result.error).map((result) => result.error),
+  ];
+
+  const currentFiles = {};
+  const filesByRelativePath = {};
+  const changes = {
+    new: [],
+    modified: [],
+    unchanged: [],
+    removed: [],
+  };
+
+  for (const file of hashedFiles) {
+    const previous = cacheState.cache.files[file.relativePath];
+    const status = !previous
+      ? "new"
+      : previous.hash === file.hash
+        ? "unchanged"
+        : "modified";
+
+    currentFiles[file.relativePath] = {
+      hash: file.hash,
+      size: file.size,
+      mtimeMs: file.mtimeMs,
+      indexedAt: previous && status === "unchanged" ? previous.indexedAt || null : null,
+    };
+    filesByRelativePath[file.relativePath] = file;
+    changes[status].push(file.relativePath);
+  }
+
+  for (const cachedPath of Object.keys(cacheState.cache.files)) {
+    if (!currentFiles[cachedPath]) {
+      changes.removed.push(cachedPath);
+    }
+  }
+
+  for (const key of Object.keys(changes)) {
+    changes[key].sort();
+  }
+
+  return {
+    root,
+    cachePath: cacheState.cachePath,
+    cacheExists: cacheState.exists,
+    previousCache: cacheState.cache,
+    currentFiles,
+    filesByRelativePath,
+    changes,
+    counts: createCounts(changes),
+    errors,
+  };
+}
+
+function createEmbeddingSummary(results) {
+  const succeeded = results.filter((result) => result.ok);
+  const failed = results.filter((result) => !result.ok);
+  const skipped = succeeded.filter((result) => result.skipped);
+  const embedded = succeeded.filter((result) => !result.skipped);
+  const chunkCount = embedded.reduce((total, result) => total + result.chunks.length, 0);
+
+  return {
+    files: results,
+    counts: {
+      queued: results.length,
+      succeeded: succeeded.length,
+      failed: failed.length,
+      skipped: skipped.length,
+      chunks: chunkCount,
+    },
+  };
+}
+
+async function embedFiles(result, pathsToEmbed, options = {}) {
+  const relativePaths = [...new Set(pathsToEmbed)].sort();
+
+  if (relativePaths.length === 0) {
+    return createEmbeddingSummary([]);
+  }
+
+  const workerCount = options.workerCount || getDefaultWorkerCount();
+  const pool = new WorkerPool({
+    size: workerCount,
+  });
+  let removeAbortListener = null;
+
+  if (options.signal) {
+    const onAbort = () => {
+      void pool.close();
+    };
+
+    if (options.signal.aborted) {
+      await pool.close();
+      throw createAbortError(options.signal);
+    }
+
+    options.signal.addEventListener("abort", onAbort, { once: true });
+    removeAbortListener = () => options.signal.removeEventListener("abort", onAbort);
+  }
+
+  try {
+    const embeddingResults = await Promise.all(relativePaths.map(async (relativePath) => {
+      throwIfAborted(options.signal);
+
+      const file = result.filesByRelativePath[relativePath];
+
+      if (!file) {
+        return null;
+      }
+
+      try {
+        const workerResult = await pool.run({
+          filePath: file.absolutePath,
+          relativePath,
+        });
+        throwIfAborted(options.signal);
+
+        return {
+          ok: true,
+          ...workerResult,
+        };
+      } catch (error) {
+        if (options.signal && options.signal.aborted) {
+          throw createAbortError(options.signal);
+        }
+
+        return {
+          ok: false,
+          filePath: file.absolutePath,
+          relativePath,
+          error: error.message,
+          chunks: [],
+        };
+      }
+    }));
+
+    return createEmbeddingSummary(embeddingResults.filter(Boolean));
+  } finally {
+    if (removeAbortListener) {
+      removeAbortListener();
+    }
+
+    await pool.close();
+  }
+}
+
+async function indexTarget(target, options = {}) {
+  const result = await collectState(target, options);
+  throwIfAborted(options.signal);
+
+  const shouldBackfillDatabase = !databaseExists(result.root);
+  const pathsToEmbed = shouldBackfillDatabase
+    ? Object.keys(result.filesByRelativePath)
+    : [...result.changes.new, ...result.changes.modified];
+  const embeddingSummary = await embedFiles(result, pathsToEmbed, options);
+  throwIfAborted(options.signal);
+
+  const databaseSummary = writeEmbeddingResults(
+    result.root,
+    result.filesByRelativePath,
+    embeddingSummary,
+    result.changes.removed,
+  );
+  throwIfAborted(options.signal);
+
+  const indexedAt = new Date().toISOString();
+  const failedPaths = new Set(
+    embeddingSummary.files
+      .filter((fileResult) => !fileResult.ok)
+      .map((fileResult) => fileResult.relativePath),
+  );
+
+  for (const filePath of [...result.changes.new, ...result.changes.modified]) {
+    if (failedPaths.has(filePath)) {
+      continue;
+    }
+
+    result.currentFiles[filePath].indexedAt = indexedAt;
+  }
+
+  for (const filePath of result.changes.unchanged) {
+    const previous = result.previousCache.files[filePath];
+    result.currentFiles[filePath].indexedAt = previous.indexedAt || indexedAt;
+  }
+
+  const filesToSave = {
+    ...result.currentFiles,
+  };
+
+  for (const filePath of failedPaths) {
+    const previous = result.previousCache.files[filePath];
+
+    if (previous) {
+      filesToSave[filePath] = previous;
+    } else {
+      delete filesToSave[filePath];
+    }
+  }
+
+  await saveCache(result.root, filesToSave);
+
+  result.embeddings = embeddingSummary;
+  result.database = databaseSummary;
+
+  return result;
+}
+
+async function statusTarget(target, options = {}) {
+  return collectState(target, options);
+}
+
+module.exports = {
+  getDefaultWorkerCount,
+  indexTarget,
+  statusTarget,
+};

package/src/parser.js

@@ -0,0 +1,116 @@
+const path = require("path");
+const Parser = require("tree-sitter");
+const JavaScript = require("tree-sitter-javascript");
+
+const javascriptParser = new Parser();
+javascriptParser.setLanguage(JavaScript);
+
+const parsers = {
+  ".cjs": parseJavaScript,
+  ".js": parseJavaScript,
+  ".jsx": parseJavaScript,
+  ".mjs": parseJavaScript,
+};
+
+function getParser(filePath) {
+  return parsers[path.extname(filePath).toLowerCase()];
+}
+
+function getNodeName(node) {
+  const nameNode = node.childForFieldName("name");
+  return nameNode ? nameNode.text : "(anonymous)";
+}
+
+function getLeadingComment(node) {
+  const previous = node.previousNamedSibling || node.previousSibling;
+
+  if (!previous || previous.type !== "comment") {
+    return "";
+  }
+
+  return previous.text.trim();
+}
+
+function createSemanticText(relativePath, chunk) {
+  const parts = [
+    `File: ${relativePath}`,
+    `Symbol: ${chunk.name}`,
+  ];
+
+  if (chunk.commentText) {
+    parts.push(`Context: ${chunk.commentText}`);
+  }
+
+  parts.push(`Code: ${chunk.code}`);
+
+  return parts.join(" | ");
+}
+
+function createChunk(relativePath, code, node) {
+  const chunkCode = code.slice(node.startIndex, node.endIndex);
+  const commentText = getLeadingComment(node);
+  const chunk = {
+    name: getNodeName(node),
+    type: node.type,
+    startLine: node.startPosition.row + 1,
+    endLine: node.endPosition.row + 1,
+    commentText,
+    code: chunkCode,
+  };
+
+  return {
+    ...chunk,
+    semanticText: createSemanticText(relativePath, chunk),
+  };
+}
+
+function parseJavaScript({ code, relativePath }) {
+  const tree = javascriptParser.parse(code);
+  const chunks = [];
+  const stack = [tree.rootNode];
+
+  while (stack.length > 0) {
+    const node = stack.pop();
+
+    if (
+      node.type === "function_declaration" ||
+      node.type === "class_declaration" ||
+      node.type === "method_definition"
+    ) {
+      chunks.push(createChunk(relativePath, code, node));
+    }
+
+    for (let index = node.childCount - 1; index >= 0; index -= 1) {
+      stack.push(node.child(index));
+    }
+  }
+
+  return chunks;
+}
+
+function parseSourceFile({ filePath, relativePath, code }) {
+  const parse = getParser(filePath);
+
+  if (!parse) {
+    return {
+      supported: false,
+      reason: `Unsupported source extension: ${path.extname(filePath).toLowerCase() || "(none)"}`,
+      chunks: [],
+    };
+  }
+
+  return {
+    supported: true,
+    chunks: parse({
+      code,
+      relativePath,
+    }),
+  };
+}
+
+module.exports = {
+  getParser,
+  parseJavaScript,
+  parseSourceFile,
+  parsers,
+};

package/src/paths.js

@@ -0,0 +1,14 @@
+const path = require("path");
+
+function toPosixPath(value) {
+  return value.split(path.sep).join("/");
+}
+
+function toRelativePosixPath(root, target) {
+  return toPosixPath(path.relative(root, target));
+}
+
+module.exports = {
+  toPosixPath,
+  toRelativePosixPath,
+};

package/src/search.js

@@ -0,0 +1,134 @@
+const path = require("path");
+const { pipeline } = require("@xenova/transformers");
+const { throwIfAborted } = require("./abort");
+const { getDatabasePath, openDatabase } = require("./database");
+
+const MODEL_ID = "Xenova/all-MiniLM-L6-v2";
+
+let embedderPromise;
+
+function getEmbedder() {
+  if (!embedderPromise) {
+    embedderPromise = pipeline("feature-extraction", MODEL_ID);
+  }
+
+  return embedderPromise;
+}
+
+function toFloat32Array(value) {
+  if (value instanceof Float32Array) {
+    return new Float32Array(value);
+  }
+
+  if (ArrayBuffer.isView(value)) {
+    return Float32Array.from(value);
+  }
+
+  if (Array.isArray(value)) {
+    return Float32Array.from(value);
+  }
+
+  throw new Error("Embedding output could not be converted to Float32Array");
+}
+
+async function embedQuery(query) {
+  const embedder = await getEmbedder();
+  const output = await embedder(query, {
+    pooling: "mean",
+    normalize: true,
+  });
+
+  return toFloat32Array(output.data);
+}
+
+function blobToFloat32Array(blob) {
+  return new Float32Array(
+    blob.buffer,
+    blob.byteOffset,
+    blob.byteLength / Float32Array.BYTES_PER_ELEMENT,
+  );
+}
+
+function cosineSimilarity(left, right) {
+  if (left.length !== right.length) {
+    return Number.NEGATIVE_INFINITY;
+  }
+
+  let dot = 0;
+  let magnitudeLeft = 0;
+  let magnitudeRight = 0;
+
+  for (let index = 0; index < left.length; index += 1) {
+    dot += left[index] * right[index];
+    magnitudeLeft += left[index] * left[index];
+    magnitudeRight += right[index] * right[index];
+  }
+
+  if (magnitudeLeft === 0 || magnitudeRight === 0) {
+    return 0;
+  }
+
+  return dot / (Math.sqrt(magnitudeLeft) * Math.sqrt(magnitudeRight));
+}
+
+function loadChunks(db) {
+  return db.prepare(`
+    SELECT
+      chunks.id,
+      files.path AS file,
+      chunks.name,
+      chunks.type,
+      chunks.startLine,
+      chunks.endLine,
+      chunks.code,
+      chunks.embedding
+    FROM chunks
+    INNER JOIN files ON files.id = chunks.file_id
+  `).all();
+}
+
+async function searchTarget(target, query, options = {}) {
+  const root = path.resolve(target || ".");
+  const limit = options.limit || 5;
+  throwIfAborted(options.signal);
+
+  const queryEmbedding = await embedQuery(query);
+  throwIfAborted(options.signal);
+
+  const db = openDatabase(root);
+
+  try {
+    const rows = loadChunks(db);
+    throwIfAborted(options.signal);
+
+    const results = rows
+      .map((row) => ({
+        id: row.id,
+        file: row.file,
+        name: row.name,
+        type: row.type,
+        startLine: row.startLine,
+        endLine: row.endLine,
+        code: row.code,
+        score: cosineSimilarity(queryEmbedding, blobToFloat32Array(row.embedding)),
+      }))
+      .filter((result) => Number.isFinite(result.score))
+      .sort((left, right) => right.score - left.score)
+      .slice(0, limit);
+
+    return {
+      root,
+      databasePath: getDatabasePath(root),
+      query,
+      results,
+    };
+  } finally {
+    db.close();
+  }
+}
+
+module.exports = {
+  blobToFloat32Array,
+  cosineSimilarity,
+  searchTarget,
+};

package/src/shutdown.js

@@ -0,0 +1,34 @@
+const cleanupTasks = new Set();
+
+let cleanupPromise = null;
+
+function registerCleanup(cleanup) {
+  cleanupTasks.add(cleanup);
+
+  return () => {
+    cleanupTasks.delete(cleanup);
+  };
+}
+
+async function runCleanup() {
+  if (cleanupPromise) {
+    return cleanupPromise;
+  }
+
+  cleanupPromise = (async () => {
+    const tasks = Array.from(cleanupTasks).reverse();
+    cleanupTasks.clear();
+    await Promise.allSettled(tasks.map((cleanup) => cleanup()));
+  })();
+
+  try {
+    return await cleanupPromise;
+  } finally {
+    cleanupPromise = null;
+  }
+}
+
+module.exports = {
+  registerCleanup,
+  runCleanup,
+};