server.mjs

/**
 * Sahon AI - Transformers.js Inference Server
 * =============================================
 * Pure JavaScript LLM inference using @huggingface/transformers.
 * Started as a subprocess by app.py (Gradio bootstrap).
 */

import { pipeline } from '@huggingface/transformers';
import http from 'http';
import { fileURLToPath } from 'url';
import { dirname, join } from 'path';

// ─── Config ───
const PORT = parseInt(process.env.NODE_PORT || '8888');
const MODEL_ID = 'Xenova/phi-3-mini-4k-instruct';

// ─── State ───
let generator = null;
let modelReady = false;
let modelError = null;
let modelProgress = 'Initializing...';

// ─── Mission Barisal Checker ───
const missionBarisal = {
  checkBias(q, r) {
    const lowerR = r.toLowerCase();
    const strongWords = ['always', 'never', 'everyone', 'nobody', 'definitely', 'absolutely'];
    const found = strongWords.filter(w => lowerR.includes(w));
    return {
      biased: found.length > 2,
      reasons: found.length > 0 ? [`Absolute language: ${found.join(', ')}`] : [],
      score: Math.max(0, 1 - found.length * 0.1)
    };
  },

  checkHallucination(r) {
    const lowerR = r.toLowerCase();
    const hedgePhrases = ['i think', 'maybe', 'perhaps', 'possibly', 'might be'];
    const found = hedgePhrases.filter(p => lowerR.includes(p));
    const words = r.split(/\s+/).length;
    let risk = 'low';
    if (words > 500 && found.length > 2) risk = 'medium';
    if (words > 1000 && found.length > 3) risk = 'high';
    return { risk, indicators: found.length ? [`Hedging: ${found.join(', ')}`] : [] };
  },

  validate(question, response) {
    if (!response || !response.trim()) return { passed: false, overall: 0 };
    const bias = this.checkBias(question, response);
    const hal = this.checkHallucination(response);
    const overall = Math.round(
      (Math.min(1, response.length / 100) * 0.3 +
       bias.score * 0.35 +
       (hal.risk === 'low' ? 1 : hal.risk === 'medium' ? 0.7 : 0.4) * 0.35) * 100
    ) / 100;
    return {
      passed: overall > 0.5,
      quality_score: overall,
      checks: { bias, hallucination: hal, length: response.split(/\s+/).length }
    };
  }
};

// ─── Load Model ───
async function initModel() {
  try {
    modelProgress = 'Loading Transformers.js model...';
    console.log(`[Sahon] Loading model: ${MODEL_ID}`);

    generator = await pipeline('text-generation', MODEL_ID, {
      dtype: 'q4',
      device: 'cpu',
      progress_callback: (p) => {
        if (p.status === 'progress') {
          const pct = Math.round(p.progress * 100);
          modelProgress = `Downloading: ${pct}%`;
          console.log(`[Sahon] ${modelProgress}`);
        }
      }
    });

    modelReady = true;
    modelProgress = 'Ready';
    console.log('[Sahon] ✅ Model loaded!');
  } catch (err) {
    modelError = err.message;
    modelProgress = `Error: ${err.message}`;
    console.error('[Sahon] ❌ Model failed:', err);
  }
}

initModel();

// ─── Build Phi-3 Prompt ───
function buildPrompt(messages) {
  let prompt = '';
  for (const msg of messages) {
    switch (msg.role) {
      case 'system': prompt += `<|system|>\n${msg.content}<|end|>\n`; break;
      case 'user':   prompt += `<|user|>\n${msg.content}<|end|>\n`; break;
      case 'assistant': prompt += `<|assistant|>\n${msg.content}<|end|>\n`; break;
    }
  }
  prompt += '<|assistant|>\n';
  return prompt;
}

// ─── Parse JSON Body ───
function parseBody(req) {
  return new Promise((resolve, reject) => {
    let body = '';
    req.on('data', chunk => body += chunk);
    req.on('end', () => {
      try { resolve(JSON.parse(body)); }
      catch (e) { reject(new Error('Invalid JSON')); }
    });
    req.on('error', reject);
  });
}

