draw-things-index

/**
 * Embedding Store using SQLite (sql.js WebAssembly)
 * 
 * Persists prompt embeddings to disk so they don't need to be
 * regenerated on every plugin restart.
 * 
 * Schema is simple: just prompt → embedding mapping with model info
 */

import initSqlJs, { Database } from 'sql.js';
import path from 'path';
import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'fs';
import { homedir } from 'os';

export interface EmbeddingStoreConfig {
  /** Path to the SQLite database file */
  dbPath?: string;
}

export interface StoredEmbedding {
  prompt: string;
  embedding: number[];
  model: string;
  createdAt: Date;
}

const DEFAULT_DB_PATH = path.join(
  homedir(),
  '.lmstudio',
  'extensions',
  'plugins',
  'ceveyne',
  'draw-things-index',
  'data',
  'embeddings.sqlite3'
);

export class EmbeddingStore {
  private db: Database | null = null;
  private config: Required<EmbeddingStoreConfig>;
  private SQL: any = null;
  private isDirty = false;

  constructor(config: EmbeddingStoreConfig = {}) {
    this.config = {
      dbPath: config.dbPath ?? DEFAULT_DB_PATH,
    };
  }

  /**
   * Initialize the database (load from disk or create new)
   */
  async init(): Promise<void> {
    if (this.db) return;

    // Initialize sql.js (WebAssembly SQLite)
    this.SQL = await initSqlJs();
    
    // Ensure directory exists
    const dir = path.dirname(this.config.dbPath);
    if (!existsSync(dir)) {
      mkdirSync(dir, { recursive: true });
    }

    // Load existing DB or create new
    if (existsSync(this.config.dbPath)) {
      console.log(`[EmbeddingStore] Loading existing DB from: ${this.config.dbPath}`);
      const buffer = readFileSync(this.config.dbPath);
      this.db = new this.SQL.Database(buffer);
    } else {
      console.log(`[EmbeddingStore] Creating new DB at: ${this.config.dbPath}`);
      this.db = new this.SQL.Database();
    }

    this.initSchema();
    
    const stats = this.getStats();
    console.log(`[EmbeddingStore] Loaded: ${stats.count} embeddings (model: ${stats.model || 'none'})`);
  }

  private initSchema(): void {
    if (!this.db) throw new Error('Database not initialized');
    
    this.db.run(`
      CREATE TABLE IF NOT EXISTS embeddings (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        prompt TEXT UNIQUE NOT NULL,
        embedding BLOB NOT NULL,
        model TEXT NOT NULL,
        created_at DATETIME DEFAULT CURRENT_TIMESTAMP
      )
    `);
    
    this.db.run(`CREATE INDEX IF NOT EXISTS idx_embeddings_prompt ON embeddings(prompt)`);
    this.db.run(`CREATE INDEX IF NOT EXISTS idx_embeddings_model ON embeddings(model)`);
    
    // Always save after schema init to ensure DB file exists
    this.forceSave();
  }

  /**
   * Save database to disk (only if dirty)
   */
  private save(): void {
    if (!this.db || !this.isDirty) return;
    this.forceSave();
  }

  /**
   * Force save database to disk (regardless of dirty flag)
   */
  private forceSave(): void {
    if (!this.db) return;
    const data = this.db.export();
    writeFileSync(this.config.dbPath, data);
    this.isDirty = false;
    console.log(`[EmbeddingStore] Saved to: ${this.config.dbPath}`);
  }

  /**
   * Get embedding for a prompt (if exists)
   */
  getEmbedding(prompt: string, model: string): number[] | null {
    if (!this.db) throw new Error('Database not initialized');
    
    const result = this.db.exec(
      `SELECT embedding FROM embeddings WHERE prompt = ? AND model = ?`,
      [prompt, model]
    );
    
    if (result.length === 0 || result[0].values.length === 0) {
      return null;
    }
    
    const blob = result[0].values[0][0] as Uint8Array;
    return this.blobToEmbedding(blob);
  }

  /**
   * Get all embeddings for a specific model
   * Returns a Map for efficient lookup
   */
  getAllEmbeddings(model: string): Map<string, number[]> {
    if (!this.db) throw new Error('Database not initialized');
    
    const result = this.db.exec(
      `SELECT prompt, embedding FROM embeddings WHERE model = ?`,
      [model]
    );
    
    const map = new Map<string, number[]>();
    
    if (result.length === 0) return map;
    
    for (const row of result[0].values) {
      const prompt = row[0] as string;
      const blob = row[1] as Uint8Array;
      map.set(prompt, this.blobToEmbedding(blob));
    }
    
    return map;
  }

