`lib/embed.js`

56 lines
'use strict';

// Ollama embedder using nomic-embed-text on Mac Studio 1
const OLLAMA_BASE = process.env.OLLAMA_BASE || 'http://192.168.1.133:11434';
const EMBED_MODEL = process.env.EMBED_MODEL || 'nomic-embed-text';
const GEN_MODEL   = process.env.GEN_MODEL   || 'qwen3:14b';

async function embed(text) {
  const ac = new AbortController();
  const t  = setTimeout(() => ac.abort(), 30_000); // 30s embed timeout
  try {
    const res = await fetch(`${OLLAMA_BASE}/api/embeddings`, {
      method: 'POST',
      headers: { 'content-type': 'application/json' },
      body: JSON.stringify({ model: EMBED_MODEL, prompt: text }),
      signal: ac.signal,
    });
    if (!res.ok) throw new Error(`Embed HTTP ${res.status}: ${await res.text()}`);
    const { embedding } = await res.json();
    return embedding; // float[]
  } finally {
    clearTimeout(t);
  }
}

async function generate(prompt) {
  const timeoutMs = parseInt(process.env.GEN_TIMEOUT_MS || '180000', 10); // 3 min default
  const ac = new AbortController();
  const t  = setTimeout(() => ac.abort(), timeoutMs);
  try {
    const res = await fetch(`${OLLAMA_BASE}/api/generate`, {
      method: 'POST',
      headers: { 'content-type': 'application/json' },
      body: JSON.stringify({ model: GEN_MODEL, prompt, stream: false }),
      signal: ac.signal,
    });
    if (!res.ok) throw new Error(`Generate HTTP ${res.status}: ${await res.text()}`);
    const { response } = await res.json();
    return response;
  } finally {
    clearTimeout(t);
  }
}

function cosineSim(a, b) {
  let dot = 0, na = 0, nb = 0;
  for (let i = 0; i < a.length; i++) {
    dot += a[i] * b[i];
    na  += a[i] * a[i];
    nb  += b[i] * b[i];
  }
  return dot / (Math.sqrt(na) * Math.sqrt(nb) + 1e-10);
}

module.exports = { embed, generate, cosineSim, OLLAMA_BASE, EMBED_MODEL, GEN_MODEL };