printhive/lib/aiRuns.js

// AI tagging runs: a queue of models tagged a few at a time. Progress, tokens and
// cost live in the database, so a run can be watched, stopped and resumed, and a
// server restart pauses it instead of spending on its own.
const fs = require('fs');
const path = require('path');
const db = require('./db');
const store = require('./store');
const ai = require('./ai');

const CONCURRENCY = 3;
const MAX_FAILURES_IN_A_ROW = 5;
const MAX_IMAGE_BYTES = 3 * 1024 * 1024;
const IMAGE_TYPES = { '.png': 'image/png', '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.webp': 'image/webp', '.gif': 'image/gif' };

let resolveConfig = async () => ai.configFrom({});
let working = false;
let rerun = false;

const httpError = (message, status) => Object.assign(new Error(message), { status });

async function init(resolver) {
  resolveConfig = resolver;
  await db.query(
    `UPDATE ai_runs SET status = 'paused', note = 'Paused when the server restarted. Resume to carry on.'
     WHERE status = 'running'`);
}

// The model's thumbnail as a data URL, when there is one small enough to send.
function thumbnailDataUrl(model) {
  if (!model.thumbnail_path) return null;
  const abs = path.resolve(model.thumbnail_path);
  const type = IMAGE_TYPES[path.extname(abs).toLowerCase()];
  try {
    if (!type || fs.statSync(abs).size > MAX_IMAGE_BYTES) return null;
    return `data:${type};base64,${fs.readFileSync(abs).toString('base64')}`;
  } catch {
    return null;
  }
}

// Names of the files inside a model's archives, as extra context.
async function entryNames(model) {
  const { rows } = await db.query(
    `SELECT e.path FROM archive_entries e JOIN files f ON f.id = e.file_id
     WHERE f.model_id = $1 ORDER BY e.id LIMIT 60`, [model.id]);
  return rows.map((r) => r.path.split('/').pop()).filter(Boolean);
}

async function createRun(modelIds, options = {}) {
  const config = await resolveConfig();
  const problem = ai.missingConfig(config);
  if (problem) throw httpError(problem, 400);
  const ids = [...new Set((modelIds || []).map(Number).filter(Number.isInteger))];
  const maxCost = Number(options.maxCost) > 0 ? Number(options.maxCost) : null;

  const run = await db.transaction(async (client) => {
    const { rows: [created] } = await client.query(
      `INSERT INTO ai_runs (provider, model, options) VALUES ($1, $2, $3) RETURNING id`,
      [config.label, config.model, {
        overwriteDescriptions: Boolean(options.overwriteDescriptions),
        maxCost,
        pricesKnown: config.prices.input != null || config.provider === 'openrouter',
      }]);
    const { rowCount } = await client.query(
      `INSERT INTO ai_run_items (run_id, model_id)
       SELECT $1, id FROM models WHERE id = ANY($2::int[]) AND deleted_at IS NULL`, [created.id, ids]);
    if (!rowCount) throw httpError('No models to tag', 400);
    await client.query('UPDATE ai_runs SET total = $2 WHERE id = $1', [created.id, rowCount]);
    return created;
  });
  kick();
  return getRun(run.id);
}

async function finish(runId, status, error = null) {
  await db.query(
    `UPDATE ai_runs SET status = $2, error = $3, finished_at = now() WHERE id = $1 AND status = 'running'`,
    [runId, status, error]);
}

async function tagOne(runId, modelId, config, { images, knownTags, overwrite }) {
  const model = await store.getModel(modelId);
  if (!model || model.deleted_at) {
    await db.query(`UPDATE ai_run_items SET status = 'error', error = 'The model was deleted', updated_at = now()
                    WHERE run_id = $1 AND model_id = $2`, [runId, modelId]);
    await db.query('UPDATE ai_runs SET failed = failed + 1 WHERE id = $1', [runId]);
    return { ok: true };
  }
  try {
    model.entryNames = await entryNames(model);
    const result = await ai.describeModel(config, model, { image: images ? thumbnailDataUrl(model) : null, knownTags });
    await store.applyAiResult(modelId, { tags: result.tags, description: result.description, overwriteDescription: overwrite });
    await db.query(`UPDATE ai_run_items SET status = 'done', tags = $3, error = NULL, updated_at = now()
                    WHERE run_id = $1 AND model_id = $2`, [runId, modelId, result.tags]);
    await db.query(
      `UPDATE ai_runs SET done = done + 1, prompt_tokens = prompt_tokens + $2,
         completion_tokens = completion_tokens + $3, cost_usd = cost_usd + $4 WHERE id = $1`,
      [runId, result.usage.prompt, result.usage.completion, result.cost || 0]);
    return { ok: true, imageRejected: result.imageRejected };
  } catch (err) {
    const fatal = ai.fatalMessage(err);
    if (fatal) return { ok: false, fatal }; // left pending, so resuming retries it
    const message = err.message.slice(0, 300);
    await db.query(`UPDATE ai_run_items SET status = 'error', error = $3, updated_at = now()
                    WHERE run_id = $1 AND model_id = $2`, [runId, modelId, message]);
    await db.query('UPDATE ai_runs SET failed = failed + 1 WHERE id = $1', [runId]);
    return { ok: false, error: message };
  }
}

