worker.js

/**
 * worker.js — Zelin Cluster Worker
 * ==================================
 * Runs RigoChat-7B-v2 via node-llama-cpp and exposes
 * an HTTP inference API for the cluster orchestrator.
 *
 * Endpoints:
 *   GET  /health      → { status, model, uptime, memMB }
 *   POST /inference    → { result, tokens, latencyMs }
 *   POST /inference/stream → NDJSON stream of tokens
 *   POST /embed        → { vector } (if embed model loaded)
 */

import http from 'http';
import os from 'os';
import path from 'path';
import { fileURLToPath } from 'url';
import { existsSync, mkdirSync, readdirSync } from 'fs';

const __dirname = path.dirname(fileURLToPath(import.meta.url));
const PORT = parseInt(process.env.PORT ?? '7860', 10);
const MODEL_DIR = process.env.MODEL_DIR ?? path.join(__dirname, 'models');
const AUTH_KEY = process.env.CLUSTER_AUTH_KEY ?? 'zelin-cluster';
const WORKER_ID = process.env.WORKER_ID ?? `worker-${Math.random().toString(36).slice(2, 6)}`;

// Model config from env or defaults
const MODEL_REPO = process.env.MODEL_REPO ?? 'IIC/RigoChat-7b-v2-GGUF';
const MODEL_FILE = process.env.MODEL_FILE ?? 'rigochat-7b-v2-Q4_K_M.gguf';
const CONTEXT_SIZE = parseInt(process.env.CONTEXT_SIZE ?? '4096', 10);
const MAX_TOKENS = parseInt(process.env.MAX_TOKENS ?? '512', 10);

// State
let model = null;
let llama = null;
let modelReady = false;
let modelLoading = false;
let modelError = null;
const startTime = Date.now();

// Inference semaphore (1 at a time for CPU)
let busy = false;
const queue = [];

function withLock(fn, timeoutMs = 120_000) {
  return new Promise((resolve, reject) => {
    const tryRun = () => {
      if (busy) {
        const t = setTimeout(() => reject(new Error('Lock timeout')), timeoutMs);
        queue.push(() => { clearTimeout(t); tryRun(); });
        return;
      }
      busy = true;
      fn()
        .then(resolve)
        .catch(reject)
        .finally(() => {
          busy = false;
          if (queue.length > 0) queue.shift()();
        });
    };
    tryRun();
  });
}

// ── Load Model ──────────────────────────────────────────────────────────────
async function loadModel() {
  if (modelReady || modelLoading) return modelReady;
  modelLoading = true;
  console.log(`[Worker ${WORKER_ID}] Loading model ${MODEL_FILE}...`);

  try {
    const { getLlama } = await import('node-llama-cpp');
    llama = await getLlama();

    mkdirSync(MODEL_DIR, { recursive: true });
    const modelPath = path.join(MODEL_DIR, MODEL_FILE);

    // Download model if needed
    if (!existsSync(modelPath)) {
      console.log(`[Worker ${WORKER_ID}] Downloading model from ${MODEL_REPO}...`);
      const { createModelDownloader } = await import('node-llama-cpp');
      const downloader = await createModelDownloader({
        modelUri: `hf:${MODEL_REPO}/${MODEL_FILE}`,
        dirPath: MODEL_DIR,
        onProgress: ({ downloadedSize, totalSize }) => {
          const pct = totalSize ? Math.round(downloadedSize / totalSize * 100) : '?';
          process.stdout.write(`\r[Worker ${WORKER_ID}] Downloading... ${pct}%`);
        },
      });
      await downloader.download();
      console.log(`\n[Worker ${WORKER_ID}] Model downloaded`);
    }

    // Find actual model file (HF may add prefix)
    let actualPath = modelPath;
    if (!existsSync(actualPath)) {
      const files = readdirSync(MODEL_DIR).filter(f => f.endsWith('.gguf'));
      const baseName = MODEL_FILE.replace('.gguf', '');
      const found = files.find(f => f.includes(baseName) || f.includes('rigochat') || f.includes('RigoChat'));
      if (found) actualPath = path.join(MODEL_DIR, found);
    }

    model = await llama.loadModel({ modelPath: actualPath, gpuLayers: 0 });
    modelReady = true;
    modelLoading = false;
    console.log(`[Worker ${WORKER_ID}] Model ready! RAM: ${Math.round(os.totalmem() / 1024 / 1024)}MB total, ${Math.round(os.freemem() / 1024 / 1024)}MB free`);
    return true;
  } catch (err) {
    modelError = err.message;
    modelLoading = false;
    console.error(`[Worker ${WORKER_ID}] Model load failed:`, err.message);
    return false;
  }
}