// ─── 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');

  if (req.method === 'OPTIONS') {
    res.writeHead(204);
    res.end();
    return;
  }

  const url = new URL(req.url, `http://${req.headers.host}`);
  const path = url.pathname;

  // ── Health ──
  if (path === '/health') {
    res.writeHead(200, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify({
      status: modelReady ? 'ok' : 'loading',
      model_ready: modelReady,
      progress: modelProgress,
      error: modelError,
    }));
    return;
  }

  // ── Check model readiness for API calls ──
  if (!modelReady && (path.startsWith('/v1/'))) {
    res.writeHead(503, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify({ error: 'Model loading', status: modelProgress }));
    return;
  }

  // ── GET /v1/models ──
  if (path === '/v1/models' && req.method === 'GET') {
    res.writeHead(200, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify({
      object: 'list',
      data: [{ id: 'phi-3-mini-4k-instruct', object: 'model', created: Math.floor(Date.now()/1000), owned_by: 'mission-barisal' }]
    }));
    return;
  }

  // ── POST /v1/chat/completions ──
  if (path === '/v1/chat/completions' && req.method === 'POST') {
    let body;
    try { body = await parseBody(req); }
    catch (e) {
      res.writeHead(400, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({ error: 'Invalid JSON' }));
      return;
    }

    const { model, messages = [], temperature = 0.7, max_tokens = 512 } = body;
    const prompt = buildPrompt(messages);

    try {
      const result = await generator(prompt, {
        max_new_tokens: max_tokens,
        temperature: temperature,
        do_sample: temperature > 0,
        return_full_text: false,
      });

      const text = result[0]?.generated_text?.trim() || '';
      const lastUser = [...messages].reverse().find(m => m.role === 'user');
      const checkResult = missionBarisal.validate(lastUser?.content || '', text);

      res.writeHead(200, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({
        id: `chatcmpl-${Date.now()}`,
        object: 'chat.completion',
        created: Math.floor(Date.now()/1000),
        model: model || 'phi-3-mini-4k-instruct',
        choices: [{
          index: 0,
          message: { role: 'assistant', content: text },
          finish_reason: 'stop',
        }],
        usage: {
          prompt_tokens: Math.ceil(prompt.length/4),
          completion_tokens: Math.ceil(text.length/4),
          total_tokens: Math.ceil((prompt.length + text.length)/4)
        },
        _mission_barisal: checkResult
      }));
    } catch (err) {
      res.writeHead(500, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({ error: err.message }));
    }
    return;
  }

  // ── POST /v1/completions ──
  if (path === '/v1/completions' && req.method === 'POST') {
    let body;
    try { body = await parseBody(req); }
    catch (e) {
      res.writeHead(400); res.end(JSON.stringify({ error: 'Invalid JSON' }));
      return;
    }

    try {
      const result = await generator(body.prompt, {
        max_new_tokens: body.max_tokens || 512,
        temperature: body.temperature || 0.7,
        do_sample: true,
        return_full_text: false,
      });
      const text = result[0]?.generated_text?.trim() || '';
      res.writeHead(200, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({
        id: `cmpl-${Date.now()}`,
        object: 'text_completion',
        created: Math.floor(Date.now()/1000),
        model: body.model || 'phi-3-mini-4k-instruct',
        choices: [{ index: 0, text, finish_reason: 'stop' }],
        usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 }
      }));
    } catch (err) {
      res.writeHead(500); res.end(JSON.stringify({ error: err.message }));
    }
    return;
  }

  // ── 404 ──
  res.writeHead(404);
  res.end('Not Found');
});

server.listen(PORT, '127.0.0.1', () => {
  console.log(`[Sahon] Transformers.js server on http://127.0.0.1:${PORT}`);
  console.log(`[Sahon] Model: ${MODEL_ID}`);
});

// Graceful shutdown
process.on('SIGTERM', () => { server.close(); process.exit(0); });
process.on('SIGINT', () => { server.close(); process.exit(0); });