OmniRoute/src/lib/memory/vectorStore.ts

// Raw SQL allowed: sqlite-vec virtual table DDL is dynamic (dim varies). See plan 21 §D5.
// Hard Rule #5 exception: sqlite-vec VIRTUAL TABLE cannot be created via src/lib/db/ domain modules
// because the table dimension (N in FLOAT[N]) depends on the active embedding model at runtime.
//
// NOTE on rowid: vec0 v0.1.9 requires BigInt when inserting explicit rowid values.
// The vec_memories table uses the *same* rowid space as the `memories` table to enable
// a simple JOIN (m.rowid = v.rowid). We do NOT use a named primary-key column because
// vec0 rejects numeric (non-BigInt) values for named PKs in this version.

import { createRequire } from "module";
import type { EmbeddingResolution } from "./embedding/types";
import {
  getMemoryVecMeta,
  setMemoryVecMeta,
  markAllMemoriesNeedReindex,
  countMemoryReindexPending,
} from "../localDb";
import { getDbInstance } from "../db/core";
import { logger } from "../../../open-sse/utils/logger.ts";
import { sanitizeErrorMessage } from "../../../open-sse/utils/error.ts";

const _require = createRequire(import.meta.url);

const log = logger("VECTOR_STORE");

// ──────────────── Types ────────────────

export interface VectorSearchHit {
  memoryId: string; // UUID (same as memories.id)
  distance: number; // L2 distance — lower = more similar
  score: number; // 1 / (1 + distance) — higher = better
}

export interface HybridRrfHit {
  memoryId: string;
  vecRank: number | null; // null if not from vector search
  ftsRank: number | null; // null if not from FTS5
  rrfScore: number; // RRF score (k=60 default)
  vecDistance: number | null;
  ftsScore: number | null;
}

export interface VectorStore {
  /** Ensure schema (sqlite-vec loaded, vec_memories created if needed, dim aligned). Idempotent. */
  ensureReady(resolution: EmbeddingResolution): Promise<{ ready: boolean; reason: string }>;
  /** Insert/update vector for a memory. */
  upsertVector(memoryId: string, vector: Float32Array): Promise<void>;
  /** Delete vector for a memory (no-op if not present). */
  deleteVector(memoryId: string): Promise<void>;
  /** KNN brute-force search. Returns top-K hits ordered by distance ASC. */
  searchVector(vector: Float32Array, topK: number, apiKeyId?: string): Promise<VectorSearchHit[]>;
  /** Hybrid RRF search (FTS5 + vector fused via Reciprocal Rank Fusion, k=60). */
  searchHybrid(
    vector: Float32Array,
    queryText: string,
    topK: number,
    apiKeyId?: string,
  ): Promise<HybridRrfHit[]>;
  /** Stats for UI Engine status. */
  stats(): Promise<{
    rowCount: number;
    needsReindex: number;
    activeDim: number | null;
    signature: string | null;
  }>;
  /** Drop and recreate vec_memories (on signature change). Marks all memories needs_reindex=1. */
  resetForSignature(signature: string, dim: number): Promise<void>;
}

// ──────────────── Constants ────────────────

const RRF_K = Number(process.env["MEMORY_RRF_K"] ?? 60);
const TOP_K_DEFAULT = Number(process.env["MEMORY_VEC_TOP_K"] ?? 20);

// ──────────────── Helpers ────────────────

/**
 * Encode a Float32Array as a Buffer of little-endian bytes.
 * sqlite-vec accepts this format for FLOAT[] column values.
 */
function encodeVector(v: Float32Array): Buffer {
  return Buffer.from(v.buffer, v.byteOffset, v.byteLength);
}

// ──────────────── Implementation ────────────────

class VectorStoreImpl implements VectorStore {
  async ensureReady(resolution: EmbeddingResolution): Promise<{ ready: boolean; reason: string }> {
    const db = getDbInstance();
    const meta = getMemoryVecMeta();

    // Signature changed (or first time with a known dim) → recreate with new dim.
    if (resolution.dimensions !== null && resolution.signature !== meta.embeddingSignature) {
      await this.resetForSignature(resolution.signature, resolution.dimensions);
      return { ready: true, reason: `vec_memories recreated with dim=${resolution.dimensions}` };
    }

    // Already marked loaded → idempotent no-op.
    if (meta.vecLoaded) {
      return { ready: true, reason: "vec_memories already ready" };
    }

    // Not yet loaded but we have a dim — create the table now.
    if (resolution.dimensions !== null) {
      const dim = meta.activeDim ?? resolution.dimensions;
      try {
        db.exec(
          `CREATE VIRTUAL TABLE IF NOT EXISTS vec_memories USING vec0(embedding FLOAT[${dim}])`,
        );
        setMemoryVecMeta({ vecLoaded: true, activeDim: dim });
        return { ready: true, reason: `vec_memories created with dim=${dim}` };
      } catch (err: unknown) {
        const msg = sanitizeErrorMessage(err instanceof Error ? err.message : String(err));
        return { ready: false, reason: `failed to create vec_memories: ${msg}` };
      }
    }

    return { ready: false, reason: "no dimensions available yet (lazy probe pending)" };
  }

