server.js

/**
 * RigoChat Worker — Distributed Inference Node
 * ================================================
 * Lightweight HF Space that loads RigoChat-7B-v2 and exposes
 * an /inference endpoint. Multiple workers can be combined
 * via Promise.any for faster responses.
 */

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

// ── Config ──────────────────────────────────────────────────────────────────
const PORT       = parseInt(process.env.PORT ?? '7860');
const MODEL_DIR  = process.env.MODEL_DIR ?? './models';
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 CTX_SIZE   = parseInt(process.env.CTX_SIZE ?? '4096');
const WORKER_ID  = process.env.SPACE_ID ?? `worker-${Math.random().toString(36).slice(2, 6)}`;
const AUTH_KEY   = process.env.WORKER_KEY ?? 'zelin-cluster'; // Simple auth

// ── State ───────────────────────────────────────────────────────────────────
let _model   = null;
let _ready   = false;
let _loading = false;
let _stats   = { requests: 0, totalMs: 0, errors: 0, avgMs: 0 };

// ── Load Model ──────────────────────────────────────────────────────────────
async function loadModel() {
  if (_ready || _loading) return;
  _loading = true;
  const start = Date.now();

  try {
    console.log(`[Worker] Loading ${MODEL_FILE}...`);
    mkdirSync(MODEL_DIR, { recursive: true });

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

    const modelPath = path.join(MODEL_DIR, MODEL_FILE);

    // Download if needed
    if (!existsSync(modelPath)) {
      // Check for HF-prefixed names
      const files = readdirSync(MODEL_DIR).filter(f => f.endsWith('.gguf'));
      const found = files.find(f => f.includes('rigochat') || f.includes('RigoChat'));

      if (!found) {
        console.log(`[Worker] Downloading ${MODEL_FILE} (~1GB)...`);
        const { createModelDownloader } = await import('node-llama-cpp');
        const dl = 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] Downloading... ${pct}%`);
          },
        });
        await dl.download();
        console.log('\n[Worker] Download complete ✅');
      }
    }

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

    console.log('[Worker] Loading model into memory...');
    _model = await llama.loadModel({ modelPath: actualPath, gpuLayers: 0 });
    _ready = true;
    _loading = false;

    const elapsed = ((Date.now() - start) / 1000).toFixed(1);
    console.log(`[Worker] ✅ Ready in ${elapsed}s — ${WORKER_ID}`);
  } catch (err) {
    _loading = false;
    console.error('[Worker] Load error:', err.message);
  }
}

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

  const { LlamaChatSession } = await import('node-llama-cpp');
  const ctx = await _model.createContext({ contextSize: CTX_SIZE });

  const systemMsg = 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 sysFinal = systemMsg.includes('español') ? systemMsg : stylePrefix + systemMsg;

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

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

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

  let result = '';
  if (lastUser) {
    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?.();

  return result.trim();
}

// ── 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') {
    res.writeHead(200, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify({
      status: _ready ? 'ready' : (_loading ? 'loading' : 'error'),
      worker: WORKER_ID,
      model: MODEL_FILE,
      uptime: process.uptime(),
      memory: Math.round(process.memoryUsage().heapUsed / 1024 / 1024) + 'MB',
      stats: _stats,
    }));
    return;
  }

  // Inference endpoint
  if (url.pathname === '/inference' && req.method === 'POST') {
    // Simple auth
    const auth = req.headers['authorization'];
    if (auth !== `Bearer ${AUTH_KEY}`) {
      res.writeHead(401, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({ error: 'Unauthorized' }));
      return;
    }

    if (!_ready) {
      res.writeHead(503, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({ error: 'Model not ready', status: _loading ? 'loading' : 'error' }));
      return;
    }

    try {
      const body = await new Promise((resolve, reject) => {
        let data = '';
        req.on('data', c => data += c);
        req.on('end', () => resolve(JSON.parse(data)));
        req.on('error', reject);
        setTimeout(() => reject(new Error('Timeout')), 30000);
      });

      const { messages, maxTokens = 300, temperature = 0.7 } = body;
      if (!messages?.length) {
        res.writeHead(400, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify({ error: 'messages required' }));
        return;
      }

      const start = Date.now();
      const result = await infer(messages, maxTokens, temperature);
      const ms = Date.now() - start;

      _stats.requests++;
      _stats.totalMs += ms;
      _stats.avgMs = Math.round(_stats.totalMs / _stats.requests);

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

  // Root
  if (url.pathname === '/') {
    res.writeHead(200, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify({
      name: 'rigochat-worker',
      version: '1.0.0',
      status: _ready ? 'ready' : (_loading ? 'loading' : 'error'),
      worker: WORKER_ID,
      model: MODEL_FILE,
      endpoints: ['/health', '/inference'],
    }));
    return;
  }

  res.writeHead(404);
  res.end('Not found');
});

// ── Start ───────────────────────────────────────────────────────────────────
server.listen(PORT, '0.0.0.0', () => {
  console.log(`[Worker] HTTP server on port ${PORT}`);
  console.log(`[Worker] Worker ID: ${WORKER_ID}`);
  console.log(`[Worker] Loading model in background...`);
  loadModel();
});