Aime/app/emb.js

// Text embeddings for semantic search (multilingual-e5-small via onnxruntime).
//
// Entirely optional: if the model directory is missing, available() is false,
// every call resolves to null and search.js silently falls back to full-text.
// That keeps the search working on a fresh install before anything is
// downloaded, instead of 500-ing.
//
// Same daemon pattern as spk.js: one long-lived python process on the A55 cores
// (4-7 are reserved for whisper/sherpa), line-delimited JSON over stdio, so the
// ~2 s model load is paid once instead of per query.
import { spawn } from 'node:child_process';
import fs from 'node:fs';
import readline from 'node:readline';
import path from 'node:path';
import * as gate from './gate.js';

const PY = process.env.AIME_PY || '/opt/aime-venv/bin/python';
const DIR = process.env.AIME_EMB_DIR || '/opt/aime-models/e5-small';
const SCRIPT = '/opt/callrec/inference/embed_text.py';
// e5 models are trained with these prefixes and lose a lot of accuracy without
// them — the asymmetric pair is what makes a short question match a long passage.
const Q_PREFIX = 'query: ';
const P_PREFIX = 'passage: ';
const BATCH = +(process.env.AIME_EMB_BATCH || 16);

let proc = null, ready = null, rl = null, pending = null, chain = Promise.resolve();

export function available() {
  return fs.existsSync(path.join(DIR, 'model.onnx')) && fs.existsSync(path.join(DIR, 'tokenizer.json'));
}

export function status() {
  return { available: available(), dir: DIR, running: !!proc };
}

function start() {
  if (proc) return ready;
  if (!available()) return Promise.reject(new Error('Embedding-Modell fehlt: ' + DIR));
  // Cores 0-3 (A55): the big cores stay free for whisper and sherpa, which are
  // the actual bottleneck. Embedding a query is a few ms either way.
  proc = spawn('taskset', ['-c', '0-3', PY, SCRIPT, DIR], { stdio: ['pipe', 'pipe', 'pipe'] });
  proc.stderr.on('data', d => process.stderr.write('[emb] ' + d));
  proc.on('exit', () => { proc = null; ready = null; rl = null; if (pending) { pending.reject(new Error('emb-Daemon beendet')); pending = null; } });
  ready = new Promise((res, rej) => {
    const to = setTimeout(() => rej(new Error('emb READY timeout')), 120000);
    rl = readline.createInterface({ input: proc.stdout });
    rl.on('line', line => {
      if (line.startsWith('READY')) { clearTimeout(to); res(); return; }
      if (!pending) return;
      const p = pending; pending = null;
      try {
        const msg = JSON.parse(line);
        if (msg.error) p.reject(new Error(msg.error)); else p.resolve(msg.vecs);
      } catch (e) { p.reject(new Error('emb: unlesbare Antwort — ' + line.slice(0, 120))); }
    });
  });
  return ready;
}

// Serialised through `chain`: the daemon handles one request at a time, and a
// concurrent write would interleave two JSON lines on the same stdin.
function request(texts) {
  const run = async () => {
    await start();
    return new Promise((resolve, reject) => {
      pending = { resolve, reject };
      const to = setTimeout(() => { if (pending) { pending = null; reject(new Error('emb timeout')); } }, 120000);
      const done = fn => v => { clearTimeout(to); fn(v); };
      pending = { resolve: done(resolve), reject: done(reject) };
      proc.stdin.write(JSON.stringify({ texts }) + '\n');
    });
  };
  chain = chain.then(run, run);
  return chain;
}

export async function embedQuery(text) {
  if (!available()) return null;
  const v = await request([Q_PREFIX + String(text || '').slice(0, 2000)]);
  return v && v[0];
}

// Returns null (not a throw) when the model is absent, so indexing just stores
// passages without vectors.
export async function embedMany(texts) {
  if (!available() || !texts || !texts.length) return null;
  const out = [];
  for (let i = 0; i < texts.length; i += BATCH) {
    const batch = texts.slice(i, i + BATCH).map(t => P_PREFIX + String(t || '').slice(0, 2000));
    out.push(...await request(batch));
  }
  return out;
}

// Beendet den Daemon und wartet auf sein Ende (siehe gate.freeRam: das e5-Modell
// hält in der Spitze 1,15 GB und muss weg, bevor rkllm lädt). Startet faul neu.
export function stop() {
  const p = proc;
  proc = null; ready = null; rl = null;
  return gate.killAndWait(p);
}