async function processRun(runId) {
  const config = await resolveConfig();
  const problem = ai.missingConfig(config);
  if (problem) { await finish(runId, 'failed', problem); return; }
  const knownTags = await store.topTags(150);
  let images = config.sendImages;
  let failuresInARow = 0;
  let lastError = null;

  for (;;) {
    const { rows: [run] } = await db.query('SELECT status, options, cost_usd FROM ai_runs WHERE id = $1', [runId]);
    if (!run || run.status !== 'running') return; // stopped, or deleted
    const { maxCost, overwriteDescriptions } = run.options || {};
    const { rows: items } = await db.query(
      `SELECT model_id FROM ai_run_items WHERE run_id = $1 AND status = 'pending' ORDER BY model_id LIMIT $2`,
      [runId, CONCURRENCY]);
    if (!items.length) { await finish(runId, 'done'); return; }
    // Checked between batches, so a run can go over by up to one batch.
    if (maxCost && Number(run.cost_usd) >= maxCost) {
      await finish(runId, 'stopped', `Stopped at your $${maxCost} spending limit`);
      return;
    }

    const results = await Promise.all(items.map((item) =>
      tagOne(runId, item.model_id, config, { images, knownTags, overwrite: overwriteDescriptions })));
    for (const r of results) {
      if (r.imageRejected && images) {
        images = false;
        await db.query('UPDATE ai_runs SET note = $2 WHERE id = $1',
          [runId, `${config.model} doesn't accept images, so models are tagged from their names and files.`]);
      }
      if (r.fatal) { await finish(runId, 'failed', r.fatal); return; }
      if (r.ok) failuresInARow = 0;
      else { failuresInARow += 1; lastError = r.error; }
    }
    if (failuresInARow >= MAX_FAILURES_IN_A_ROW) {
      await finish(runId, 'failed', `Stopped after ${failuresInARow} failures in a row. Last: ${lastError}`);
      return;
    }
  }
}

// Works through running runs, one at a time. Safe to call often.
function kick() {
  if (working) { rerun = true; return; }
  working = true;
  (async () => {
    for (;;) {
      const { rows: [run] } = await db.query(`SELECT id FROM ai_runs WHERE status = 'running' ORDER BY id LIMIT 1`);
      if (!run) return;
      try {
        await processRun(run.id);
      } catch (err) {
        console.error(`AI run ${run.id} failed:`, err.message);
        await finish(run.id, 'failed', err.message.slice(0, 300));
      }
    }
  })()
    .catch((err) => console.error('AI runs stopped:', err.message))
    .finally(() => {
      working = false;
      if (rerun) { rerun = false; kick(); }
    });
}

const shapeRun = (r) => r && ({
  ...r,
  prompt_tokens: Number(r.prompt_tokens),
  completion_tokens: Number(r.completion_tokens),
  cost_usd: Number(r.cost_usd),
});

async function listRuns(limit = 15) {
  const { rows } = await db.query('SELECT * FROM ai_runs ORDER BY id DESC LIMIT $1', [limit]);
  return rows.map(shapeRun);
}

// A run with its latest results and errors, for watching it.
async function getRun(id) {
  const { rows: [run] } = await db.query('SELECT * FROM ai_runs WHERE id = $1', [id]);
  if (!run) return null;
  const { rows: recent } = await db.query(
    `SELECT i.model_id, m.name, i.status, i.tags, i.error FROM ai_run_items i
     LEFT JOIN models m ON m.id = i.model_id
     WHERE i.run_id = $1 AND i.status <> 'pending' ORDER BY i.updated_at DESC LIMIT 12`, [id]);
  return { ...shapeRun(run), recent };
}

async function stop(id) {
  const { rowCount } = await db.query(
    `UPDATE ai_runs SET status = 'stopped', finished_at = now(), note = NULL
     WHERE id = $1 AND status IN ('running', 'paused')`, [id]);
  return rowCount > 0;
}

// Carries on a paused, stopped or failed run, retrying models that errored.
async function resume(id) {
  const resumed = await db.transaction(async (client) => {
    const { rows: [run] } = await client.query(
      `SELECT id FROM ai_runs WHERE id = $1 AND status IN ('paused', 'stopped', 'failed') FOR UPDATE`, [id]);
    if (!run) return false;
    const retried = await client.query(
      `UPDATE ai_run_items SET status = 'pending', error = NULL WHERE run_id = $1 AND status = 'error'`, [id]);
    const { rows: [left] } = await client.query(
      `SELECT COUNT(*)::int AS n FROM ai_run_items WHERE run_id = $1 AND status = 'pending'`, [id]);
    if (!left.n) return false;
    await client.query(
      `UPDATE ai_runs SET status = 'running', error = NULL, note = NULL, finished_at = NULL,
         failed = GREATEST(failed - $2, 0) WHERE id = $1`, [id, retried.rowCount]);
    return true;
  });
  if (resumed) kick();
  return resumed;
}

async function hasActiveRun() {
  const { rows } = await db.query(`SELECT 1 FROM ai_runs WHERE status = 'running' LIMIT 1`);
  return rows.length > 0;
}

module.exports = { init, createRun, listRuns, getRun, stop, resume, kick, hasActiveRun, thumbnailDataUrl };