// ── Inference ───────────────────────────────────────────────────────────────
async function runInference(messages, maxTokens = 300, temperature = 0.7) {
  if (!modelReady) throw new Error('Model not ready');

  return withLock(async () => {
    const { LlamaChatSession } = await import('node-llama-cpp');
    const ctx = await model.createContext({ contextSize: CONTEXT_SIZE });

    // Build system prompt with style
    const baseSystem = messages.find(m => m.role === 'system')?.content ?? '';
    const stylePrefix = 'Responde en español casual argentino. Máx 2 líneas. Sin mayúsculas al inicio. Sin punto final. Sin emojis salvo que sea muy gracioso.\n\n';
    const systemFinal = baseSystem.includes('español casual') ? baseSystem : stylePrefix + baseSystem;

    const session = new LlamaChatSession({
      contextSequence: ctx.getSequence(),
      systemPrompt: systemFinal,
    });

    const userMessages = messages.filter(m => m.role !== 'system');
    const lastUser = userMessages[userMessages.length - 1];
    let result = '';

    if (lastUser) {
      // Inject recent history
      for (const msg of userMessages.slice(0, -1)) {
        if (msg.role === 'user') {
          await session.prompt(msg.content ?? '', { maxTokens: 1 }).catch(() => {});
        }
      }

      result = await session.prompt(lastUser.content ?? '', {
        maxTokens,
        temperature,
        topP: 0.9,
        topK: 40,
        minP: 0.05,
        repeatPenalty: {
          penalty: 1.35,
          lastTokens: 96,
          frequencyPenalty: 0.1,
          presencePenalty: 0.05,
        },
      });
    }

    session.dispose?.();
    ctx.dispose?.();

    let cleaned = result.trim();
    // Clean artifacts
    cleaned = cleaned.replace(/^PASO \d+:?\s*/i, '');
    cleaned = cleaned.replace(/^Respuesta final:?\s*/i, '');
    cleaned = cleaned.replace(/^Mi respuesta:?\s*/i, '');
    cleaned = cleaned.replace(/^Zelin:?\s*/i, '');

    return cleaned || result.trim();
  });
}

// ── Two-Pass Thinking (for complex queries) ────────────────────────────────
async function runThinking(messages, maxTokens = 400, temperature = 0.6) {
  if (!modelReady) throw new Error('Model not ready');

  return withLock(async () => {
    const { LlamaChatSession } = await import('node-llama-cpp');
    const ctx = await model.createContext({ contextSize: CONTEXT_SIZE });

    const baseSystem = messages.find(m => m.role === 'system')?.content ?? '';
    const stylePrefix = 'Responde en español casual argentino. Máx 2 líneas. Sin mayúsculas al inicio. Sin punto final.\n\n';
    const systemFinal = baseSystem.includes('español casual') ? baseSystem : stylePrefix + baseSystem;

    const session = new LlamaChatSession({
      contextSequence: ctx.getSequence(),
      systemPrompt: systemFinal,
    });

    const userMessages = messages.filter(m => m.role !== 'system');
    const lastUser = userMessages[userMessages.length - 1];

    if (!lastUser) throw new Error('No user message');

    // Inject history
    for (const msg of userMessages.slice(0, -1)) {
      if (msg.role === 'user') {
        await session.prompt(msg.content ?? '', { maxTokens: 1 }).catch(() => {});
      }
    }

    // Pass 1: Think
    const thinkPrompt = `PIENSA INTERNAMENTE (NO muestres esto). Analiza:
1. ¿Qué te preguntan? Intención real.
2. ¿Hay ironía o sarcasmo?
3. ¿Cómo respondería Zelin (argentina, casual, minúsculas)?
PIENSA y luego darás tu respuesta final.`;

    try {
      await session.prompt(thinkPrompt + '\n\nMensaje: ' + (lastUser.content ?? ''), {
        maxTokens: 200,
        temperature: 0.3,
        topP: 0.85,
        topK: 30,
        repeatPenalty: { penalty: 1.15, lastTokens: 64 },
      });

      // Pass 2: Respond
      const result = await session.prompt('Ahora da SOLO tu respuesta final como Zelin (1-2 líneas, español casual, minúsculas, sin punto final):', {
        maxTokens,
        temperature,
        topP: 0.9,
        topK: 40,
        repeatPenalty: { penalty: 1.35, lastTokens: 96, frequencyPenalty: 0.1, presencePenalty: 0.05 },
      });

      session.dispose?.();
      ctx.dispose?.();

      let cleaned = result.trim();
      cleaned = cleaned.replace(/^PASO \d+:?\s*/i, '');
      cleaned = cleaned.replace(/^Respuesta final:?\s*/i, '');
      cleaned = cleaned.replace(/^Zelin:?\s*/i, '');
      return cleaned || result.trim();
    } catch {
      session.dispose?.();
      ctx.dispose?.();
      throw new Error('Thinking inference failed');
    }
  });
}

