printhive/lib/autotag.js

// Automatic tagging: strict whole-word rules over model names, file names, and the
// names (and immediate folders) of files inside archives. Rule tags are recorded
// as such (model_tags.source = 'rule'), so a pass only adds or removes its own
// tags; ones you typed are never touched.
const crypto = require('crypto');
const db = require('./db');
const store = require('./store');

// Bump when the built-in rules change, so every model is checked again.
const RULES_VERSION = 1;

const BUILT_IN_RULES = [
  { group: 'Part types', tag: 'bracket', phrases: ['bracket'] },
  { group: 'Part types', tag: 'holder', phrases: ['holder'] },
  { group: 'Part types', tag: 'mount', phrases: ['mount'] },
  { group: 'Part types', tag: 'hinge', phrases: ['hinge'] },
  { group: 'Part types', tag: 'drawer', phrases: ['drawer'] },
  { group: 'Part types', tag: 'baseplate', phrases: ['baseplate', 'base plate'] },
  { group: 'Part types', tag: 'hook', phrases: ['hook'] },
  { group: 'Part types', tag: 'clip', phrases: ['clip'] },
  { group: 'Part types', tag: 'organizer', phrases: ['organizer', 'organiser'] },
  { group: 'Part types', tag: 'keychain', phrases: ['keychain', 'key chain', 'keyring', 'key ring'] },
  { group: 'Part types', tag: 'gridfinity', phrases: ['gridfinity'] },
  { group: 'Mechanisms', tag: 'print-in-place', phrases: ['print in place', 'printinplace', 'pip'] },
  { group: 'Mechanisms', tag: 'articulated', phrases: ['articulated', 'flexi'] },
  { group: 'Mechanisms', tag: 'threaded', phrases: ['threaded', 'thread'] },
  { group: 'Mechanisms', tag: 'gear', phrases: ['gear'] },
  { group: 'Materials', tag: 'pla', phrases: ['pla'] },
  { group: 'Materials', tag: 'petg', phrases: ['petg'] },
  { group: 'Materials', tag: 'abs', phrases: ['abs'] },
  { group: 'Materials', tag: 'asa', phrases: ['asa'] },
  { group: 'Materials', tag: 'tpu', phrases: ['tpu'] },
  { group: 'Materials', tag: 'nylon', phrases: ['nylon'] },
  { group: 'Techniques', tag: 'vase mode', phrases: ['vase mode', 'vasemode', 'spiral vase'] },
  { group: 'Techniques', tag: 'fuzzy skin', phrases: ['fuzzy skin', 'fuzzyskin'] },
  { group: 'Techniques', tag: 'multicolor', phrases: ['multicolor', 'multicolour', 'multi color', 'multi colour'] },
];

// Words in a name: splits on punctuation, camelCase and letter/number boundaries.
// "PETGWallHook_v2.stl" -> ["petg", "wall", "hook", "v", "2", "stl"]
function tokenize(text) {
  return String(text || '')
    .replace(/(?<=[a-z])(?=[A-Z])/g, ' ')
    .replace(/(?<=[A-Z])(?=[A-Z][a-z])/g, ' ')
    .replace(/(?<=[A-Za-z])(?=\d)|(?<=\d)(?=[A-Za-z])/g, ' ')
    .toLowerCase()
    .split(/[^a-z0-9]+/)
    .filter(Boolean);
}

const sameWord = (token, word) => token === word || token === `${word}s` || token === `${word}es`;

function hasPhrase(tokens, words) {
  if (!words.length) return false;
  for (let i = 0; i + words.length <= tokens.length; i++) {
    if (words.every((w, j) => sameWord(tokens[i + j], w))) return true;
  }
  return false;
}

const normalizeTag = (tag) => String(tag).trim().toLowerCase().replace(/\s+/g, ' ').slice(0, 40);

// rules: [{ tag, phrases }] -> compiled once for matching.
function compile(rules) {
  return rules.map((r) => ({ tag: normalizeTag(r.tag), phrases: r.phrases.map(tokenize).filter((w) => w.length) }));
}

const COMPILED_BUILT_INS = compile(BUILT_IN_RULES);

