hnswsqlite/src/cli/index.ts

Vector search with HNSWlib and SQLite in TypeScript.

#!/usr/bin/env node

import { Command } from 'commander';
import { addDocument } from './commands/add';
import { searchDocuments } from './commands/search';
import { deleteDocument, listDocuments } from './commands/delete';
import { VectorStore } from '../vectorStore';
import { initDatabase } from './utils/db';

const program = new Command();

program
  .name('hnswsqlite')
  .description('CLI for HNSWSQLite vector store')
  .version('0.1.0');

// Global options
program
  .option('-d, --database <path>', 'path to the SQLite database', 'vectors.db')
  .option('--dim <dimension>', 'dimension of the vectors', '1536')
  .option('--verbose', 'enable verbose output', false);

// Add command
program
  .command('add <text> [embedding...]')
  .description('Add a document with optional embedding')
  .action(async (text, embedding, options) => {
    const { database, dim } = options.parent.opts();
    const store = new VectorStore(database, parseInt(dim, 10));
    
    try {
      const docId = await addDocument(store, text, embedding);
      console.log(`✅ Document added with ID: ${docId}`);
    } catch (error) {
      if (error instanceof Error) {
        console.error('❌ Error adding document:', error.message);
      } else {
        console.error('❌ Error adding document:', error);
      }
      process.exit(1);
    } finally {
      store.close();
    }
  });

// Search command
program
  .command('search <query>')
  .description('Search for similar documents')
  .option('-k, --k <number>', 'number of results to return', '5')
  .action(async (query, options) => {
    const { database, dim } = options.parent.opts();
    const store = new VectorStore(database, parseInt(dim, 10));
    
    try {
      // In a real app, you'd generate an embedding from the query
      // For now, we'll use a dummy embedding
      const dummyEmbedding = new Array(parseInt(dim, 10)).fill(0.1);
      const results = await searchDocuments(store, dummyEmbedding, parseInt(options.k, 10));
      
      if (results.length === 0) {
        console.log('No matching documents found.');
        return;
      }
      
      console.log('\nSearch Results:');
      results.forEach((doc: {id: number, text: string}, index: number) => {
        console.log(`\n${index + 1}. ID: ${doc.id}`);
        console.log(`   Text: ${doc.text}`);
        console.log('   ---');
      });
    } catch (error) {
      if (error instanceof Error) {
        console.error('❌ Error searching documents:', error.message);
      } else {
        console.error('❌ Error searching documents:', error);
      }
      process.exit(1);
    } finally {
      store.close();
    }
  });

// Delete command
program
  .command('delete <id>')
  .description('Delete a document by ID')
  .action(async (id, options) => {
    const { database, dim } = options.parent.opts();
    const store = new VectorStore(database, parseInt(dim, 10));
    
    try {
      const success = await deleteDocument(store, parseInt(id, 10));
      if (success) {
        console.log(`✅ Document ${id} deleted successfully`);
      } else {
        console.log(`❌ Document ${id} not found or could not be deleted`);
      }
    } catch (error) {
      if (error instanceof Error) {
        console.error('❌ Error deleting document:', error.message);
      } else {
        console.error('❌ Error deleting document:', error);
      }
      process.exit(1);
    } finally {
      store.close();
    }
  });

// List command
program
  .command('list')
  .description('List all documents')
  .action(async (options) => {
    const { database, dim } = options.parent.opts();
    const store = new VectorStore(database, parseInt(dim, 10));
    
    try {
      const docs = await listDocuments(store);
      if (docs.length === 0) {
        console.log('No documents found.');
        return;
      }
      
      console.log('\nDocuments:');
      docs.forEach((doc: {id: number, text: string}, index: number) => {
        console.log(`\n${index + 1}. ID: ${doc.id}`);
        console.log(`   Text: ${doc.text}`);
        console.log('   ---');
      });
    } catch (error) {
      if (error instanceof Error) {
        console.error('❌ Error listing documents:', error.message);
      } else {
        console.error('❌ Error listing documents:', error);
      }
      process.exit(1);
    } finally {
      store.close();
    }
  });

// Initialize command
program
  .command('init')
  .description('Initialize a new database')
  .action((options) => {
    const { database, dim } = options.parent.opts();
    initDatabase(database, parseInt(dim, 10));
    console.log(`✅ Database initialized at ${database}`);
  });

// Handle unknown commands
program.on('command:*', () => {
  console.error('Invalid command. Use --help for a list of commands.');
  process.exit(1);
});

// Parse command line arguments
program.parse(process.argv);

// Show help if no arguments
if (!process.argv.slice(2).length) {
  program.outputHelp();
}