vecfinder

vecfinder.js (23199B)

---

 #!/usr/bin/env node
 
 const fs = require('node:fs/promises');
 const path = require('node:path');
 const http = require('node:http');
 const readline = require('node:readline');
 const { stdin, stdout, stderr } = process;
 const EventEmitter = require('node:events');
 const { execSync, spawn } = require('node:child_process');
 
 const ONE_GIB = 1073741824;
 
 let config = {
   endpoint: 'http://localhost:1234/v1/embeddings',
   model: 'text-embedding-qwen3-0.6b-text-embedding',
   apiKey: '',
   parallelism: 2,
   cacheMaxSize: ONE_GIB,
 };
 
 let CACHE_DIR = null;
 let CONFIG_PATH = null;
 
 async function getWritableCacheDir(cacheBase) {
   try {
     await fs.mkdir(cacheBase, { recursive: true });
     await fs.access(cacheBase, fs.constants.W_OK);
     return cacheBase;
   } catch {
     return null;
   }
 }
 
 async function createSystemCacheDir() {
   const cacheBase = '/var/cache/vecfinder';
 
   try {
     const vecfinderGroupExists = await new Promise((resolve) => {
       execSync('getent group vecfinder > /dev/null 2>&1', { stdio: 'ignore' });
       resolve(true);
     }).catch(() => false);
 
     if (!vecfinderGroupExists) {
       return null;
     }
 
     const stat = await fs.stat('/var/cache').catch(() => null);
     if (!stat || !(stat.mode & fs.constants.S_IWUSR)) {
       return null;
     }
 
     try {
       await fs.mkdir(cacheBase, { recursive: true });
     } catch (err) {
       if (err.code !== 'EEXIST') return null;
     }
 
     const cacheStat = await fs.stat(cacheBase).catch(() => null);
     if (!cacheStat) return null;
 
     const mode = cacheStat.mode & 0o777;
     const isGroupWritable = !!(mode & fs.constants.S_IWGRP);
 
     if (isGroupWritable) {
       return await getWritableCacheDir(cacheBase);
     }
 
     try {
       execSync(`chmod 0770 "${cacheBase}"`, { stdio: 'ignore' });
     } catch {
       return null;
     }
 
     return await getWritableCacheDir(cacheBase);
   } catch {
     return null;
   }
 }
 
 async function discoverConfig() {
   const userConfig = path.join(process.env.HOME || process.env.USERPROFILE, '.config', 'vecfinder', 'config.json');
   const systemConfig1 = '/etc/vecfinder/config.json';
   const systemConfig2 = '/etc/vecfinder.json';
 
   if (await fileExists(userConfig)) {
     CONFIG_PATH = userConfig;
     CACHE_DIR = path.join(process.env.HOME || process.env.USERPROFILE, '.cache', 'vecfinder');
     return;
   }
 
   if (await fileExists(systemConfig1) || await fileExists(systemConfig2)) {
     const systemConfig = await fileExists(systemConfig1) ? systemConfig1 : systemConfig2;
     CONFIG_PATH = systemConfig;
 
     try {
       const data = JSON.parse(await fs.readFile(systemConfig, 'utf8'));
       if (data.cacheDir && typeof data.cacheDir === 'string') {
         CACHE_DIR = await getWritableCacheDir(data.cacheDir);
       }
 
       if (!CACHE_DIR) {
         CACHE_DIR = await createSystemCacheDir();
       }
     } catch {}
 
     if (!CACHE_DIR) {
       CACHE_DIR = path.join(process.env.HOME || process.env.USERPROFILE, '.cache', 'vecfinder');
     }
     return;
   }
 
   CONFIG_PATH = userConfig;
   CACHE_DIR = path.join(process.env.HOME || process.env.USERPROFILE, '.cache', 'vecfinder');
 }
 
 async function loadConfig() {
   await discoverConfig();
 
   if (!CONFIG_PATH || !(await fileExists(CONFIG_PATH))) {
     return;
   }
 
   try {
     const data = await fs.readFile(CONFIG_PATH, 'utf8');
     config = { ...config, ...JSON.parse(data) };
   } catch {}
 }
 
 let hasBat = null;
 
 function checkBat() {
   if (hasBat !== null) return hasBat;
   try {
     execSync('which bat', { stdio: 'ignore' });
     hasBat = true;
   } catch {
     hasBat = false;
   }
   return hasBat;
 }
 
 function stripAnsi(str) {
   return str.replace(/\x1b\[[0-9;]*m/g, '');
 }
 
 async function fileExists(filePath) {
   try {
     await fs.access(filePath);
     return true;
   } catch {
     return false;
   }
 }
 
 // Cache
 function cachePath(filePath) {
   return path.join(CACHE_DIR, filePath);
 }
 
 async function getCachedVector(filePath) {
   try {
     const cp = cachePath(filePath);
     const [cacheStat, fileStat] = await Promise.all([fs.stat(cp), fs.stat(filePath)]);
     if (fileStat.mtimeMs > cacheStat.mtimeMs) return null;
     const buf = await fs.readFile(cp);
     return Array.from(new Float64Array(buf.buffer, buf.byteOffset, buf.byteLength / 8));
   } catch {
     return null;
   }
 }
 
 async function setCachedVector(filePath, vector) {
   try {
     const cp = cachePath(filePath);
     await fs.mkdir(path.dirname(cp), { recursive: true });
     await fs.writeFile(cp, Buffer.from(new Float64Array(vector).buffer));
     enforceCacheLimit();
   } catch {}
 }
 
 async function enforceCacheLimit() {
   if (!CACHE_DIR) return;
 
   try {
     const files = [];
     let totalSize = 0;
 
     const pending = [CACHE_DIR];
     while (pending.length) {
       const current = pending.pop();
       try {
         const entries = await fs.readdir(current, { withFileTypes: true });
         for (const entry of entries) {
           const fullPath = path.join(current, entry.name);
           if (entry.isDirectory()) {
             pending.push(fullPath);
           } else {
             const stat = await fs.stat(fullPath);
             files.push({ path: fullPath, size: stat.size });
             totalSize += stat.size;
           }
         }
       } catch {}
     }
 
     while (totalSize > config.cacheMaxSize && files.length > 0) {
       const i = Math.floor(Math.random() * files.length);
       const [removed] = files.splice(i, 1);
       await fs.unlink(removed.path);
       totalSize -= removed.size;
     }
   } catch {}
 }
 
 // Embedding
 let currentQueryId = 0;
 
 function normalize(vec) {
   let norm = 0;
   for (let i = 0; i < vec.length; i++) norm += vec[i] * vec[i];
   norm = Math.sqrt(norm);
   if (norm === 0) return vec;
   return vec.map(v => v / norm);
 }
 
 function dot(a, b) {
   let sum = 0;
   for (let i = 0; i < a.length; i++) sum += a[i] * b[i];
   return sum;
 }
 
 async function fetchEmbedding(text) {
   return new Promise((resolve, reject) => {
     const req = http.request(config.endpoint, {
       method: 'POST',
       headers: {
         'Content-Type': 'application/json',
         ...(config.apiKey ? { 'Authorization': `Bearer ${config.apiKey}` } : {}),
       },
     }, (res) => {
       let data = '';
       res.on('data', chunk => data += chunk);
       res.on('end', () => {
         try {
           const json = JSON.parse(data);
           if (!json.data?.[0]?.embedding) {
             return reject(new Error('Invalid embedding response'));
           }
           resolve(json.data[0].embedding);
         } catch (e) {
           reject(e);
         }
       });
     });
     req.on('error', reject);
     req.write(JSON.stringify({ model: config.model, input: text }));
     req.end();
   });
 }
 
 // Document store with events
 class DocumentStore extends EventEmitter {
   constructor() {
     super();
     this.documents = new Map();
     this.queryVec = null;
     this.consumersActive = 0;
   }
 
   async addDocument(id, content, metadata = {}, vector) {
     if (!vector) {
       const raw = await fetchEmbedding(content);
       vector = normalize(raw);
     }
     this.documents.set(id, { content, vector, metadata });
     this.emit('documentAdded');
   }
 
   setQuery(queryText) {
     currentQueryId++;
     const queryId = currentQueryId;
 
     if (queryText.trim() === '') {
       this.queryVec = null;
       this.emit('queryUpdated');
       return;
     }
 
     fetchEmbedding(queryText).then(raw => {
       if (queryId !== currentQueryId) return;
       this.queryVec = normalize(raw);
       this.emit('queryUpdated');
     }).catch(err => {
       console.error(`Embedding error: ${err.message}`);
     });
   }
 
   getResults() {
     const results = [];
     for (const [id, doc] of this.documents.entries()) {
       const score = this.queryVec ? dot(this.queryVec, doc.vector) : 0;
       results.push({
         id,
         score,
         content: doc.content,
         metadata: doc.metadata,
       });
     }
 
     if (this.queryVec) {
       results.sort((a, b) => b.score - a.score);
     } else {
       results.sort((a, b) => a.id.localeCompare(b.id));
     }
 
     return results;
   }
 }
 
 // File indexing
 async function isTextFile(filePath) {
   try {
     const chunk = await fs.readFile(filePath, { length: 512 });
     for (let i = 0; i < chunk.length; i++) {
       if (chunk[i] === 0) return false;
     }
     return true;
   } catch {
     return false;
   }
 }
 
 async function walkDir(dirPath) {
   const ignore = new Set(['node_modules', '.git', 'dist', 'build', 'coverage', '.next', '.nuxt']);
   const files = [];
   const pending = [dirPath];
 
   while (pending.length) {
     const current = pending.pop();
     try {
       const entries = await fs.readdir(current, { withFileTypes: true });
       for (const entry of entries) {
         const fullPath = path.join(current, entry.name);
         if (entry.isDirectory()) {
           if (!ignore.has(entry.name)) pending.push(fullPath);
         } else {
           files.push(fullPath);
         }
       }
     } catch {}
   }
 
   return files;
 }
 
 async function indexDirectory(store, dirPath) {
   const files = await walkDir(dirPath);
   const queue = [];
   let totalFiles = 0;
   let completedFiles = 0;
 
   for (const filePath of files) {
     const stats = await fs.stat(filePath).catch(() => null);
     if (!stats || stats.size > 1024 * 1024) continue;
     const text = await isTextFile(filePath);
     if (!text) continue;
     queue.push(filePath);
     totalFiles++;
   }
 
   function consumer() {
     if (queue.length === 0 && store.consumersActive === 0) return;
 
     if (store.consumersActive >= config.parallelism || queue.length === 0) return;
 
     store.consumersActive++;
     const filePath = queue.shift();
 
     (async () => {
       try {
         const content = await fs.readFile(filePath, 'utf8');
         const id = path.relative(dirPath, filePath);
 
         let vector = await getCachedVector(filePath);
         if (vector) {
           await store.addDocument(id, content, { path: filePath }, vector);
         } else {
           const raw = await fetchEmbedding(content);
           vector = normalize(raw);
           await store.addDocument(id, content, { path: filePath }, vector);
           await setCachedVector(filePath, vector);
         }
       } catch {}
 
       completedFiles++;
       store.consumersActive--;
       store.emit('documentAdded');
       if (queue.length > 0 || store.consumersActive > 0) {
         setImmediate(consumer);
       }
     })();
   }
 
   for (let i = 0; i < Math.min(config.parallelism, totalFiles); i++) {
     setImmediate(consumer);
   }
 
   return new Promise(resolve => {
     const checkComplete = () => {
       if (completedFiles >= totalFiles && store.consumersActive === 0 && queue.length === 0) {
         store.removeListener('documentAdded', checkComplete);
         resolve();
       }
     };
     store.on('documentAdded', checkComplete);
     checkComplete();
   });
 }
 
 // UI
 class UI {
   constructor(store, useBatFlag = null) {
     this.store = store;
     this.query = '';
     this.selectedIndex = 0;
     this.listOffset = 0;
     this.previewOffset = 0;
     this.rl = null;
     this._previewWrapped = null;
     this._previewFileId = null;
     this._previewWidth = null;
     this.spinnerFrame = 0;
     this.spinnerInterval = null;
     this.useBat = useBatFlag !== null ? useBatFlag : checkBat();
     this._batOutput = null;
     this._batFileId = null;
     this._lastRenderState = null;
     this._needsFullRedraw = false;
 
     this.store.on('documentAdded', () => {
       this._needsFullRedraw = true;
     });
     this.store.on('queryUpdated', () => {
       this.selectedIndex = 0;
       this.listOffset = 0;
       this.previewOffset = 0;
       this._needsFullRedraw = true;
       this.render();
     });
   }
 
   init() {
     this.rl = readline.createInterface({ input: stdin, output: stderr, terminal: true });
     stdin.on('keypress', (str, key) => this.handleKeyPress(str, key));
     this.rl.on('line', () => this.handleEnter());
     this.rl.on('close', () => this.cleanup());
     this.spinnerInterval = setInterval(() => {
       this.spinnerFrame++;
       this.render();
     }, 80);
     this.render();
   }
 
   cleanup() {
     clearInterval(this.spinnerInterval);
     stdin.removeAllListeners('keypress');
     if (this.rl) this.rl.close();
   }
 
   adjustListOffset(maxHeight) {
     const results = this.store.getResults();
     const topThreshold = this.listOffset + 2;
     const bottomThreshold = this.listOffset + maxHeight - 3;
 
     if (this.selectedIndex < topThreshold) {
       this.listOffset = Math.max(0, this.selectedIndex - 2);
     } else if (this.selectedIndex > bottomThreshold) {
       this.listOffset = Math.min(results.length - maxHeight, this.selectedIndex - (maxHeight - 3));
     }
 
     this.listOffset = Math.max(0, Math.min(this.listOffset, Math.max(0, results.length - maxHeight)));
   }
 
   getMaxPreviewOffset(maxHeight) {
     if (this.useBat && this._batOutput) {
       return Math.max(0, this._batOutput.split('\n').length - maxHeight);
     }
     return this._previewWrapped ? Math.max(0, this._previewWrapped.length - maxHeight) : 0;
   }
 
   handleKeyPress(str, key) {
     if (!key) return;
 
     if (key.name === 'enter') {
       this.handleEnter();
       return;
     }
 
     if (key.name === 'up') {
       if (this.selectedIndex > 0) {
         this.selectedIndex--;
         this.previewOffset = 0;
         const height = process.stdout.rows || 24;
         const maxHeight = height - 7;
         this.adjustListOffset(maxHeight);
         this.render();
       }
       return;
     }
 
     if (key.name === 'down') {
       const results = this.store.getResults();
       if (this.selectedIndex < results.length - 1) {
         this.selectedIndex++;
         this.previewOffset = 0;
         const height = process.stdout.rows || 24;
         const maxHeight = height - 7;
         this.adjustListOffset(maxHeight);
         this.render();
       }
       return;
     }
 
     if (key.name === 'home') {
       this.selectedIndex = 0;
       this.previewOffset = 0;
       const height = process.stdout.rows || 24;
       const maxHeight = height - 7;
       this.adjustListOffset(maxHeight);
       this.render();
       return;
     }
 
     if (key.name === 'end') {
       const results = this.store.getResults();
       this.selectedIndex = results.length - 1;
       this.previewOffset = 0;
       const height = process.stdout.rows || 24;
       const maxHeight = height - 7;
       this.adjustListOffset(maxHeight);
       this.render();
       return;
     }
 
     if (key.name === 'pageup') {
       this.previewOffset = Math.max(0, this.previewOffset - 10);
       this.render();
       return;
     }
 
     if (key.name === 'pagedown') {
       const height = process.stdout.rows || 24;
       const maxHeight = height - 7;
       const maxPreviewOffset = this.getMaxPreviewOffset(maxHeight);
       this.previewOffset = Math.min(maxPreviewOffset, this.previewOffset + 10);
       this.render();
       return;
     }
 
     if (key.name === 'backspace' || key.name === 'delete') {
       this.query = this.query.slice(0, -1);
       this.store.setQuery(this.query);
       return;
     }
 
     if (str && str.length === 1 && !key.ctrl && !key.meta) {
       this.query += str;
       this.store.setQuery(this.query);
       return;
     }
   }
 
   handleEnter() {
     this.cleanup();
     const results = this.store.getResults();
     if (results.length > 0 && this.selectedIndex < results.length) {
       stdout.write(results[this.selectedIndex].id + '\n');
     } else {
       stdout.write('\n');
     }
     process.exit(0);
   }
 
   render() {
     const width = process.stdout.columns || 80;
     const height = process.stdout.rows || 24;
 
     const leftWidth = Math.floor(width * 0.35);
     const rightWidth = width - leftWidth - 1;
     const headerHeight = 3;
     const maxHeight = height - 4 - headerHeight;
 
     const results = this.store.getResults();
     const frames = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'];
     const spinner = this.store.consumersActive > 0 ? frames[this.spinnerFrame % frames.length] : '\u2713';
 
     const titleDecorWidth = '┤ VecFinder ├'.length;
     const dashSpace = leftWidth - titleDecorWidth;
     const titleCenterLeft = Math.floor(dashSpace / 2);
     const titleCenterRight = dashSpace - titleCenterLeft;
 
     const selectedResult = results[this.selectedIndex];
 
     const visibleIds = [];
     for (let row = 0; row < maxHeight; row++) {
       const result = results[row + this.listOffset];
       visibleIds.push(result ? result.id : null);
     }
 
     const currentState = {
       width,
       height,
       leftWidth,
       rightWidth,
       query: this.query,
       selectedIndex: this.selectedIndex,
       listOffset: this.listOffset,
       previewOffset: this.previewOffset,
       spinner,
       consumersActive: this.store.consumersActive,
       selectedFileId: selectedResult?.id,
       visibleIds,
       batFileId: this._batFileId,
       batOutputHash: this._batOutput ? this._batOutput.slice(0, 100) : null,
     };
 
     if (!this._needsFullRedraw && this._lastRenderState) {
       const s = this._lastRenderState;
       if (
         s.width === currentState.width &&
         s.height === currentState.height &&
         s.leftWidth === currentState.leftWidth &&
         s.rightWidth === currentState.rightWidth &&
         s.query === currentState.query &&
         s.selectedIndex === currentState.selectedIndex &&
         s.listOffset === currentState.listOffset &&
         s.previewOffset === currentState.previewOffset &&
         s.spinner === currentState.spinner &&
         s.consumersActive === currentState.consumersActive &&
         s.selectedFileId === currentState.selectedFileId &&
         s.visibleIds.length === currentState.visibleIds.length &&
         s.visibleIds.every((id, i) => id === currentState.visibleIds[i]) &&
         s.batFileId === currentState.batFileId
       ) {
         return;
       }
     }
     this._needsFullRedraw = false;
     this._lastRenderState = currentState;
 
     stderr.write('\x1b[2J\x1b[H');
 
     stderr.write('─'.repeat(titleCenterLeft) + '┤ VecFinder ├' + '─'.repeat(titleCenterRight) + '┬' + '─'.repeat(rightWidth) + '\n');
 
     let leftHeader = 'Files ' + spinner;
     leftHeader = leftHeader.padEnd(leftWidth, ' ').slice(0, leftWidth);
 
     let rightHeader = '';
     if (selectedResult) {
       const title = 'File: ' + selectedResult.id;
       rightHeader = title.padEnd(rightWidth, ' ').slice(0, rightWidth);
     } else {
       rightHeader = ''.padEnd(rightWidth, ' ');
     }
 
     stderr.write(leftHeader + '│' + rightHeader + '\n');
 
     stderr.write('─'.repeat(leftWidth) + '┼' + '─'.repeat(rightWidth) + '\n');
 
     for (let row = 0; row < maxHeight; row++) {
       const result = results[row + this.listOffset];
       const isSelected = (row + this.listOffset === this.selectedIndex);
 
       let left;
       if (result) {
         let name = result.id;
         const maxName = leftWidth - 10;
         if (name.length > maxName) name = '...' + name.slice(-(maxName - 3));
         if (this.query && result.score !== undefined) {
           const score = result.score.toFixed(3);
           left = ' ' + name.padEnd(leftWidth - score.length - 3, ' ') + score + ' ';
         } else {
           left = ' ' + name;
         }
       } else {
         left = '';
       }
       left = left.padEnd(leftWidth, ' ').slice(0, leftWidth);
 
       let right = '';
       const isBatMode = this.useBat && selectedResult?.metadata?.path;
       if (selectedResult) {
         if (isBatMode) {
           if (!this._batOutput || this._batFileId !== selectedResult.id) {
             try {
               this._batOutput = execSync(`bat --color=always --style=numbers --terminal-width=${rightWidth} --wrap=character "${selectedResult.metadata.path}"`, {
                 maxBuffer: 1024 * 1024,
                 encoding: 'utf8',
               });
             } catch {
               this._batOutput = null;
             }
             this._batFileId = selectedResult.id;
           }
           if (this._batOutput) {
             const lines = this._batOutput.split('\n');
             const visibleLines = lines.slice(this.previewOffset, this.previewOffset + maxHeight);
             const line = visibleLines[row] || '';
             const stripped = stripAnsi(line);
             const padding = rightWidth - stripped.length;
             right = line + ' '.repeat(Math.max(0, padding)) + '\x1b[0m';
           }
         } else {
           if (!this._previewWrapped || this._previewFileId !== selectedResult.id || this._previewWidth !== rightWidth) {
             const lines = selectedResult.content.split('\n');
             this._previewWrapped = [];
             for (const line of lines) {
               if (line.length === 0) {
                 this._previewWrapped.push('');
               } else {
                 for (let i = 0; i < line.length; i += rightWidth) {
                   this._previewWrapped.push(line.slice(i, i + rightWidth));
                 }
               }
             }
             this._previewFileId = selectedResult.id;
             this._previewWidth = rightWidth;
           }
           right = this._previewWrapped[row + this.previewOffset] || '';
           right = right.padEnd(rightWidth, ' ').slice(0, rightWidth);
         }
       }
 
       if (isSelected) {
         stderr.write(`\x1b[7m${left}\x1b[0m│${right}`);
       } else {
         stderr.write(`${left}│${right}`);
       }
       stderr.write('\n');
     }
 
     stderr.write('─'.repeat(leftWidth) + '┴' + '─'.repeat(rightWidth) + '\n');
     stderr.write(`Query: ${this.query}`);
   }
 }
 
 function printUsage() {
   stdout.write(`Usage: vecfinder [options]
        vecfinder --query <text> [-k N]
 
 Options:
   -h, --help          Show this help message
   --query <text>      Search directory for best match (non-interactive)
   -k <number>         Number of results to return (default: 1, requires --query)
   --no-bat            Disable bat syntax highlighting (interactive mode)
 `);
 }
 
 // Main
 async function main() {
   const args = process.argv.slice(2);
   const queryArgIdx = args.indexOf('--query');
   const hasQueryFlag = queryArgIdx >= 0;
   const queryText = hasQueryFlag ? (args[queryArgIdx + 1] || null) : null;
   const kArgIdx = args.indexOf('-k');
   const k = kArgIdx >= 0 ? Math.max(1, parseInt(args[kArgIdx + 1], 10) || 1) : 1;
   const showHelp = args.includes('--help') || args.includes('-h');
 
   await loadConfig();
   await fs.mkdir(CACHE_DIR, { recursive: true });
 
   if (showHelp) {
     printUsage();
     process.exit(0);
   }
 
   const store = new DocumentStore();
 
   if (hasQueryFlag) {
     if (!queryText) {
       printUsage();
       stderr.write('Error: --query requires a text argument\n');
       process.exit(1);
     }
 
     await indexDirectory(store, process.cwd());
     const raw = await fetchEmbedding(queryText);
     store.queryVec = normalize(raw);
     const results = store.getResults().slice(0, k);
     for (const r of results) {
       stdout.write(r.id + '\n');
     }
     process.exit(0);
   }
 
   const useBatFlag = args.includes('--no-bat') ? false : null;
   const ui = new UI(store, useBatFlag);
 
   ui.init();
 
   indexDirectory(store, process.cwd());
 }
 
 main().catch(err => {
   console.error(`Fatal error: ${err.message}`);
   process.exit(1);
 });