hnswsqlite/src/db.ts

Vector search with HNSWlib and SQLite in TypeScript.

import Database from 'better-sqlite3';

export class SQLiteDB {
  private db: Database.Database;

  constructor(dbPath: string) {
    this.db = new Database(dbPath);
    this._initTables();
  }

  private _initTables() {
    this.db.prepare(`CREATE TABLE IF NOT EXISTS documents (
      id INTEGER PRIMARY KEY AUTOINCREMENT,
      text TEXT,
      embedding BLOB
    )`).run();
  }

  addDocument(text: string, embedding: Buffer): number {
    const result = this.db.prepare('INSERT INTO documents (text, embedding) VALUES (?, ?)').run(text, embedding);
    return Number(result.lastInsertRowid);
  }

  getDocument(id: number): {id: number, text: string, embedding: Buffer} | undefined {
    return this.db.prepare('SELECT * FROM documents WHERE id = ?').get(id) as {id: number, text: string, embedding: Buffer} | undefined;
  }

  getAllEmbeddings(): Array<{id: number, embedding: Buffer}> {
    return this.db.prepare('SELECT id, embedding FROM documents').all() as Array<{id: number, embedding: Buffer}>;
  }

  close() {
    this.db.close();
  }
}