serve_gemma.js

const http = require('http'), fs = require('fs'), path = require('path');
const fetch = globalThis.fetch || require('node-fetch'); // Node 18+ has built-in fetch
const PORT = 8150;
const ROOT = __dirname;
const SPLITS = path.join(ROOT, 'model_splits');

const MIME = {
  '.html': 'text/html', '.js': 'text/javascript', '.mjs': 'text/javascript',
  '.json': 'application/json', '.gguf': 'application/octet-stream',
  '.wasm': 'application/wasm', '.ts': 'text/javascript',
};

const ENTITY_DIR = path.join(ROOT, '..', 'entity');

async function extractKeywords(text) {
  // Try local xLAM at :8093 for intent-based keyword extraction
  try {
    const resp = await fetch('http://127.0.0.1:8093/v1/chat/completions', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({
        model: 'xlam',
        messages: [{ role: 'user', content: `Extract 3-5 keyword phrases (1-3 words each) from this message that capture the user's intent. Return JSON: {"keywords":["k1","k2","k3"]}\n\nMessage: "${text}"` }],
        temperature: 0.1,
        max_tokens: 60,
      }),
    });
    if (resp.ok) {
      const data = await resp.json();
      const content = data.choices?.[0]?.message?.content || '';
      const match = content.match(/\{[\s\S]*"keywords"[\s\S]*\}/);
      if (match) {
        const parsed = JSON.parse(match[0]);
        if (Array.isArray(parsed.keywords)) return parsed.keywords.map(k => k.toLowerCase());
      }
    }
  } catch (e) { /* xLAM not available, fall back */ }
  // Fallback: extract names and nouns (4+ chars, skip common words)
  const stop = new Set(['what','that','this','with','from','have','your','been','does','were','they','their','about','would','could','should','there','where','which','these','those','before','after','other','being','still','never','always','remember']);
  return (text.toLowerCase().match(/[a-z]{4,}/g) || []).filter(w => !stop.has(w)).slice(0, 5);
}

function grepEntity(text, keywords) {
  const unique = [...new Set(keywords)];
  if (unique.length === 0) return [];

  const results = [];
  function scanDir(dir) {
    try {
      const entries = fs.readdirSync(dir, { withFileTypes: true });
      for (const e of entries) {
        const fp = path.join(dir, e.name);
        if (e.isDirectory()) { scanDir(fp); continue; }
        if (!/\.(md|json|txt)$/.test(e.name)) continue;
        try {
          const content = fs.readFileSync(fp, 'utf8');
          const lines = content.split('\n');
          for (let i = 0; i < lines.length; i++) {
            const line = lines[i].trim();
            if (line.length < 10 || line.length > 300) continue;
            const lower = line.toLowerCase();
            for (const kw of unique) {
              if (lower.includes(kw)) {
                results.push(line);
                break;
              }
            }
          }
        } catch (e) {}
      }
    } catch (e) {}
  }
  scanDir(ENTITY_DIR);
  return [...new Set(results)].slice(0, 6);
}

http.createServer((req, res) => {
  // Handle grep API
  if (req.method === 'POST' && req.url === '/api/grep') {
    let body = '';
    req.on('data', c => body += c);
    req.on('end', async () => {
      try {
        const { text } = JSON.parse(body);
        const keywords = await extractKeywords(text);
        console.log('[grep] keywords:', keywords);
        const results = grepEntity(text, keywords);
        console.log('[grep] found:', results.length, 'snippets');
        res.writeHead(200, { 'Content-Type': 'application/json', 'Access-Control-Allow-Origin': '*' });
        res.end(JSON.stringify({ results, keywords }));
      } catch (e) {
        res.writeHead(500, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify({ error: e.message }));
      }
    });
    return;
  }

  if (req.method === 'OPTIONS') {
    res.writeHead(200, {
      'Access-Control-Allow-Origin': '*',
      'Access-Control-Allow-Methods': 'GET,HEAD',
      'Access-Control-Allow-Headers': 'Content-Type,Range',
      'Cross-Origin-Embedder-Policy': 'require-corp',
      'Cross-Origin-Opener-Policy': 'same-origin',
    });
    return res.end();
  }

  let p = decodeURIComponent(req.url.split('?')[0]);
  if (p === '/') p = '/index.html';

  let fp;
  if (p.startsWith('/model/')) {
    fp = path.join(SPLITS, p.slice('/model/'.length));
  } else {
    // Resolve relative to project root, allowing ../mamba_webgpu/ paths
    fp = path.resolve(ROOT, '.' + p);
  }

  fs.stat(fp, (e, st) => {
    if (e) { res.writeHead(404); return res.end('not found: ' + fp); }
    const total = st.size;
    const range = req.headers.range;
    const ct = MIME[path.extname(fp)] || 'application/octet-stream';

    const headers = {
      'Content-Type': ct,
      'Accept-Ranges': 'bytes',
      'Access-Control-Allow-Origin': '*',
      'Cross-Origin-Embedder-Policy': 'require-corp',
      'Cross-Origin-Opener-Policy': 'same-origin',
    };

    if (range) {
      const m = range.match(/bytes=(\d+)-(\d*)/);
      if (m) {
        const start = parseInt(m[1]);
        const end = m[2] ? parseInt(m[2]) : total - 1;
        headers['Content-Range'] = 'bytes ' + start + '-' + end + '/' + total;
        headers['Content-Length'] = end - start + 1;
        res.writeHead(206, headers);
        fs.createReadStream(fp, { start, end }).pipe(res);
        return;
      }
    }

    headers['Content-Length'] = total;
    res.writeHead(200, headers);
    fs.createReadStream(fp).pipe(res);
  });
}).listen(PORT, () => {
  console.log('Gemma WebGPU on :' + PORT);
  console.log('Model splits: ' + SPLITS);
  const files = fs.readdirSync(SPLITS).filter(f => f.endsWith('.gguf'));
  console.log('Split files: ' + files.length);
  console.log('CORS + COEP/COOP headers enabled for wllama multi-threading');
});