smart-coding-mcp 1.4.1 → 2.0.0

package/lib/ast-chunker.js

@@ -0,0 +1,273 @@
+/**
+ * AST-based Code Chunker
+ * 
+ * Uses Tree-sitter to parse code and chunk at semantic boundaries
+ * (functions, classes, methods) instead of arbitrary line splits.
+ */
+
+import Parser from 'web-tree-sitter';
+import path from 'path';
+import fs from 'fs/promises';
+import { fileURLToPath } from 'url';
+import { smartChunk } from './utils.js'; // Fallback
+
+const __dirname = path.dirname(fileURLToPath(import.meta.url));
+
+// Mapping of file extensions to Tree-sitter language names
+const LANGUAGE_MAP = {
+  js: 'javascript',
+  mjs: 'javascript',
+  cjs: 'javascript',
+  jsx: 'javascript',
+  ts: 'typescript',
+  tsx: 'typescript',
+  py: 'python',
+  go: 'go',
+  rs: 'rust',
+  rb: 'ruby',
+  java: 'java',
+  c: 'c',
+  cpp: 'cpp',
+  h: 'c',
+  hpp: 'cpp'
+};
+
+// Node types that represent semantic boundaries
+const SEMANTIC_NODES = {
+  javascript: ['function_declaration', 'arrow_function', 'class_declaration', 'method_definition', 'export_statement'],
+  typescript: ['function_declaration', 'arrow_function', 'class_declaration', 'method_definition', 'export_statement'],
+  python: ['function_definition', 'class_definition', 'decorated_definition'],
+  go: ['function_declaration', 'method_declaration', 'type_declaration'],
+  rust: ['function_item', 'impl_item', 'struct_item', 'enum_item'],
+  ruby: ['method', 'class', 'module'],
+  java: ['method_declaration', 'class_declaration', 'interface_declaration'],
+  c: ['function_definition', 'struct_specifier'],
+  cpp: ['function_definition', 'class_specifier', 'struct_specifier']
+};
+
+export class ASTChunker {
+  constructor(config) {
+    this.config = config;
+    this.parser = null;
+    this.languages = new Map();
+    this.initialized = false;
+  }
+
+  /**
+   * Initialize Tree-sitter parser
+   */
+  async init() {
+    if (this.initialized) return;
+    
+    try {
+      await Parser.init();
+      this.parser = new Parser();
+      this.initialized = true;
+      console.error('[AST] Tree-sitter parser initialized');
+    } catch (error) {
+      console.error('[AST] Failed to initialize Tree-sitter:', error.message);
+      throw error;
+    }
+  }
+
+  /**
+   * Load a language grammar
+   */
+  async loadLanguage(langName) {
+    if (this.languages.has(langName)) {
+      return this.languages.get(langName);
+    }
+
+    try {
+      // Try to find the WASM file in node_modules
+      const possiblePaths = [
+        path.join(__dirname, '..', 'node_modules', `tree-sitter-${langName}`, `tree-sitter-${langName}.wasm`),
+        path.join(__dirname, '..', 'node_modules', 'tree-sitter-wasms', 'out', `tree-sitter-${langName}.wasm`),
+        path.join(__dirname, '..', 'grammars', `tree-sitter-${langName}.wasm`)
+      ];
+
+      for (const wasmPath of possiblePaths) {
+        try {
+          await fs.access(wasmPath);
+          const language = await Parser.Language.load(wasmPath);
+          this.languages.set(langName, language);
+          if (this.config.verbose) {
+            console.error(`[AST] Loaded ${langName} grammar from ${wasmPath}`);
+          }
+          return language;
+        } catch {
+          continue;
+        }
+      }
+
+      console.error(`[AST] No grammar found for ${langName}`);
+      return null;
+    } catch (error) {
+      console.error(`[AST] Failed to load ${langName}:`, error.message);
+      return null;
+    }
+  }
+
+  /**
+   * Get the language name from file extension
+   */
+  getLanguageForFile(file) {
+    const ext = path.extname(file).slice(1).toLowerCase();
+    return LANGUAGE_MAP[ext] || null;
+  }
+
+  /**
+   * Chunk code using AST analysis
+   */
+  async chunk(content, file) {
+    // Initialize if needed
+    if (!this.initialized) {
+      await this.init();
+    }
+
+    const langName = this.getLanguageForFile(file);
+    
+    // Fall back to smart chunking if language not supported
+    if (!langName) {
+      if (this.config.verbose) {
+        console.error(`[AST] No AST support for ${path.extname(file)}, using smart chunking`);
+      }
+      return smartChunk(content, file, this.config);
+    }
+
+    const language = await this.loadLanguage(langName);
+    
+    // Fall back if grammar not available
+    if (!language) {
+      return smartChunk(content, file, this.config);
+    }
+
+    try {
+      this.parser.setLanguage(language);
+      const tree = this.parser.parse(content);
+      const chunks = [];
+      const lines = content.split('\n');
+      const semanticNodes = SEMANTIC_NODES[langName] || [];
+
+      // Walk the AST and extract semantic chunks
+      this.walkTree(tree.rootNode, (node) => {
+        if (semanticNodes.includes(node.type)) {
+          const startLine = node.startPosition.row;
+          const endLine = node.endPosition.row;
+          
+          // Skip very small nodes (< 3 lines)
+          if (endLine - startLine < 2) return;
+          
+          // Extract the text for this node
+          const chunkLines = lines.slice(startLine, endLine + 1);
+          const text = chunkLines.join('\n');
+          
+          // Skip if too large (will be handled by split)
+          const targetTokens = this.config.chunkSize * 4; // Rough estimate
+          if (text.length > targetTokens * 4) {
+            // Split large nodes
+            this.splitLargeNode(node, lines, chunks);
+          } else {
+            chunks.push({
+              text,
+              startLine: startLine + 1, // 1-indexed
+              endLine: endLine + 1,
+              nodeType: node.type
+            });
+          }
+        }
+      });
+
+      // If no semantic chunks found, fall back to smart chunking
+      if (chunks.length === 0) {
+        return smartChunk(content, file, this.config);
+      }
+
+      // Sort by start line
+      chunks.sort((a, b) => a.startLine - b.startLine);
+
+      // Merge small gaps and remove overlaps
+      return this.mergeAndCleanChunks(chunks, lines);
+
+    } catch (error) {
+      console.error(`[AST] Parse error for ${file}:`, error.message);
+      return smartChunk(content, file, this.config);
+    }
+  }
+
+  /**
+   * Walk the AST tree and call callback for each node
+   */
+  walkTree(node, callback) {
+    callback(node);
+    for (let i = 0; i < node.childCount; i++) {
+      this.walkTree(node.child(i), callback);
+    }
+  }
+
+  /**
+   * Split large AST nodes into smaller chunks
+   */
+  splitLargeNode(node, lines, chunks) {
+    const chunkSize = this.config.chunkSize || 25;
+    const startLine = node.startPosition.row;
+    const endLine = node.endPosition.row;
+    
+    for (let i = startLine; i <= endLine; i += chunkSize) {
+      const chunkEnd = Math.min(i + chunkSize - 1, endLine);
+      const chunkLines = lines.slice(i, chunkEnd + 1);
+      
+      chunks.push({
+        text: chunkLines.join('\n'),
+        startLine: i + 1,
+        endLine: chunkEnd + 1,
+        nodeType: node.type + '_part'
+      });
+    }
+  }
+
+  /**
+   * Merge small chunks and clean up overlaps
+   */
+  mergeAndCleanChunks(chunks, lines) {
+    const cleaned = [];
+    const minSize = 5; // Minimum lines per chunk
+
+    for (const chunk of chunks) {
+      // Skip if overlaps with previous
+      if (cleaned.length > 0) {
+        const prev = cleaned[cleaned.length - 1];
+        if (chunk.startLine <= prev.endLine) {
+          // Extend previous chunk if this one extends further
+          if (chunk.endLine > prev.endLine) {
+            prev.endLine = chunk.endLine;
+            const extendedLines = lines.slice(prev.startLine - 1, prev.endLine);
+            prev.text = extendedLines.join('\n');
+          }
+          continue;
+        }
+      }
+
+      // Add to cleaned list
+      cleaned.push(chunk);
+    }
+
+    return cleaned;
+  }
+}
+
+/**
+ * Factory function to get the appropriate chunker based on config
+ */
+export function getChunker(config) {
+  if (config.chunkingMode === 'ast') {
+    return new ASTChunker(config);
+  }
+  
+  // Return a wrapper that uses smartChunk
+  return {
+    async chunk(content, file) {
+      return smartChunk(content, file, config);
+    }
+  };
+}

