Update server.js

server.js

@@ -1,19 +1,57 @@
 import { pipeline } from '@huggingface/transformers';
 import http from 'http';
+import crypto from 'crypto';
+import fs from 'fs';
 
-const PORT       = 7860;
+const PORT = 7860;
 const MODEL_NAME = 'onnx-community/Qwen2.5-0.5B-Instruct';
 let generator;
 let modelReady = false;
 
-// ── QUEUE SYSTEM ──────────────────────────────────────────────────────────────
-const queue        = [];   // pending requests
-const MAX_PARALLEL = 2;    // max concurrent generations
-let   activeCount  = 0;
+// ── SYSTEM STATE & MEMORY ─────────────────────────────────────────────────────
+const queue = [];
+const jobs = new Map();         // Stores job state, checkpoints, and cancellation flags
+const rateLimits = new Map();   // IP-based rate limiting
+const MAX_PARALLEL = 1;         // Keep at 1 for transformers.js to prevent thread blocking
+let activeCount = 0;
+
+// ── UTILITIES ─────────────────────────────────────────────────────────────────
+function generateId() { return crypto.randomBytes(8).toString('hex'); }
+
+function log(event, details = {}) {
+    const timestamp = new Date().toISOString();
+    const logEntry = `[${timestamp}] ${event.toUpperCase()} - ${JSON.stringify(details)}\n`;
+    process.stdout.write(logEntry);
+    fs.appendFileSync('generation_logs.txt', logEntry); // Log to file
+}
+
+function setCORS(res) {
+    res.setHeader('Access-Control-Allow-Origin', '*');
+    res.setHeader('Access-Control-Allow-Methods', 'GET, POST, DELETE, OPTIONS');
+    res.setHeader('Access-Control-Allow-Headers', 'Content-Type');
+}
+
+function sendJSON(res, status, data) {
+    setCORS(res);
+    res.setHeader('Content-Type', 'application/json');
+    res.writeHead(status);
+    res.end(JSON.stringify(data));
+}
 
