Hugging Face's logo

Spaces:
automindy
/
fiemwl
Build error

App Files Community
fiemwl / src /lib /memory /vectorStore.ts
automindy's picture
automindy
Upload 1980 files
6111b2b verified
Raw History Blame Contribute Delete
13.3 kB
// 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;
}