package/lib/config.js

@@ -62,7 +62,10 @@ const DEFAULT_CONFIG = {
   watchFiles: false,
   verbose: false,
   workerThreads: "auto", // "auto" = CPU cores - 1, or set a number
-  embeddingModel: "Xenova/all-MiniLM-L6-v2",
+  embeddingModel: "nomic-ai/nomic-embed-text-v1.5",
+  embeddingDimension: 256, // MRL dimension: 64, 128, 256, 512, 768
+  device: "auto", // "cpu", "webgpu", or "auto"
+  chunkingMode: "smart", // "smart", "ast", or "line"
   semanticWeight: 0.7,
   exactMatchBoost: 1.5,
   smartIndexing: true
@@ -237,6 +240,42 @@ export async function loadConfig(workspaceDir = null) {
     }
   }
   
+  // MRL embedding dimension
+  if (process.env.SMART_CODING_EMBEDDING_DIMENSION !== undefined) {
+    const value = parseInt(process.env.SMART_CODING_EMBEDDING_DIMENSION, 10);
+    const validDims = [64, 128, 256, 512, 768];
+    if (validDims.includes(value)) {
+      config.embeddingDimension = value;
+      console.error(`[Config] Using embedding dimension: ${value}`);
+    } else {
+      console.error(`[Config] Invalid SMART_CODING_EMBEDDING_DIMENSION: ${value}, using default (must be 64, 128, 256, 512, or 768)`);
+    }
+  }
+  
+  // Device selection
+  if (process.env.SMART_CODING_DEVICE !== undefined) {
+    const value = process.env.SMART_CODING_DEVICE.trim().toLowerCase();
+    const validDevices = ['cpu', 'webgpu', 'auto'];
+    if (validDevices.includes(value)) {
+      config.device = value;
+      console.error(`[Config] Using device: ${value}`);
+    } else {
+      console.error(`[Config] Invalid SMART_CODING_DEVICE: ${value}, using default (must be 'cpu', 'webgpu', or 'auto')`);
+    }
+  }
+  
+  // Chunking mode
+  if (process.env.SMART_CODING_CHUNKING_MODE !== undefined) {
+    const value = process.env.SMART_CODING_CHUNKING_MODE.trim().toLowerCase();
+    const validModes = ['smart', 'ast', 'line'];
+    if (validModes.includes(value)) {
+      config.chunkingMode = value;
+      console.error(`[Config] Using chunking mode: ${value}`);
+    } else {
+      console.error(`[Config] Invalid SMART_CODING_CHUNKING_MODE: ${value}, using default (must be 'smart', 'ast', or 'line')`);
+    }
+  }
+  
   return config;
 }
 

