openprompt-lang 1.2.7 → 1.4.0

package/src/embeddings/vector-store.js

@@ -0,0 +1,505 @@
+// @use(kind, contract, limit, deps, pattern)
+// @kind(module)
+// @contract(in: ChunkWithVector -> out: void, sideEffect: SQLite writes)
+// @limit(lines: 380)
+// @deps(better-sqlite3, ../persistence/sqlite/connection)
+// @pattern(repository)
+
+/**
+ * Almacenamiento vectorial en SQLite con búsqueda por similitud de coseno.
+ *
+ * Para ~5000 chunks, la búsqueda lineal por coseno es suficiente.
+ * Si el volumen supera 100k chunks, considerar índice HNSW.
+ *
+ * Serialización de vectores:
+ *   - Almacenamiento: Float32Array → Buffer (4 bytes por float)
+ *   - Recuperación: Buffer → Float32Array
+ */
+
+import { open } from "../persistence/sqlite/connection.js"
+import { ensureSchema } from "../persistence/sqlite/schema.js"
+
+// ─── Constantes ────────────────────────────────────────────────────
+
+const DEFAULT_DB_PATH = "./.opencode/opl.db"
+const DEFAULT_TOP_K = 10
+const DEFAULT_MIN_SCORE = 0.0
+
+// ─── Schema ────────────────────────────────────────────────────────
+
+const VECTOR_STORE_SQL = `
+CREATE TABLE IF NOT EXISTS embeddings (
+    id          TEXT PRIMARY KEY,
+    doc_id      TEXT NOT NULL,
+    doc_title   TEXT NOT NULL DEFAULT '',
+    chapter_idx INTEGER NOT NULL DEFAULT 0,
+    chapter_title TEXT NOT NULL DEFAULT '',
+    chunk_idx   INTEGER NOT NULL DEFAULT 0,
+    content     TEXT NOT NULL,
+    vector      BLOB NOT NULL,
+    dimension   INTEGER NOT NULL DEFAULT 768,
+    tokens      INTEGER NOT NULL DEFAULT 0,
+    model       TEXT NOT NULL DEFAULT '',
+    strategy    TEXT NOT NULL DEFAULT 'section',
+    metadata    TEXT NOT NULL DEFAULT '{}',
+    created_at  TEXT NOT NULL DEFAULT (datetime('now'))
+);
+
+CREATE INDEX IF NOT EXISTS idx_embeddings_doc_id ON embeddings(doc_id);
+CREATE INDEX IF NOT EXISTS idx_embeddings_created ON embeddings(created_at);
+
+CREATE VIRTUAL TABLE IF NOT EXISTS embeddings_fts USING fts5(
+    content,
+    content=embeddings,
+    content_rowid=rowid
+);
+`
+
+// ─── Serialización de vectores ─────────────────────────────────────
+
+/**
+ * Convierte un array de números a Buffer para almacenamiento SQLite.
+ * Usa Float32 (4 bytes por float).
+ *
+ * @param {number[]} vector
+ * @returns {Buffer}
+ */
+function vectorToBuffer(vector) {
+  const float32 = new Float32Array(vector)
+  return Buffer.from(float32.buffer)
+}
+
+/**
+ * Convierte un Buffer SQLite a array de números.
+ *
+ * @param {Buffer} buffer
+ * @returns {number[]}
+ */
+function bufferToVector(buffer) {
+  const float32 = new Float32Array(buffer.buffer, buffer.byteOffset, buffer.byteLength / 4)
+  return Array.from(float32)
+}
+
+// ─── Coso similitud (re-exportada aquí para conveniencia) ──────────
+
+/**
+ * Calcula la similitud de coseno entre dos vectores.
+ *
+ * @param {number[]} a
+ * @param {number[]} b
+ * @returns {number}
+ */
+export function cosineSimilarity(a, b) {
+  if (a.length !== b.length) {
+    throw new Error(`Dimensión incorrecta: vector A tiene ${a.length}, vector B tiene ${b.length}.`)
+  }
+  if (a.length === 0) return 0
+
+  let dotProduct = 0
+  let normA = 0
+  let normB = 0
+
+  for (let i = 0; i < a.length; i++) {
+    dotProduct += a[i] * b[i]
+    normA += a[i] * a[i]
+    normB += b[i] * b[i]
+  }
+
+  const magnitude = Math.sqrt(normA) * Math.sqrt(normB)
+  if (magnitude === 0) return 0
+  return dotProduct / magnitude
+}
+
+// ─── Helpers de BD ─────────────────────────────────────────────────
+
+/**
+ * Obtiene o crea la conexión a la BD de vectores.
+ *
+ * @param {string} dbPath
+ * @returns {{ db: import('better-sqlite3').Database }}
+ */
+function getDb(dbPath) {
+  const conn = open(dbPath || DEFAULT_DB_PATH)
+  ensureSchema(conn.db)
+  conn.db.exec(VECTOR_STORE_SQL)
+  return { db: conn.db }
+}
+
+/**
+ * Serializa metadatos de chunk para almacenar en metadata JSON.
+ *
+ * @param {Object} chunk
+ * @returns {string}
+ */
+function serializeMetadata(chunk) {
+  return JSON.stringify(chunk.metadata || {})
+}
+
+// ─── API Pública ───────────────────────────────────────────────────
+
+/**
+ * Guarda un chunk con su vector embedding en SQLite.
+ *
+ * @param {Object} chunk - Chunk del documento (de chunker.js)
+ * @param {number[]} vector - Vector embedding
+ * @param {string} model - Modelo usado para el embedding
+ * @param {Object} [options]
+ * @param {string} [options.dbPath]
+ * @returns {{ id: string, success: boolean }}
+ */
+export function storeEmbedding(chunk, vector, model, options = {}) {
+  const { db } = getDb(options.dbPath)
+
+  if (!chunk || !chunk.id) {
+    throw new Error("Chunk inválido: se requiere un objeto con id")
+  }
+  if (!vector || vector.length === 0) {
+    throw new Error("Vector inválido: se requiere un array no vacío")
+  }
+
+  const vectorBuffer = vectorToBuffer(vector)
+  const metadata = serializeMetadata(chunk)
+
+  const insert = db.transaction(() => {
+    const stmt = db.prepare(`
+      INSERT OR REPLACE INTO embeddings
+        (id, doc_id, doc_title, chapter_idx, chapter_title, chunk_idx,
+         content, vector, dimension, tokens, model, strategy, metadata)
+      VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+    `)
+
+    stmt.run(
+      chunk.id,
+      chunk.docId || "",
+      chunk.docTitle || "",
+      chunk.chapterIdx || 0,
+      chunk.chapterTitle || "",
+      chunk.chunkIndex || 0,
+      chunk.content || "",
+      vectorBuffer,
+      vector.length,
+      chunk.tokens || 0,
+      model || "",
+      chunk.strategy || "section",
+      metadata
+    )
+
+    // Sincronizar FTS5
+    try {
+      const row = db.prepare(`SELECT rowid FROM embeddings WHERE id = ?`).get(chunk.id)
+      if (row) {
+        db.prepare(`INSERT INTO embeddings_fts(rowid, content) VALUES (?, ?)`).run(
+          row.rowid,
+          chunk.content || ""
+        )
+      }
+    } catch {
+      db.exec(`INSERT INTO embeddings_fts(embeddings_fts) VALUES('rebuild')`)
+    }
+  })
+
+  insert()
+
+  return { id: chunk.id, success: true }
+}
+
+/**
+ * Guarda múltiples embeddings en una transacción SQLite.
+ *
+ * @param {Array<{ chunk: Object, vector: number[], model: string }>} entries
+ * @param {Object} [options]
+ * @param {string} [options.dbPath]
+ * @returns {{ count: number, durationMs: number }}
+ */
+export function storeEmbeddingsBatch(entries, options = {}) {
+  const { db } = getDb(options.dbPath)
+
+  if (!Array.isArray(entries) || entries.length === 0) {
+    return { count: 0, durationMs: 0 }
+  }
+
+  const start = Date.now()
+
+  const stmt = db.prepare(`
+    INSERT OR REPLACE INTO embeddings
+      (id, doc_id, doc_title, chapter_idx, chapter_title, chunk_idx,
+       content, vector, dimension, tokens, model, strategy, metadata)
+    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+  `)
+
+  const insertMany = db.transaction((items) => {
+    for (const item of items) {
+      const { chunk, vector, model } = item
+      const vectorBuffer = vectorToBuffer(vector)
+
+      stmt.run(
+        chunk.id,
+        chunk.docId || "",
+        chunk.docTitle || "",
+        chunk.chapterIdx || 0,
+        chunk.chapterTitle || "",
+        chunk.chunkIndex || 0,
+        chunk.content || "",
+        vectorBuffer,
+        vector.length,
+        chunk.tokens || 0,
+        model || "",
+        chunk.strategy || "section",
+        serializeMetadata(chunk)
+      )
+    }
+  })
+
+  insertMany(entries)
+
+  // Reconstruir índice FTS5
+  try {
+    db.exec(`INSERT INTO embeddings_fts(embeddings_fts) VALUES('rebuild')`)
+  } catch {
+    // FTS5 puede no estar disponible, no es crítico
+  }
+
+  const durationMs = Date.now() - start
+  return { count: entries.length, durationMs }
+}
+
+/**
+ * Busca los top-K chunks más similares por similitud de coseno.
+ *
+ * Estrategia: carga todos los vectores y hace búsqueda lineal.
+ * Para ~5000 chunks, esto es < 50ms en Node.js.
+ *
+ * @param {number[]} queryVector - Vector de la consulta
+ * @param {Object} [options]
+ * @param {number} [options.topK=10] - Máximo de resultados
+ * @param {number} [options.minScore=0.0] - Score mínimo para incluir
+ * @param {string} [options.docFilter] - Filtrar por doc_id
+ * @param {string} [options.dbPath]
+ * @returns {Array<{
+ *   id: string, docId: string, docTitle: string,
+ *   chapterIdx: number, chapterTitle: string,
+ *   content: string, score: number,
+ *   tokens: number, model: string
+ * }>}
+ */
+export function searchSimilar(queryVector, options = {}) {
+  const { db } = getDb(options.dbPath)
+  const topK = options.topK || DEFAULT_TOP_K
+  const minScore = options.minScore ?? DEFAULT_MIN_SCORE
+  const docFilter = options.docFilter || null
+
+  if (!queryVector || queryVector.length === 0) {
+    return []
+  }
+
+  // Construir query base
+  let sql = `SELECT id, doc_id, doc_title, chapter_idx, chapter_title, 
+                    content, vector, dimension, tokens, model, metadata
+             FROM embeddings`
+  const params = []
+
+  if (docFilter) {
+    sql += ` WHERE doc_id = ?`
+    params.push(docFilter)
+  }
+
+  sql += ` ORDER BY created_at DESC`
+
+  const rows = db.prepare(sql).all(...params)
+
+  if (rows.length === 0) return []
+
+  // Calcular similitud de coseno para cada fila
+  const scored = []
+  for (const row of rows) {
+    const vector = bufferToVector(row.vector)
+    const score = cosineSimilarity(queryVector, vector)
+
+    if (score >= minScore) {
+      scored.push({
+        id: row.id,
+        docId: row.doc_id,
+        docTitle: row.doc_title,
+        chapterIdx: row.chapter_idx,
+        chapterTitle: row.chapter_title,
+        content: row.content,
+        score,
+        tokens: row.tokens,
+        model: row.model,
+        metadata: row.metadata,
+      })
+    }
+  }
+
+  // Ordenar por score descendente y tomar top-K
+  scored.sort((a, b) => b.score - a.score)
+  return scored.slice(0, topK)
+}
+
+/**
+ * Busca por texto primero (FTS5) y reordena por similitud de coseno.
+ *
+ * 1. Busca en FTS5 por palabras clave → candidatos
+ * 2. Calcula similitud de coseno para cada candidato
+ * 3. Reordena por score
+ *
+ * @param {string} queryText - Texto para búsqueda FTS5
+ * @param {number[]} queryVector - Vector para reordenamiento semántico
+ * @param {Object} [options]
+ * @param {number} [options.topK=10]
+ * @param {string} [options.dbPath]
+ * @returns {Promise<Array>}
+ */
+export function hybridSearch(queryText, queryVector, options = {}) {
+  const { db } = getDb(options.dbPath)
+  const topK = options.topK || DEFAULT_TOP_K
+
+  if (!queryText || queryText.trim().length === 0) {
+    return searchSimilar(queryVector, options)
+  }
+
+  try {
+    // Buscar en FTS5
+    const ftsQuery = queryText
+      .split(/\s+/)
+      .filter((w) => w.length > 2)
+      .map((w) => `${w}*`)
+      .join(" OR ")
+
+    if (!ftsQuery) return searchSimilar(queryVector, options)
+
+    const ftsResults = db
+      .prepare(
+        `SELECT rank, rowid FROM embeddings_fts WHERE embeddings_fts MATCH ? ORDER BY rank LIMIT ?`
+      )
+      .all(ftsQuery, topK * 3)
+
+    if (ftsResults.length === 0) {
+      return searchSimilar(queryVector, options)
+    }
+
+    // Obtener vectores completos de los candidatos FTS (por rowid)
+    const rowids = ftsResults.map((r) => r.rowid)
+    const placeholders = rowids.map(() => "?").join(",")
+
+    const rows = db
+      .prepare(
+        `SELECT id, doc_id, doc_title, chapter_idx, chapter_title,
+                content, vector, dimension, tokens, model
+         FROM embeddings WHERE rowid IN (${placeholders})`
+      )
+      .all(...rowids)
+
+    if (rows.length === 0) return []
+
+    // Reordenar por similitud de coseno
+    const scored = []
+    for (const row of rows) {
+      const vector = bufferToVector(row.vector)
+      const score = cosineSimilarity(queryVector, vector)
+      scored.push({
+        id: row.id,
+        docId: row.doc_id,
+        docTitle: row.doc_title,
+        chapterIdx: row.chapter_idx,
+        chapterTitle: row.chapter_title,
+        content: row.content,
+        score,
+        tokens: row.tokens,
+        model: row.model,
+      })
+    }
+
+    scored.sort((a, b) => b.score - a.score)
+    return scored.slice(0, topK)
+  } catch {
+    // Si FTS5 falla, caer a búsqueda por coseno puro
+    return searchSimilar(queryVector, options)
+  }
+}
+
+/**
+ * Elimina todos los embeddings de un documento específico.
+ *
+ * @param {string} docId - ID del documento a eliminar
+ * @param {Object} [options]
+ * @param {string} [options.dbPath]
+ * @returns {{ deleted: number }}
+ */
+export function deleteDocumentEmbeddings(docId, options = {}) {
+  const { db } = getDb(options.dbPath)
+
+  if (!docId) {
+    throw new Error("docId es requerido para eliminar embeddings")
+  }
+
+  // Eliminar primero del FTS
+  try {
+    db.prepare(
+      `DELETE FROM embeddings_fts WHERE rowid IN (SELECT rowid FROM embeddings WHERE doc_id = ?)`
+    ).run(docId)
+  } catch {
+    // FTS puede no existir
+  }
+
+  const result = db.prepare(`DELETE FROM embeddings WHERE doc_id = ?`).run(docId)
+
+  return { deleted: result.changes }
+}
+
+/**
+ * Obtiene estadísticas del índice de embeddings.
+ *
+ * @param {Object} [options]
+ * @param {string} [options.dbPath]
+ * @returns {{
+ *   totalEmbeddings: number,
+ *   totalDocs: number,
+ *   dimension: number,
+ *   model: string,
+ *   lastIndexed: string|null,
+ *   storageBytes: number
+ * }}
+ */
+export function getEmbeddingStats(options = {}) {
+  const { db } = getDb(options.dbPath)
+
+  const countResult = db.prepare(`SELECT COUNT(*) as count FROM embeddings`).get()
+  const totalEmbeddings = countResult.count
+
+  if (totalEmbeddings === 0) {
+    return {
+      totalEmbeddings: 0,
+      totalDocs: 0,
+      dimension: 0,
+      model: "",
+      lastIndexed: null,
+      storageBytes: 0,
+    }
+  }
+
+  const docsResult = db.prepare(`SELECT COUNT(DISTINCT doc_id) as count FROM embeddings`).get()
+  const lastResult = db.prepare(`SELECT MAX(created_at) as last FROM embeddings`).get()
+  const modelResult = db
+    .prepare(
+      `SELECT model, COUNT(*) as count FROM embeddings GROUP BY model ORDER BY count DESC LIMIT 1`
+    )
+    .get()
+  const dimResult = db
+    .prepare(
+      `SELECT dimension, COUNT(*) as count FROM embeddings GROUP BY dimension ORDER BY count DESC LIMIT 1`
+    )
+    .get()
+
+  // Estimar tamaño en bytes
+  const vectorSize = db.prepare(`SELECT SUM(LENGTH(vector)) as total FROM embeddings`).get()
+
+  return {
+    totalEmbeddings,
+    totalDocs: docsResult.count,
+    dimension: dimResult ? dimResult.dimension : 0,
+    model: modelResult ? modelResult.model : "",
+    lastIndexed: lastResult ? lastResult.last : null,
+    storageBytes: vectorSize ? vectorSize.total + totalEmbeddings * 512 : 0, // 512 aprox por metadata
+  }
+}