// ── HTTP Server ─────────────────────────────────────────────────────────────
const server = http.createServer(async (req, res) => {
  // CORS
  res.setHeader('Access-Control-Allow-Origin', '*');
  res.setHeader('Access-Control-Allow-Methods', 'GET, POST, OPTIONS');
  res.setHeader('Access-Control-Allow-Headers', 'Content-Type, Authorization');
  if (req.method === 'OPTIONS') { res.writeHead(204); res.end(); return; }

  const url = new URL(req.url, `http://localhost:${PORT}`);

  // ── Health Check ──────────────────────────────────────────────────────
  if (url.pathname === '/health' && req.method === 'GET') {
    const mem = process.memoryUsage();
    res.writeHead(200, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify({
      status: modelReady ? 'ready' : (modelLoading ? 'loading' : 'error'),
      workerId: WORKER_ID,
      model: MODEL_FILE,
      uptime: Math.round((Date.now() - startTime) / 1000),
      memMB: Math.round(mem.rss / 1024 / 1024),
      busy,
      error: modelError,
    }));
    return;
  }

  // ── Auth check for POST endpoints ─────────────────────────────────────
  if (req.method === 'POST') {
    const auth = req.headers['authorization'];
    if (auth !== `Bearer ${AUTH_KEY}`) {
      res.writeHead(401, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({ error: 'Unauthorized' }));
      return;
    }
  }

  // ── Inference ─────────────────────────────────────────────────────────
  if (url.pathname === '/inference' && req.method === 'POST') {
    if (!modelReady) {
      res.writeHead(503, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({ error: 'Model not ready', status: modelLoading ? 'loading' : 'error' }));
      return;
    }

    try {
      const body = await readBody(req);
      const { messages, maxTokens = 300, temperature = 0.7, thinking = false } = JSON.parse(body);

      if (!messages || !Array.isArray(messages)) {
        res.writeHead(400, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify({ error: 'messages array required' }));
        return;
      }

      const startMs = Date.now();
      const result = thinking
        ? await runThinking(messages, maxTokens, temperature)
        : await runInference(messages, maxTokens, temperature);
      const latencyMs = Date.now() - startMs;

      res.writeHead(200, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({
        result,
        worker: WORKER_ID,
        latencyMs,
        tokens: result?.split(/\s+/).length ?? 0,
        thinking,
      }));
    } catch (err) {
      res.writeHead(500, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({ error: err.message, worker: WORKER_ID }));
    }
    return;
  }

  // ── Status (detailed) ─────────────────────────────────────────────────
  if (url.pathname === '/status' && req.method === 'GET') {
    res.writeHead(200, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify({
      workerId: WORKER_ID,
      model: MODEL_FILE,
      modelRepo: MODEL_REPO,
      ready: modelReady,
      loading: modelLoading,
      error: modelError,
      busy,
      uptime: Math.round((Date.now() - startTime) / 1000),
      memMB: Math.round(process.memoryUsage().rss / 1024 / 1024),
      totalMemMB: Math.round(os.totalmem() / 1024 / 1024),
      freeMemMB: Math.round(os.freemem() / 1024 / 1024),
      cpus: os.cpus().length,
    }));
    return;
  }

  // ── Root ──────────────────────────────────────────────────────────────
  if (url.pathname === '/' && req.method === 'GET') {
    res.writeHead(200, { 'Content-Type': 'text/html' });
    res.end(`<html><body><h1>Zelin Worker ${WORKER_ID}</h1><p>Status: ${modelReady ? 'READY' : (modelLoading ? 'LOADING' : 'ERROR')}</p><p>Model: ${MODEL_FILE}</p></body></html>`);
    return;
  }

  // 404
  res.writeHead(404, { 'Content-Type': 'application/json' });
  res.end(JSON.stringify({ error: 'Not found' }));
});

function readBody(req) {
  return new Promise((resolve, reject) => {
    let data = '';
    req.on('data', chunk => data += chunk);
    req.on('end', () => resolve(data));
    req.on('error', reject);
  });
}

// ── Start ───────────────────────────────────────────────────────────────────
server.listen(PORT, '0.0.0.0', () => {
  console.log(`[Worker ${WORKER_ID}] HTTP server on :${PORT}`);
  loadModel().then(ok => {
    if (ok) console.log(`[Worker ${WORKER_ID}] Ready for inference!`);
    else console.error(`[Worker ${WORKER_ID}] FAILED to load model`);
  });
});