package/lib/embedding-worker.js

@@ -1,12 +1,38 @@
 import { parentPort, workerData } from "worker_threads";
-import { pipeline } from "@xenova/transformers";
+import { pipeline, layer_norm } from "@huggingface/transformers";
 
 let embedder = null;
+const VALID_DIMENSIONS = [64, 128, 256, 512, 768];
 
 // Initialize the embedding model once when worker starts
 async function initializeEmbedder() {
   if (!embedder) {
-    embedder = await pipeline("feature-extraction", workerData.embeddingModel);
+    const modelName = workerData.embeddingModel || 'nomic-ai/nomic-embed-text-v1.5';
+    const dimension = workerData.embeddingDimension || 256;
+    const targetDim = VALID_DIMENSIONS.includes(dimension) ? dimension : 256;
+    const isNomic = modelName.includes('nomic');
+    
+    const extractor = await pipeline("feature-extraction", modelName);
+    
+    if (isNomic) {
+      // MRL embedder with dimension slicing
+      embedder = async function(text, options = {}) {
+        let embeddings = await extractor(text, { pooling: 'mean' });
+        embeddings = layer_norm(embeddings, [embeddings.dims[1]])
+          .slice(null, [0, targetDim])
+          .normalize(2, -1);
+        return { data: embeddings.data };
+      };
+      embedder.dimension = targetDim;
+    } else {
+      // Legacy embedder (MiniLM etc.)
+      embedder = async function(text, options = {}) {
+        return await extractor(text, { pooling: 'mean', normalize: true });
+      };
+      embedder.dimension = 384;
+    }
+    
+    embedder.modelName = modelName;
   }
   return embedder;
 }