  async upsertVector(memoryId: string, vector: Float32Array): Promise<void> {
    const db = getDbInstance();

    // Map UUID memoryId → INTEGER rowid (the rowid is used as the FK into vec_memories).
    const row = db.prepare("SELECT rowid FROM memories WHERE id = ?").get(memoryId) as
      | { rowid: number }
      | undefined;

    if (!row) {
      throw new Error(`memory not found: ${memoryId}`);
    }

    // vec0 v0.1.9 requires BigInt for explicit rowid insertion — plain numbers are rejected.
    // INSERT OR REPLACE is not supported by vec0 — use DELETE + INSERT for upsert semantics.
    db.prepare("DELETE FROM vec_memories WHERE rowid = ?").run(BigInt(row.rowid));
    db.prepare("INSERT INTO vec_memories(rowid, embedding) VALUES (?, ?)").run(
      BigInt(row.rowid),
      encodeVector(vector),
    );
  }

  async deleteVector(memoryId: string): Promise<void> {
    const db = getDbInstance();
    db.prepare(
      "DELETE FROM vec_memories WHERE rowid = (SELECT rowid FROM memories WHERE id = ?)",
    ).run(memoryId);
  }

  async searchVector(
    vector: Float32Array,
    topK: number,
    apiKeyId?: string,
  ): Promise<VectorSearchHit[]> {
    const db = getDbInstance();
    const k = topK > 0 ? topK : TOP_K_DEFAULT;

    const rows = db
      .prepare(
        `SELECT m.id AS memory_id, v.distance
         FROM vec_memories v
         JOIN memories m ON m.rowid = v.rowid
         WHERE v.embedding MATCH ?
           AND ($apiKeyId IS NULL OR m.api_key_id = $apiKeyId)
           AND k = ?
         ORDER BY v.distance ASC`,
      )
      .all(encodeVector(vector), { apiKeyId: apiKeyId ?? null }, k) as Array<{
        memory_id: string;
        distance: number;
      }>;

    return rows.map((r) => ({
      memoryId: r.memory_id,
      distance: r.distance,
      score: 1 / (1 + r.distance),
    }));
  }

  async searchHybrid(
    vector: Float32Array,
    queryText: string,
    topK: number,
    apiKeyId?: string,
  ): Promise<HybridRrfHit[]> {
    const db = getDbInstance();
    const k = topK > 0 ? topK : TOP_K_DEFAULT;
    const rrfK = RRF_K;

    // SQLite does not support FULL OUTER JOIN — use UNION ALL + GROUP BY (RRF recipe).
    // Reference: https://alexgarcia.xyz/blog/2024/sqlite-vec-hybrid-search/
    const rows = db
      .prepare(
        `WITH vec_results AS (
           SELECT m.id AS memory_id,
                  ROW_NUMBER() OVER (ORDER BY v.distance ASC) AS vec_rank,
                  v.distance AS vec_distance
           FROM vec_memories v
           JOIN memories m ON m.rowid = v.rowid
           WHERE v.embedding MATCH ?
             AND ($apiKeyId IS NULL OR m.api_key_id = $apiKeyId)
             AND k = ?
         ),
         fts_results AS (
           SELECT m.id AS memory_id,
                  ROW_NUMBER() OVER (ORDER BY fts.rank ASC) AS fts_rank,
                  fts.rank AS fts_score
           FROM memory_fts fts
           JOIN memories m ON m.memory_id = fts.rowid
           WHERE fts.memory_fts MATCH ?
             AND ($apiKeyId IS NULL OR m.api_key_id = $apiKeyId)
           LIMIT ?
         ),
         fused AS (
           SELECT
             memory_id,
             MAX(vec_rank)      AS vec_rank,
             MAX(fts_rank)      AS fts_rank,
             MAX(vec_distance)  AS vec_distance,
             MAX(fts_score)     AS fts_score,
             SUM(rrf_contrib)   AS rrf_score
           FROM (
             SELECT memory_id, vec_rank, NULL AS fts_rank, vec_distance,
                    NULL AS fts_score, 1.0 / (${rrfK} + vec_rank) AS rrf_contrib
             FROM vec_results
             UNION ALL
             SELECT memory_id, NULL, fts_rank, NULL, fts_score, 1.0 / (${rrfK} + fts_rank)
             FROM fts_results
           )
           GROUP BY memory_id
         )
         SELECT memory_id, vec_rank, fts_rank, vec_distance, fts_score, rrf_score
         FROM fused
         ORDER BY rrf_score DESC
         LIMIT ?`,
      )
      .all(
        encodeVector(vector),
        { apiKeyId: apiKeyId ?? null },
        k,
        queryText,
        k,
        k,
      ) as Array<{
        memory_id: string;
        vec_rank: number | null;
        fts_rank: number | null;
        vec_distance: number | null;
        fts_score: number | null;
        rrf_score: number;
      }>;

    return rows.map((r) => ({
      memoryId: r.memory_id,
      vecRank: r.vec_rank,
      ftsRank: r.fts_rank,
      rrfScore: r.rrf_score,
      vecDistance: r.vec_distance,
      ftsScore: r.fts_score,
    }));
  }