// Tags for a set of names (each matched separately, so a phrase never spans two).
function matchTags(texts, customRules = []) {
  const lists = (Array.isArray(texts) ? texts : [texts]).map(tokenize);
  const rules = [...COMPILED_BUILT_INS, ...compile(customRules.map((r) => ({ tag: r.tag, phrases: [r.phrase] })))];
  const tags = new Set();
  for (const rule of rules) {
    if (rule.tag && rule.phrases.some((words) => lists.some((tokens) => hasPhrase(tokens, words)))) tags.add(rule.tag);
  }
  return [...tags];
}

// ---------------------------------------------------------------- database
async function isEnabled() {
  return (await store.getSetting('autotag_enabled')) !== 'false';
}

async function listRules() {
  const { rows } = await db.query('SELECT id, phrase, tag FROM tag_rules ORDER BY id');
  return rows;
}

const saltFor = (rules) => crypto.createHash('md5').update(JSON.stringify([RULES_VERSION, rules])).digest('hex');

// Models whose names, files or archive contents changed since their rule tags
// were worked out (or whose rules changed).
const PENDING = `
  FROM (
    SELECT m.id, m.name, m.autotag_key,
      COALESCE((SELECT string_agg(f.original_name, E'\\n' ORDER BY f.id) FROM files f WHERE f.model_id = m.id), '') AS files,
      COALESCE((SELECT string_agg(e.path, E'\\n' ORDER BY e.id)
                FROM archive_entries e JOIN files f ON f.id = e.file_id WHERE f.model_id = m.id), '') AS entries
    FROM models m
  ) d
  WHERE d.autotag_key IS DISTINCT FROM md5($1 || d.name || E'\\n' || d.files || E'\\n' || d.entries)`;

async function run() {
  if (!(await isEnabled())) return false;
  const rules = await listRules();
  const salt = saltFor(rules.map((r) => [r.phrase, r.tag]));
  const { rows } = await db.query(`SELECT d.id, d.name, d.files, d.entries ${PENDING} ORDER BY d.id LIMIT 100`, [salt]);
  if (!rows.length) return false;
  for (const row of rows) {
    const entryNames = row.entries.split('\n').filter(Boolean).map((p) => p.split('/').slice(-2).join('/'));
    const texts = [row.name, ...row.files.split('\n'), ...entryNames.flatMap((p) => p.split('/'))];
    await store.replaceRuleTags(row.id, matchTags(texts, rules));
    await db.query(`UPDATE models SET autotag_key = md5($2 || $3 || E'\\n' || $4 || E'\\n' || $5) WHERE id = $1`,
      [row.id, salt, row.name, row.files, row.entries]);
  }
  return true;
}

async function pendingCount() {
  if (!(await isEnabled())) return 0;
  const rules = await listRules();
  const { rows } = await db.query(`SELECT COUNT(*)::int AS n ${PENDING}`, [saltFor(rules.map((r) => [r.phrase, r.tag]))]);
  return rows[0].n;
}

// Checks every model again on the next pass.
async function reapply() {
  await db.query('UPDATE models SET autotag_key = NULL');
}

// Turning it off keeps the tags already applied unless removeTags is set.
async function setEnabled(enabled, { removeTags = false } = {}) {
  await store.setSetting('autotag_enabled', enabled ? 'true' : 'false');
  await reapply();
  return !enabled && removeTags ? store.removeTagsFrom('rule') : 0;
}

async function addRule(phrase, tag) {
  const cleanPhrase = String(phrase || '').trim().slice(0, 80);
  const cleanTag = normalizeTag(tag || '');
  if (!tokenize(cleanPhrase).length) throw Object.assign(new Error('Enter a word or phrase to look for'), { status: 400 });
  if (!cleanTag) throw Object.assign(new Error('Enter the tag to add'), { status: 400 });
  await db.query('INSERT INTO tag_rules (phrase, tag) VALUES ($1, $2) ON CONFLICT DO NOTHING', [cleanPhrase, cleanTag]);
}

async function removeRule(id) {
  await db.query('DELETE FROM tag_rules WHERE id = $1', [id]);
}

module.exports = {
  BUILT_IN_RULES, tokenize, matchTags,
  run, pendingCount, reapply, isEnabled, setEnabled, listRules, addRule, removeRule,
};