@@ -65,3 +91,4 @@ initializeEmbedder().then(() => {
 }).catch((error) => {
   parentPort.postMessage({ type: "error", error: error.message });
 });
+

package/lib/mrl-embedder.js

@@ -0,0 +1,133 @@
+/**
+ * MRL (Matryoshka Representation Learning) Embedder
+ * 
+ * Provides flexible embedding dimensions (64, 128, 256, 512, 768) using
+ * nomic-embed-text-v1.5 with layer normalization and dimension slicing.
+ */
+
+import { pipeline, layer_norm } from '@huggingface/transformers';
+
+// Valid MRL dimensions for nomic-embed-text-v1.5
+const VALID_DIMENSIONS = [64, 128, 256, 512, 768];
+
+/**
+ * Create an MRL-enabled embedder with configurable output dimensions
+ * 
+ * @param {string} modelName - Model identifier (e.g., 'nomic-ai/nomic-embed-text-v1.5')
+ * @param {object} options - Configuration options
+ * @param {number} options.dimension - Target embedding dimension (64, 128, 256, 512, 768)
+ * @param {string} options.device - Device to use ('cpu', 'webgpu', 'auto')
+ * @returns {Function} Embedder function compatible with existing codebase
+ */
+export async function createMRLEmbedder(modelName, options = {}) {
+  const dimension = options.dimension || 256;
+  const device = options.device || 'cpu';
+  
+  // Validate dimension
+  if (!VALID_DIMENSIONS.includes(dimension)) {
+    console.error(`[MRL] Invalid dimension ${dimension}, using 256. Valid: ${VALID_DIMENSIONS.join(', ')}`);
+  }
+  
+  const targetDim = VALID_DIMENSIONS.includes(dimension) ? dimension : 256;
+  
+  console.error(`[MRL] Loading ${modelName} (output: ${targetDim}d, device: ${device})`);
+  
+  // Detect best device if auto
+  const finalDevice = device === 'auto' ? detectBestDevice() : device;
+  
+  // Create the feature extraction pipeline
+  const pipelineOptions = {};
+  if (finalDevice === 'webgpu') {
+    pipelineOptions.device = 'webgpu';
+  }
+  
+  const extractor = await pipeline('feature-extraction', modelName, pipelineOptions);
+  
+  console.error(`[MRL] Model loaded on ${finalDevice}`);
+  
+  /**
+   * Embed text with MRL dimension slicing
+   * Compatible with existing embedder(text, options) signature
+   */
+  async function embed(text, embedOptions = {}) {
+    // Generate full 768d embedding
+    let embeddings = await extractor(text, { pooling: 'mean' });
+    
+    // Apply MRL: layer_norm -> slice -> normalize
+    embeddings = layer_norm(embeddings, [embeddings.dims[1]])
+      .slice(null, [0, targetDim])
+      .normalize(2, -1);
+    
+    // Return in format compatible with existing code (has .data property)
+    return {
+      data: embeddings.data,
+      dims: [embeddings.dims[0], targetDim]
+    };
+  }
+  
+  // Attach metadata
+  embed.modelName = modelName;
+  embed.dimension = targetDim;
+  embed.device = finalDevice;
+  
+  return embed;
+}
+
+/**
+ * Detect best available device for inference
+ */
+function detectBestDevice() {
+  // WebGPU check (browser environment)
+  if (typeof navigator !== 'undefined' && navigator.gpu) {
+    return 'webgpu';
+  }
+  
+  // Node.js with experimental WebGPU (Node 20+)
+  // This would require --experimental-webgpu flag
+  // For now, default to CPU in Node.js
+  return 'cpu';
+}
+
+/**
+ * Create a legacy-compatible embedder (384d, MiniLM)
+ * Used as fallback if MRL model fails to load
+ */
+export async function createLegacyEmbedder(modelName = 'Xenova/all-MiniLM-L6-v2') {
+  console.error(`[Embedder] Loading legacy model: ${modelName}`);
+  const extractor = await pipeline('feature-extraction', modelName);
+  
+  async function embed(text, options = {}) {
+    const output = await extractor(text, { pooling: 'mean', normalize: true });
+    return output;
+  }
+  
+  embed.modelName = modelName;
+  embed.dimension = 384;
+  embed.device = 'cpu';
+  
+  return embed;
+}
+
+/**
+ * Smart embedder factory - picks MRL or legacy based on config
+ */
+export async function createEmbedder(config) {
+  const model = config.embeddingModel || 'nomic-ai/nomic-embed-text-v1.5';
+  const dimension = config.embeddingDimension || 256;
+  const device = config.device || 'cpu';
+  
+  // Use MRL for nomic models
+  if (model.includes('nomic')) {
+    try {
+      return await createMRLEmbedder(model, { dimension, device });
+    } catch (err) {
+      console.error(`[Embedder] MRL model failed: ${err.message}, falling back to legacy`);
+      return await createLegacyEmbedder();
+    }
+  }
+  
+  // Use legacy for MiniLM and other models
+  return await createLegacyEmbedder(model);
+}
+
+export { VALID_DIMENSIONS };