  /**
   * Store embedding for a prompt
   */
  setEmbedding(prompt: string, embedding: number[], model: string): void {
    if (!this.db) throw new Error('Database not initialized');
    
    const blob = this.embeddingToBlob(embedding);
    
    // Upsert (INSERT OR REPLACE)
    this.db.run(
      `INSERT OR REPLACE INTO embeddings (prompt, embedding, model, created_at) 
       VALUES (?, ?, ?, CURRENT_TIMESTAMP)`,
      [prompt, blob, model]
    );
    
    this.isDirty = true;
  }

  /**
   * Store multiple embeddings at once (more efficient)
   */
  setEmbeddings(entries: Array<{ prompt: string; embedding: number[] }>, model: string): void {
    if (!this.db) throw new Error('Database not initialized');
    
    for (const { prompt, embedding } of entries) {
      const blob = this.embeddingToBlob(embedding);
      this.db.run(
        `INSERT OR REPLACE INTO embeddings (prompt, embedding, model, created_at) 
         VALUES (?, ?, ?, CURRENT_TIMESTAMP)`,
        [prompt, blob, model]
      );
    }
    
    this.isDirty = true;
    this.save(); // Save after batch
  }

  /**
   * Delete all embeddings for a specific model
   * (useful when model changes)
   */
  clearModel(model: string): void {
    if (!this.db) throw new Error('Database not initialized');
    
    this.db.run(`DELETE FROM embeddings WHERE model = ?`, [model]);
    this.isDirty = true;
    this.save();
    
    console.log(`[EmbeddingStore] Cleared all embeddings for model: ${model}`);
  }

  /**
   * Delete all embeddings
   */
  clearAll(): void {
    if (!this.db) throw new Error('Database not initialized');
    
    this.db.run(`DELETE FROM embeddings`);
    this.isDirty = true;
    this.save();
    
    console.log(`[EmbeddingStore] Cleared all embeddings`);
  }

  /**
   * Get statistics about stored embeddings
   */
  getStats(): { count: number; model: string | null; dimension: number | null; dbSizeBytes: number } {
    if (!this.db) return { count: 0, model: null, dimension: null, dbSizeBytes: 0 };
    
    const countResult = this.db.exec(`SELECT COUNT(*) FROM embeddings`);
    const count = countResult[0]?.values[0]?.[0] as number || 0;
    
    // Get most recent model
    const modelResult = this.db.exec(`SELECT model FROM embeddings ORDER BY created_at DESC LIMIT 1`);
    const model = modelResult[0]?.values[0]?.[0] as string || null;
    
    // Get dimension from first embedding
    let dimension: number | null = null;
    if (count > 0) {
      const embResult = this.db.exec(`SELECT embedding FROM embeddings LIMIT 1`);
      if (embResult.length > 0 && embResult[0].values.length > 0) {
        const blob = embResult[0].values[0][0] as Uint8Array;
        dimension = blob.length / 4; // Float32 = 4 bytes
      }
    }
    
    const dbSizeBytes = existsSync(this.config.dbPath) 
      ? require('fs').statSync(this.config.dbPath).size 
      : 0;
    
    return { count, model, dimension, dbSizeBytes };
  }

  /**
   * Flush any pending changes to disk
   */
  flush(): void {
    this.isDirty = true;
    this.save();
  }

  /**
   * Close database connection
   */
  close(): void {
    if (this.db) {
      this.save();
      this.db.close();
      this.db = null;
    }
  }

  // ─────────────────────────────────────────────────────────────
  // Blob conversion (Float32Array ↔ Uint8Array)
  // ─────────────────────────────────────────────────────────────

  private embeddingToBlob(embedding: number[]): Uint8Array {
    const float32 = new Float32Array(embedding);
    return new Uint8Array(float32.buffer);
  }

  private blobToEmbedding(blob: Uint8Array): number[] {
    const float32 = new Float32Array(
      blob.buffer,
      blob.byteOffset,
      blob.byteLength / 4
    );
    return Array.from(float32);
  }
}