// Semantic search index: each model's text (name, description, notes, tags, file // names) embedded locally, refreshed in the background whenever that text changes. const db = require('./db'); const embeddingModel = require('./embeddingModel'); const BATCH = 16; const MAX_CHARS = 2000; // the model reads ~512 tokens; the start of the text matters most const RETRY_AFTER_MS = 10 * 60 * 1000; // Part of every text fingerprint, so switching embedding model re-indexes everything. const SALT = `${embeddingModel.MODEL}:1:`; let lastError = null; let retryAt = 0; let vectors = null; // [{ id, vec: Float32Array }], rebuilt after the index changes let vectorsLoadedAt = 0; const VECTORS_MAX_AGE_MS = 60 * 1000; // so trashed and restored models catch up const DOCS = ` SELECT m.id, m.embedding_key, concat_ws(E'\\n', m.name, NULLIF(btrim(m.description), ''), (SELECT string_agg(tg.name, ', ' ORDER BY tg.name) FROM model_tags mt JOIN tags tg ON tg.id = mt.tag_id WHERE mt.model_id = m.id), NULLIF(btrim(m.notes), ''), m.source_site, d.name, (SELECT string_agg(f.original_name, ', ' ORDER BY f.id) FROM files f WHERE f.model_id = m.id), (SELECT string_agg(e.path, ', ' ORDER BY e.id) FROM archive_entries e JOIN files f ON f.id = e.file_id WHERE f.model_id = m.id) ) AS text FROM models m LEFT JOIN designers d ON d.id = m.designer_id WHERE m.deleted_at IS NULL`; const PENDING = `FROM (${DOCS}) doc WHERE doc.embedding_key IS DISTINCT FROM md5($1 || doc.text)`; async function run() { if (Date.now() < retryAt) return false; const { rows } = await db.query(`SELECT doc.id, doc.text ${PENDING} ORDER BY doc.id LIMIT $2`, [SALT, BATCH]); if (!rows.length) return false; let vecs; try { vecs = await embeddingModel.embed(rows.map((r) => r.text.slice(0, MAX_CHARS))); lastError = null; } catch (err) { lastError = err.message; retryAt = Date.now() + RETRY_AFTER_MS; console.warn('Semantic search indexing paused:', err.message); return false; } for (let i = 0; i < rows.length; i++) { await db.query('UPDATE models SET embedding = $2, embedding_key = md5($3 || $4) WHERE id = $1', [rows[i].id, vecs[i], SALT, rows[i].text]); } vectors = null; return true; } async function loadVectors() { if (!vectors || Date.now() - vectorsLoadedAt > VECTORS_MAX_AGE_MS) { const { rows } = await db.query('SELECT id, embedding FROM models WHERE embedding IS NOT NULL AND deleted_at IS NULL'); vectors = rows.map((r) => ({ id: r.id, vec: Float32Array.from(r.embedding) })); vectorsLoadedAt = Date.now(); } return vectors; } // Models ranked by how close their meaning is to the query: [{ id, score }], // score = cosine similarity (vectors are unit length). async function search(query, limit = 200) { const [q] = await embeddingModel.embed([embeddingModel.QUERY_PREFIX + String(query).slice(0, 300)]); const list = await loadVectors(); return list .map(({ id, vec }) => { let score = 0; for (let i = 0; i < vec.length; i++) score += vec[i] * q[i]; return { id, score: Math.round(score * 1000) / 1000 }; }) .sort((a, b) => b.score - a.score) .slice(0, limit); } async function status() { const { rows: [counts] } = await db.query( `SELECT (SELECT COUNT(*)::int FROM models WHERE deleted_at IS NULL AND embedding IS NOT NULL) AS indexed, (SELECT COUNT(*)::int ${PENDING}) AS pending`, [SALT]); return { ...counts, error: lastError }; } // Something outside the worker changed which models are searchable (trash, restore). function invalidate() { vectors = null; } module.exports = { run, search, status, invalidate };