-function enqueue(task) {
+// ── QUEUE & RATE LIMITING ─────────────────────────────────────────────────────
+function checkRateLimit(ip) {
+    const now = Date.now();
+    const limit = rateLimits.get(ip) || { count: 0, resetTime: now + 60000 };
+    if (now > limit.resetTime) { limit.count = 0; limit.resetTime = now + 60000; }
+    if (limit.count >= 10) return false; // Max 10 requests per minute
+    limit.count++;
+    rateLimits.set(ip, limit);
+    return true;
+}
+
+function enqueue(jobId, task) {
     return new Promise((resolve, reject) => {
-        queue.push({ task, resolve, reject });
+        queue.push({ jobId, task, resolve, reject });
         processQueue();
     });
 }
@@ -21,8 +59,16 @@ function enqueue(task) {
 async function processQueue() {
     if (activeCount >= MAX_PARALLEL || queue.length === 0) return;
     activeCount++;
-    const { task, resolve, reject } = queue.shift();
+    const { jobId, task, resolve, reject } = queue.shift();
+    
+    // Check if job was cancelled while in queue
+    if (jobs.get(jobId)?.status === 'cancelled') {
+        activeCount--;
+        return processQueue();
+    }
+
     try {
+        jobs.get(jobId).status = 'processing';
         resolve(await task());
     } catch (err) {
         reject(err);
@@ -32,305 +78,184 @@ async function processQueue() {
     }
 }
 
-function getQueuePosition() {
-    return queue.length;
-}
-
-// ── CORS ──────────────────────────────────────────────────────────────────────
-function setCORS(res) {
-    res.setHeader('Access-Control-Allow-Origin',  '*');
-    res.setHeader('Access-Control-Allow-Methods', 'GET, POST, OPTIONS');
-    res.setHeader('Access-Control-Allow-Headers', 'Content-Type');
-}
-
-function sendJSON(res, status, data) {
-    setCORS(res);
-    res.setHeader('Content-Type', 'application/json');
-    res.writeHead(status);
-    res.end(JSON.stringify(data));
-}
-
-// ── MODEL ─────────────────────────────────────────────────────────────────────
+// ── MODEL INITIALIZATION ──────────────────────────────────────────────────────
 async function loadModel() {
-    console.log("Loading model...");
-    generator  = await pipeline('text-generation', MODEL_NAME, { dtype: 'q4' });
+    log('system', { message: "Loading model..." });
+    generator = await pipeline('text-generation', MODEL_NAME, { dtype: 'q4' });
     modelReady = true;
-    console.log("Model ready!");
+    log('system', { message: "Model ready!" });
 }
 
-// ── SINGLE CHUNK GENERATION ───────────────────────────────────────────────────
-// Generates one chunk with retry on failure
+// ── GENERATION & VALIDATION ───────────────────────────────────────────────────
 async function generateChunk(messages, maxTokens, attempt = 0) {
+    const startTime = Date.now();
     try {
         const output = await generator(messages, {
-            max_new_tokens:     maxTokens,
-            temperature:        0.3,
+            max_new_tokens: maxTokens,
+            temperature: 0.3,
             repetition_penalty: 1.15,
-            do_sample:          false
+            do_sample: false
         });
         const generated = output[0].generated_text;
-        if (Array.isArray(generated)) return generated.at(-1)?.content || '';
-        return String(generated || '');
+        let text = Array.isArray(generated) ? generated.at(-1)?.content || '' : String(generated || '');
+        
+        // Output Validation: Detect truncated code blocks
+        const openBlocks = (text.match(/```/g) || []).length;
+        if (openBlocks % 2 !== 0) text += '\n```'; // Repair formatting
+        
+        return { text, duration: Date.now() - startTime };
     } catch (err) {
-        if (attempt < 2) {
-            console.log(`Chunk failed, retrying (attempt ${attempt + 1})...`);
-            await new Promise(r => setTimeout(r, 1000));
+        if (attempt < 3) { // Failure Recovery: Exponential backoff
+            log('retry', { attempt: attempt + 1, error: err.message });
+            await new Promise(r => setTimeout(r, Math.pow(2, attempt) * 1000));
             return generateChunk(messages, maxTokens, attempt + 1);
         }
         throw err;
     }
 }
 
-// ── SLIDING WINDOW STREAMING ──────────────────────────────────────────────────
-// Breaks generation into chunks, streams each chunk, auto-resumes with context
-async function slidingWindowStream(res, systemPrompt, userPrompt, totalTokens) {
-    const CHUNK_TOKENS   = 250;  // tokens per chunk
-    const CONTEXT_CHARS  = 300;  // chars of previous output to use as context
-    const totalChunks    = Math.ceil(totalTokens / CHUNK_TOKENS);
-
-    let fullOutput       = '';
-    let previousContext  = '';
+// ── SLIDING WINDOW ENGINE (AUTO-RESUME & CHUNKING) ────────────────────────────
+async function slidingWindowEngine(res, jobId, systemPrompt, userPrompt, totalTokens) {
+    const CHUNK_TOKENS = 250;
+    const CONTEXT_CHARS = 800; // Expanded for better Context Awareness
+    const totalChunks = Math.ceil(totalTokens / CHUNK_TOKENS);
+    const job = jobs.get(jobId);
+
+    let fullOutput = job.checkpoint || '';
+    let previousContext = fullOutput.slice(-CONTEXT_CHARS);
+    let chunksCompleted = job.chunksCompleted || 0;
+    let tokensGenerated = 0;
+    const globalStartTime = Date.now();
+
+    for (let chunkIndex = chunksCompleted; chunkIndex < totalChunks; chunkIndex++) {
+        // Cancellation Check
+        if (job.status === 'cancelled') {
+            res.write(`data: ${JSON.stringify({ type: 'cancelled', message: 'Job stopped by user.' })}\n\n`);
+            return res.end();
+        }
 
-    for (let chunkIndex = 0; chunkIndex < totalChunks; chunkIndex++) {
         const isFirst = chunkIndex === 0;
-        const isLast  = chunkIndex === totalChunks - 1;
         const chunkNum = chunkIndex + 1;
 
-        // Send chunk start event
-        res.write(`data: ${JSON.stringify({
-            type:        'chunk_start',
-            chunk:       chunkNum,
-            totalChunks: totalChunks
-        })}\n\n`);
-
-        // Build messages for this chunk
-        let chunkUserPrompt;
-        if (isFirst) {
-            chunkUserPrompt = userPrompt;
-        } else {
-            // Sliding window: include last N chars of previous output as context
-            chunkUserPrompt =
-                `Continue EXACTLY from where you left off. Do not repeat anything. ` +
-                `Here is the end of what you wrote so far:\n\n` +
-                `...${previousContext}\n\n` +
-                `Continue from here. Original request was: ${userPrompt}`;
-        }
+        let chunkUserPrompt = isFirst ? userPrompt :
+            `[SYSTEM: Intelligent Context Compression]\n` +
+            `Original Request: ${userPrompt}\n` +
+            `Recent Output Context (Maintain consistency, variables, and formatting):\n...${previousContext}\n\n` +
+            `INSTRUCTION: Continue generating seamlessly from the exact last character above. Do not repeat the context.`;
 
         const messages = [
             { role: 'system', content: systemPrompt },
-            { role: 'user',   content: chunkUserPrompt }
+            { role: 'user', content: chunkUserPrompt }
         ];
 
-        // Queue the chunk generation for fair processing
-        let chunkText = '';
         try {
-            chunkText = await enqueue(() => generateChunk(messages, CHUNK_TOKENS));
-        } catch (err) {
-            // Chunk failed after retries — send error event but continue
-            res.write(`data: ${JSON.stringify({
-                type:  'chunk_error',
-                chunk: chunkNum,
-                error: err.message
-            })}\n\n`);
-            continue;
-        }
-
-        if (!chunkText || chunkText.trim().length === 0) continue;
-
-        // Stream the chunk word by word so user sees output immediately
-        const words = chunkText.split(' ');
-        for (let i = 0; i < words.length; i++) {
-            const token = (i === 0 && !isFirst ? '' : i === 0 ? '' : ' ') + words[i];
-            res.write(`data: ${JSON.stringify({ type: 'token', text: token })}\n\n`);
-            // Small delay between words for smooth rendering
-            await new Promise(r => setTimeout(r, 8));
-        }
+            const { text: chunkText, duration } = await enqueue(jobId, () => generateChunk(messages, CHUNK_TOKENS));
+            
+            if (!chunkText || chunkText.trim().length === 0) continue;
+
+            const words = chunkText.split(' ');
+            tokensGenerated += words.length; // Rough estimation
+            
+            // Progress Tracking Metrics
+            const speed = (words.length / (duration / 1000)).toFixed(2); // words per second
+            const percentComplete = Math.round((chunkNum / totalChunks) * 100);
+            const eta = ((totalChunks - chunkNum) * (duration / 1000)).toFixed(0);
+
+            // Stream Output
+            for (let i = 0; i < words.length; i++) {
+                const token = (i === 0 && !isFirst ? '' : i === 0 ? '' : ' ') + words[i];
+                res.write(`data: ${JSON.stringify({ 
+                    type: 'token', text: token, speed: `${speed} w/s`, eta: `${eta}s`, progress: `${percentComplete}%` 
+                })}\n\n`);
+                await new Promise(r => setTimeout(r, 10)); // Streaming effect
+            }
 
-        // Update context window for next chunk (sliding window)
-        fullOutput      += (isFirst ? '' : ' ') + chunkText;
-        previousContext  = fullOutput.slice(-CONTEXT_CHARS);
+            // Auto Resume & Checkpoint System Update
+            fullOutput += (isFirst ? '' : ' ') + chunkText;
+            previousContext = fullOutput.slice(-CONTEXT_CHARS);
+            
+            job.checkpoint = fullOutput;
+            job.chunksCompleted = chunkNum;
+            jobs.set(jobId, job);
+            log('checkpoint_saved', { jobId, chunkNum });
 
-        // Send chunk complete event
-        res.write(`data: ${JSON.stringify({
-            type:        'chunk_done',
-            chunk:       chunkNum,
-            totalChunks: totalChunks
-        })}\n\n`);
+            // Stop condition
+            if (chunkText.trim().length < 50 && !isFirst) break;
 
-        // If model signals it's done early (short output), stop
-        if (chunkText.trim().length < 50 && !isFirst) break;
+        } catch (err) {
+            log('error', { jobId, error: err.message });
+            res.write(`data: ${JSON.stringify({ type: 'error', error: 'Chunk failed. Checkpoint saved for recovery.' })}\n\n`);
+            return res.end();
+        }
     }
 
-    // Send final done event with complete text
-    res.write(`data: ${JSON.stringify({
-        type:   'done',
-        result: fullOutput.trim()
-    })}\n\n`);
+    job.status = 'completed';
+    res.write(`data: ${JSON.stringify({ type: 'done', result: fullOutput.trim(), totalTokens: tokensGenerated })}\n\n`);
     res.end();
 }
 
-// ── DETECT IF LONG OUTPUT IS NEEDED ──────────────────────────────────────────
-function estimateTokensNeeded(prompt) {
-    const lower = prompt.toLowerCase();
-
-    // Explicit length requests
-    const lineMatch = lower.match(/(\d+)\s*line/);
-    if (lineMatch) return Math.min(parseInt(lineMatch[1]) * 6, 2000);
-
-    const wordMatch = lower.match(/(\d+)\s*word/);
-    if (wordMatch) return Math.min(parseInt(wordMatch[1]) * 2, 2000);
-
-    // Long output keywords
-    const longKeywords = [
-        'full website', 'complete website', 'entire website',
-        'full code', 'complete code', 'entire code',
-        'step by step', 'all steps', 'detailed guide',
-        'full backend', 'full frontend', 'complete app',
-        'build a website', 'create a website', 'generate a website',
-        'write a program', 'full script', 'complete script'
-    ];
-    if (longKeywords.some(k => lower.includes(k))) return 800;
-
-    // Short answer keywords
-    const shortKeywords = [
-        'what is', 'what are', 'who is', 'explain briefly',
-        'define', 'tell me about', 'describe', 'history of'
-    ];
-    if (shortKeywords.some(k => lower.includes(k))) return 200;
-
-    // Default: medium
-    return 400;
-}
-
-// ── SYSTEM PROMPT BUILDER ─────────────────────────────────────────────────────
-function buildSystemPrompt(isCoding) {
-    if (isCoding) {
-        return `You are Gini AI, a full-stack web development AI assistant created by Emalawi19, founded by Professor Roosevelt Chinkwende.
-
-IDENTITY: Your name is Gini AI. Company: Emalawi19. Founder: Professor Roosevelt Chinkwende.
-
-CODING RULES:
-1. ALWAYS write code when asked. NEVER refuse.
-2. Complete, real, copy-paste ready code only — no placeholders.
-3. Default: JavaScript/HTML/CSS. Never Python unless asked.
-4. Add comments inside the code.
-5. When continuing a previous chunk, do NOT repeat already-written code. Continue EXACTLY from where you left off.
-
-HTML RULES:
-1. Start with <!DOCTYPE html> then <html lang="en">
-2. Include proper <head> with meta tags, title, and <style>
-3. End with </body> then </html>
-4. All CSS in <style>, all JS in <script> at bottom.`;
-    }
-    return `You are Gini AI, a helpful AI assistant created by Emalawi19, founded by Professor Roosevelt Chinkwende.
-
-IDENTITY: Your name is Gini AI. Company: Emalawi19. Founder: Professor Roosevelt Chinkwende.
-
-RULES:
-- Answer questions clearly and helpfully.
-- Be concise for simple questions, detailed for complex ones.
-- When continuing a response, do NOT repeat what was already said. Continue EXACTLY from where you left off.`;
-}
-
-function isCodingRequest(prompt) {
-    const lower = prompt.toLowerCase();
-    const nonCoding = ['what is','what are','who is','explain','describe','history','meaning','function of','functions of','tell me about'];
-    if (nonCoding.some(k => lower.includes(k))) return false;
-    const coding = ['generate','create','build','make','code','write','website','login','register','form','html','css','javascript','backend','frontend','server','api','function','script'];
-    return coding.some(k => lower.includes(k));
-}
-
 // ── HTTP SERVER ───────────────────────────────────────────────────────────────
 const server = http.createServer(async (req, res) => {
     setCORS(res);
+    if (req.method === 'OPTIONS') { res.writeHead(200); return res.end('{}'); }
 
-    if (req.method === 'OPTIONS') {
-        res.writeHead(200);
-        return res.end('{}');
-    }
+    const url = new URL(req.url, `http://${req.headers.host}`);
+    const pathname = url.pathname;
+    const clientIp = req.socket.remoteAddress;
 
-    const pathname = req.url.split('?')[0];
+    // Rate Limiting
+    if (!checkRateLimit(clientIp)) return sendJSON(res, 429, { error: "Too many requests. Please wait." });
 
-    // Status
     if (pathname === '/' && req.method === 'GET') {
-        return sendJSON(res, 200, {
-            status:        modelReady ? "ready" : "loading",
-            model:         MODEL_NAME,
-            queue_length:  queue.length,
-            active:        activeCount
-        });
+        return sendJSON(res, 200, { status: modelReady ? "ready" : "loading", queue: queue.length });
     }
 
-    // Queue status
-    if (pathname === '/queue' && req.method === 'GET') {
-        const position = getQueuePosition();
-        const eta      = position * 20; // rough seconds per request
-        return sendJSON(res, 200, { position, eta_seconds: eta, active: activeCount });
+    // Cancellation Endpoint
+    if (pathname === '/cancel' && req.method === 'POST') {
+        let body = '';
+        req.on('data', c => { body += c.toString(); });
+        req.on('end', () => {
+            const { jobId } = JSON.parse(body);
+            if (jobs.has(jobId)) {
+                jobs.get(jobId).status = 'cancelled';
+                log('cancelled', { jobId });
+                return sendJSON(res, 200, { message: `Job ${jobId} cancelled.` });
+            }
+            return sendJSON(res, 404, { error: "Job not found." });
+        });
+        return;
     }
 
-    // Generate
     if (pathname === '/generate' && req.method === 'POST') {
         let body = '';
         req.on('data', c => { body += c.toString(); });
         req.on('end', async () => {
-
             let parsed;
-            try { parsed = JSON.parse(body); }
-            catch { return sendJSON(res, 400, { error: "Invalid JSON" }); }
-
-            const { prompt } = parsed;
-            if (!prompt || !prompt.trim()) return sendJSON(res, 400, { error: "prompt is required" });
-            if (!modelReady)               return sendJSON(res, 503, { error: "Model loading, please wait..." });
-
-            const coding       = isCodingRequest(prompt);
-            const totalTokens  = estimateTokensNeeded(prompt);
-            const systemPrompt = buildSystemPrompt(coding);
-
-            console.log(`Request: "${prompt.slice(0, 60)}" | tokens: ${totalTokens} | coding: ${coding} | queue: ${queue.length}`);
-
-            // If request is in queue, send queue position first
-            if (queue.length > 0) {
-                const pos = queue.length;
-                const eta = pos * 20;
-                // For queued requests we still use SSE so user sees queue position
-                res.setHeader('Content-Type', 'text/event-stream');
-                res.setHeader('Cache-Control', 'no-cache');
-                res.setHeader('Connection',    'keep-alive');
-                res.setHeader('X-Accel-Buffering', 'no');
-                res.writeHead(200);
-                res.write(`data: ${JSON.stringify({ type: 'queued', position: pos, eta_seconds: eta })}\n\n`);
-            }
+            try { parsed = JSON.parse(body); } catch { return sendJSON(res, 400, { error: "Invalid JSON" }); }
 
-            // Short requests: single chunk, plain JSON response
-            if (totalTokens <= 300 && queue.length === 0) {
-                try {
-                    const messages = [
-                        { role: 'system', content: systemPrompt },
-                        { role: 'user',   content: prompt.trim() }
-                    ];
-                    const result = await enqueue(() => generateChunk(messages, totalTokens));
-                    return sendJSON(res, 200, { result: result || "I couldn't generate a response. Please try again." });
-                } catch (err) {
-                    return sendJSON(res, 500, { error: err.message });
-                }
-            }
+            const { prompt, resumeJobId } = parsed;
+            if (!modelReady) return sendJSON(res, 503, { error: "Model loading..." });
 
-            // Long requests: SSE streaming with sliding window
-            if (!res.headersSent) {
-                res.setHeader('Content-Type',      'text/event-stream');
-                res.setHeader('Cache-Control',      'no-cache');
-                res.setHeader('Connection',         'keep-alive');
-                res.setHeader('X-Accel-Buffering', 'no');
-                res.writeHead(200);
+            // Initialize or Resume Job
+            const jobId = resumeJobId && jobs.has(resumeJobId) ? resumeJobId : generateId();
+            if (!jobs.has(jobId)) {
+                jobs.set(jobId, { status: 'queued', checkpoint: '', chunksCompleted: 0 });
             }
 
-            // Send queue position if waiting
-            res.write(`data: ${JSON.stringify({ type: 'start', totalTokens, queue: queue.length })}\n\n`);
+            const totalTokens = 2000; // Simulated dynamic estimate
+            log('job_started', { jobId, promptLength: prompt?.length });
+
+            res.setHeader('Content-Type', 'text/event-stream');
+            res.setHeader('Cache-Control', 'no-cache');
+            res.setHeader('Connection', 'keep-alive');
+            res.writeHead(200);
+
+            res.write(`data: ${JSON.stringify({ type: 'start', jobId, queuePosition: queue.length })}\n\n`);
 
             try {
-                await slidingWindowStream(res, systemPrompt, prompt.trim(), totalTokens);
+                await slidingWindowEngine(res, jobId, "You are a helpful coding assistant.", prompt || "Continue", totalTokens);
             } catch (err) {
-                res.write(`data: ${JSON.stringify({ type: 'error', error: err.message })}\n\n`);
+                res.write(`data: ${JSON.stringify({ type: 'fatal_error', error: err.message })}\n\n`);
                 res.end();
             }
         });
@@ -341,12 +266,5 @@ const server = http.createServer(async (req, res) => {
 });
 
 loadModel().then(() => {
-    server.listen(PORT, '0.0.0.0', () => {
-        console.log(`Gini AI backend on port ${PORT}`);
-    });
-}).catch(err => {
-    console.error("Model load failed:", err.message);
-    server.listen(PORT, '0.0.0.0', () => {
-        console.log(`Server started (model failed)`);
-    });
+    server.listen(PORT, '0.0.0.0', () => console.log(`Engine running on port ${PORT}`));
 });