  async stats(): Promise<{
    rowCount: number;
    needsReindex: number;
    activeDim: number | null;
    signature: string | null;
  }> {
    let rowCount = 0;
    try {
      const db = getDbInstance();
      const row = db.prepare("SELECT COUNT(*) AS cnt FROM vec_memories").get() as
        | { cnt: number }
        | undefined;
      rowCount = row?.cnt ?? 0;
    } catch {
      // vec_memories may not exist yet — not an error, just 0 rows.
      rowCount = 0;
    }

    const needsReindex = countMemoryReindexPending();
    const meta = getMemoryVecMeta();

    return {
      rowCount,
      needsReindex,
      activeDim: meta.activeDim,
      signature: meta.embeddingSignature,
    };
  }

  async resetForSignature(signature: string, dim: number): Promise<void> {
    const db = getDbInstance();

    // DROP + CREATE is intentionally destructive — triggers lazy backfill via F5.
    db.exec("DROP TABLE IF EXISTS vec_memories");
    db.exec(`CREATE VIRTUAL TABLE vec_memories USING vec0(embedding FLOAT[${dim}])`);

    markAllMemoriesNeedReindex();
    setMemoryVecMeta({
      activeDim: dim,
      embeddingSignature: signature,
      lastResetAt: new Date().toISOString(),
      vecLoaded: true,
    });
  }
}

// ──────────────── Singleton ────────────────

let _instance: VectorStore | null | undefined = undefined; // undefined = not yet attempted

/**
 * Singleton instance (lazy-initialized).
 * Returns null if sqlite-vec is unavailable (e.g. WASM / cloud backend).
 * Callers should degrade gracefully to FTS5 keyword search when this returns null.
 */
export function getVectorStore(): VectorStore | null {
  if (_instance !== undefined) {
    return _instance;
  }

  // Test seam: VECTOR_STORE_DISABLE_VEC=true forces null (simulates cloud/WASM environment).
  if (process.env["VECTOR_STORE_DISABLE_VEC"] === "true") {
    log.warn(
      "VECTOR_STORE_DISABLE_VEC is set — sqlite-vec disabled. Degrading to FTS5 keyword search.",
    );
    _instance = null;
    return null;
  }

  const db = getDbInstance();
  const raw = db.raw as { loadExtension?: (path: string) => void } | null;

  // sqlite-vec must be loaded as a native extension on the better-sqlite3 raw handle.
  // The SqliteAdapter wrapper does not expose loadExtension directly.
  if (!raw || typeof raw.loadExtension !== "function") {
    log.warn(
      "sqlite-vec not loaded: db driver does not support loadExtension (cloud/WASM backend). " +
        "Degrading to FTS5 keyword search.",
    );
    _instance = null;
    return null;
  }

  try {
    const sqliteVec = _require("sqlite-vec") as { load: (db: unknown) => void };
    sqliteVec.load(raw);
    log.info("sqlite-vec loaded successfully");
    _instance = new VectorStoreImpl();
  } catch (err: unknown) {
    const safeMsg = sanitizeErrorMessage(err instanceof Error ? err.message : String(err));
    log.warn(`sqlite-vec failed to load: ${safeMsg}. Degrading to FTS5 keyword search.`);
    _instance = null;
  }

  return _instance;
}

/**
 * Reset the singleton cache (for tests only — allows re-initialization between tests).
 * @internal
 */
export function _resetVectorStoreSingleton(): void {
  _instance = undefined;
}