package/lib/tokenizer.js

@@ -10,6 +10,10 @@
  * Each model has its own maximum sequence length
  */
 export const MODEL_TOKEN_LIMITS = {
+  // MRL / Nomic models (longer context)
+  "nomic-ai/nomic-embed-text-v1.5": 8192,
+  "nomic-ai/nomic-embed-text-v1": 2048,
+  
   // Sentence Transformers / MiniLM family
   "Xenova/all-MiniLM-L6-v2": 256,
   "Xenova/all-MiniLM-L12-v2": 256,

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "smart-coding-mcp",
-  "version": "1.4.1",
+  "version": "2.0.0",
   "description": "An extensible MCP server that enhances coding productivity with AI-powered features including semantic code search, intelligent indexing, and more, using local LLMs",
   "type": "module",
   "main": "index.js",
@@ -45,11 +45,13 @@
   "homepage": "https://github.com/omar-haris/smart-coding-mcp#readme",
   "license": "MIT",
   "dependencies": {
+    "@huggingface/transformers": "^3.8.1",
     "@modelcontextprotocol/sdk": "^1.0.4",
-    "@xenova/transformers": "^2.17.2",
     "chokidar": "^3.5.3",
+    "fastembed": "^2.1.0",
     "fdir": "^6.5.0",
-    "glob": "^10.3.10"
+    "glob": "^10.3.10",
+    "web-tree-sitter": "^0.24.6"
   },
   "engines": {
     "node": ">=18.0.0"

package/test/ast-chunker.test.js

@@ -0,0 +1,105 @@
+/**
+ * Tests for AST Chunker
+ * 
+ * Tests the AST-based code chunking functionality:
+ * - Tree-sitter initialization
+ * - Language detection
+ * - Semantic chunking vs smart chunking fallback
+ */
+
+import { describe, it, expect, beforeAll } from 'vitest';
+import { ASTChunker, getChunker } from '../lib/ast-chunker.js';
+import { loadConfig } from '../lib/config.js';
+
+describe('AST Chunker', () => {
+  let config;
+  
+  beforeAll(async () => {
+    config = await loadConfig();
+  });
+
+  describe('Chunker Factory', () => {
+    it('should return AST chunker when mode is ast', () => {
+      const chunker = getChunker({ ...config, chunkingMode: 'ast' });
+      expect(chunker).toBeInstanceOf(ASTChunker);
+    });
+
+    it('should return smart chunker wrapper when mode is smart', () => {
+      const chunker = getChunker({ ...config, chunkingMode: 'smart' });
+      expect(typeof chunker.chunk).toBe('function');
+      expect(chunker).not.toBeInstanceOf(ASTChunker);
+    });
+  });
+
+  describe('Language Detection', () => {
+    it('should detect JavaScript files', () => {
+      const chunker = new ASTChunker(config);
+      expect(chunker.getLanguageForFile('test.js')).toBe('javascript');
+      expect(chunker.getLanguageForFile('test.mjs')).toBe('javascript');
+      expect(chunker.getLanguageForFile('test.jsx')).toBe('javascript');
+    });
+
+    it('should detect TypeScript files', () => {
+      const chunker = new ASTChunker(config);
+      expect(chunker.getLanguageForFile('test.ts')).toBe('typescript');
+      expect(chunker.getLanguageForFile('test.tsx')).toBe('typescript');
+    });
+
+    it('should detect Python files', () => {
+      const chunker = new ASTChunker(config);
+      expect(chunker.getLanguageForFile('test.py')).toBe('python');
+    });
+
+    it('should return null for unsupported files', () => {
+      const chunker = new ASTChunker(config);
+      expect(chunker.getLanguageForFile('test.sql')).toBeNull();
+      expect(chunker.getLanguageForFile('test.md')).toBeNull();
+    });
+  });
+
+  describe('Fallback Behavior', () => {
+    it('should fall back to smart chunking for unsupported languages', async () => {
+      const chunker = new ASTChunker(config);
+      const sqlContent = 'SELECT * FROM users WHERE id = 1;';
+      
+      const chunks = await chunker.chunk(sqlContent, 'query.sql');
+      expect(Array.isArray(chunks)).toBe(true);
+    });
+
+    it('should handle empty content', async () => {
+      const chunker = new ASTChunker(config);
+      const chunks = await chunker.chunk('', 'empty.js');
+      expect(Array.isArray(chunks)).toBe(true);
+    });
+  });
+
+  describe('JavaScript Chunking', () => {
+    it('should chunk JavaScript functions', async () => {
+      const chunker = new ASTChunker(config);
+      const jsCode = `
+function add(a, b) {
+  return a + b;
+}
+
+function multiply(a, b) {
+  return a * b;
+}
+
+class Calculator {
+  constructor() {
+    this.result = 0;
+  }
+  
+  add(n) {
+    this.result += n;
+    return this;
+  }
+}
+`;
+      
+      const chunks = await chunker.chunk(jsCode, 'calc.js');
+      expect(Array.isArray(chunks)).toBe(true);
+      // Should have found some chunks (exact number depends on Tree-sitter grammar availability)
+    });
+  });
+});