printhive/lib/embeddings.js

// 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 };