Update server.js

server.js

@@ -1,273 +1,352 @@
-import { pipeline, TextStreamer } from '@huggingface/transformers';
+import { pipeline } from '@huggingface/transformers';
 import http from 'http';
-import fs from 'fs';
-import path from 'path';
 
-const PORT          = 7860;
-const MODEL_NAME    = 'onnx-community/Qwen2.5-0.5B-Instruct';
-const KNOWLEDGE_DIR = './knowledge';
+const PORT       = 7860;
+const MODEL_NAME = 'onnx-community/Qwen2.5-0.5B-Instruct';
 let generator;
-let knowledgeBase   = [];
-
-// ── KNOWLEDGE ─────────────────────────────────────────────────────────────────
-function loadKnowledge() {
-    if (!fs.existsSync(KNOWLEDGE_DIR)) { fs.mkdirSync(KNOWLEDGE_DIR); return; }
-    const files = fs.readdirSync(KNOWLEDGE_DIR).filter(f => f.endsWith('.txt'));
-    knowledgeBase = [];
-    for (const file of files) {
-        const content = fs.readFileSync(path.join(KNOWLEDGE_DIR, file), 'utf-8');
-        const chunks  = splitChunks(content, 800, 80);
-        chunks.forEach(c => knowledgeBase.push({ source: file, text: c }));
-    }
-    console.log(`Knowledge loaded: ${files.length} files, ${knowledgeBase.length} chunks`);
-}
-
-function splitChunks(text, size, overlap) {
-    const chunks = [];
-    let start    = 0;
-    while (start < text.length) {
-        chunks.push(text.slice(start, start + size));
-        start += size - overlap;
-    }
-    return chunks;
-}
+let modelReady = false;
 
-function retrieveContext(prompt, topK = 5) {
-    if (knowledgeBase.length === 0) return '';
+// ── QUEUE SYSTEM ──────────────────────────────────────────────────────────────
+const queue        = [];   // pending requests
+const MAX_PARALLEL = 2;    // max concurrent generations
+let   activeCount  = 0;
 
-    const words = prompt.toLowerCase()
-        .split(/\W+/)
-        .filter(w => w.length > 2);
-
-    const scored = knowledgeBase.map(chunk => {
-        const lower = chunk.text.toLowerCase();
-        let score   = 0;
-        for (const w of words) {
-            const matches = (lower.match(new RegExp(w, 'g')) || []).length;
-            score += matches;
-        }
-        return { ...chunk, score };
+function enqueue(task) {
+    return new Promise((resolve, reject) => {
+        queue.push({ task, resolve, reject });
+        processQueue();
     });
+}
 
-    const top = scored
-        .filter(c => c.score > 0)
-        .sort((a, b) => b.score - a.score)
-        .slice(0, topK);
+async function processQueue() {
+    if (activeCount >= MAX_PARALLEL || queue.length === 0) return;
+    activeCount++;
+    const { task, resolve, reject } = queue.shift();
+    try {
+        resolve(await task());
+    } catch (err) {
+        reject(err);
+    } finally {
+        activeCount--;
+        processQueue();
+    }
+}
 
-    if (top.length === 0) return '';
+function getQueuePosition() {
+    return queue.length;
+}
 
-    return top.map(c => c.text).join('\n\n---\n\n');
+// ── 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 buildMessages(prompt) {
-    const ragContext = retrieveContext(prompt, 5);
-    console.log(`RAG chunks found: ${ragContext.length} chars`);
-
-    const systemPrompt = ragContext
-        ? `You are Mlimi Connect AI, a free agricultural advisor for Malawian farmers.
-KNOWLEDGE BASE — THIS IS YOUR ONLY SOURCE OF INFORMATION. USE ONLY THIS:
-===START OF KNOWLEDGE===
-${ragContext}
-===END OF KNOWLEDGE===
-STRICT RULES:
-1. Answer ONLY using the knowledge provided above between ===START=== and ===END===.
-2. Do NOT add information from outside the knowledge base.
-3. Do NOT be vague. Give specific details: variety names, exact spacing, fertilizer amounts, timing.
-4. Structure your answer clearly with numbered steps.
-5. If the knowledge above does not contain the answer, say: "I don't have specific information on that in my knowledge base."
-6. ONLY answer agriculture questions. For anything else say: "I can only help with farming questions."`
-
-        : `You are Mlimi Connect AI, a free agricultural advisor for Malawian farmers.
-I don't have specific notes on that topic in my knowledge base yet.
-Give a brief, honest answer based on general Malawian agricultural knowledge.
-Keep it practical and specific to Malawi's conditions.
-ONLY answer agriculture questions.`;
-
-    return {
-        messages: [
-            { role: 'system', content: systemPrompt },
-            { role: 'user',   content: prompt        }
-        ],
-        ragContext
-    };
+function sendJSON(res, status, data) {
+    setCORS(res);
+    res.setHeader('Content-Type', 'application/json');
+    res.writeHead(status);
+    res.end(JSON.stringify(data));
 }
 
 // ── MODEL ─────────────────────────────────────────────────────────────────────
 async function loadModel() {
     console.log("Loading model...");
-    generator = await pipeline('text-generation', MODEL_NAME, { dtype: 'q4' });
+    generator  = await pipeline('text-generation', MODEL_NAME, { dtype: 'q4' });
+    modelReady = true;
     console.log("Model ready!");
 }
 
-// Non-streaming (kept for backward compatibility with /generate)
-async function generateResponse(messages) {
-    const output = await generator(messages, {
-        max_new_tokens:     600,
-        temperature:        0.2,
-        repetition_penalty: 1.15,
-        do_sample:          false
-    });
-    const generated = output[0].generated_text;
-    if (Array.isArray(generated)) return generated.at(-1)?.content || '';
-    return String(generated || '');
+// ── SINGLE CHUNK GENERATION ───────────────────────────────────────────────────
+// Generates one chunk with retry on failure
+async function generateChunk(messages, maxTokens, attempt = 0) {
+    try {
+        const output = await generator(messages, {
+            max_new_tokens:     maxTokens,
+            temperature:        0.3,
+            repetition_penalty: 1.15,
+            do_sample:          false
+        });
+        const generated = output[0].generated_text;
+        if (Array.isArray(generated)) return generated.at(-1)?.content || '';
+        return String(generated || '');
+    } catch (err) {
+        if (attempt < 2) {
+            console.log(`Chunk failed, retrying (attempt ${attempt + 1})...`);
+            await new Promise(r => setTimeout(r, 1000));
+            return generateChunk(messages, maxTokens, attempt + 1);
+        }
+        throw err;
+    }
 }
 
-// Streaming version — calls onToken(token) for every generated token as it arrives
-async function generateResponseStream(messages, onToken) {
-    const tokenizer = generator.tokenizer;
-
-    const streamer = new TextStreamer(tokenizer, {
-        skip_prompt: true,
-        skip_special_tokens: true,
-        callback_function: (text) => {
-            if (text) onToken(text);
+// ── 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  = '';
+
+    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}`;
         }
-    });
 
-    const output = await generator(messages, {
-        max_new_tokens:     600,
-        temperature:        0.2,
-        repetition_penalty: 1.15,
-        do_sample:          false,
-        streamer:           streamer
-    });
+        const messages = [
+            { role: 'system', content: systemPrompt },
+            { 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;
+        }
 
-    const generated = output[0].generated_text;
-    if (Array.isArray(generated)) return generated.at(-1)?.content || '';
-    return String(generated || '');
-}
+        if (!chunkText || chunkText.trim().length === 0) continue;
 
-// ── SERVER ────────────────────────────────────────────────────────────────────
-const server = http.createServer(async (req, res) => {
-    res.setHeader('Access-Control-Allow-Origin',  '*');
-    res.setHeader('Access-Control-Allow-Methods', 'GET, POST, OPTIONS');
-    res.setHeader('Access-Control-Allow-Headers', 'Content-Type');
+        // 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));
+        }
 
-    if (req.method === 'OPTIONS') { res.writeHead(200); return res.end(); }
+        // Update context window for next chunk (sliding window)
+        fullOutput      += (isFirst ? '' : ' ') + chunkText;
+        previousContext  = fullOutput.slice(-CONTEXT_CHARS);
 
-    const pathname = req.url.split('?')[0];
+        // Send chunk complete event
+        res.write(`data: ${JSON.stringify({
+            type:        'chunk_done',
+            chunk:       chunkNum,
+            totalChunks: totalChunks
+        })}\n\n`);
 
-    if (pathname === '/' && req.method === 'GET') {
-        res.setHeader('Content-Type', 'application/json');
-        res.writeHead(200);
-        return res.end(JSON.stringify({
-            status: "running",
-            model:  MODEL_NAME,
-            knowledge_chunks: knowledgeBase.length
-        }));
+        // If model signals it's done early (short output), stop
+        if (chunkText.trim().length < 50 && !isFirst) break;
     }
 
-    if (pathname === '/reload-knowledge' && req.method === 'POST') {
-        loadKnowledge();
-        res.setHeader('Content-Type', 'application/json');
-        res.writeHead(200);
-        return res.end(JSON.stringify({ message: `Reloaded: ${knowledgeBase.length} chunks` }));
+    // Send final done event with complete text
+    res.write(`data: ${JSON.stringify({
+        type:   'done',
+        result: fullOutput.trim()
+    })}\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.
 
-    // ── Existing non-streaming endpoint (unchanged behavior) ──────────────
-    if (pathname === '/generate' && req.method === 'POST') {
-        let body = '';
-        req.on('data', c => { body += c.toString(); });
-        req.on('end', async () => {
-            res.setHeader('Content-Type', 'application/json');
-            try {
-                const { prompt } = JSON.parse(body);
+IDENTITY: Your name is Gini AI. Company: Emalawi19. Founder: Professor Roosevelt Chinkwende.
 
-                if (!generator) {
-                    res.writeHead(503);
-                    return res.end(JSON.stringify({ error: "Model still loading..." }));
-                }
+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.`;
+}
 
-                console.log(`Query: "${prompt}"`);
+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));
+}
 
-                const { messages, ragContext } = buildMessages(prompt);
-                let result = await generateResponse(messages);
+// ── HTTP SERVER ───────────────────────────────────────────────────────────────
+const server = http.createServer(async (req, res) => {
+    setCORS(res);
 
-                if (!result || result.trim().length < 10) {
-                    result = ragContext
-                        ? `Here is what my knowledge base says:\n\n${ragContext.slice(0, 800)}`
-                        : "I don't have specific information on that topic. Please ask your local agricultural extension officer.";
-                }
+    if (req.method === 'OPTIONS') {
+        res.writeHead(200);
+        return res.end('{}');
+    }
 
-                console.log(`Response: ${result.slice(0, 80)}...`);
-                res.writeHead(200);
-                res.end(JSON.stringify({ result }));
+    const pathname = req.url.split('?')[0];
 
-            } catch (err) {
-                console.error("Error:", err.message);
-                res.writeHead(500);
-                res.end(JSON.stringify({ error: err.message }));
-            }
+    // Status
+    if (pathname === '/' && req.method === 'GET') {
+        return sendJSON(res, 200, {
+            status:        modelReady ? "ready" : "loading",
+            model:         MODEL_NAME,
+            queue_length:  queue.length,
+            active:        activeCount
         });
-        return;
     }
 
-    // ── NEW: streaming endpoint (SSE) ──────────────────────────────────────
-    if (pathname === '/generate-stream' && req.method === 'POST') {
+    // 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 });
+    }
+
+    // Generate
+    if (pathname === '/generate' && req.method === 'POST') {
         let body = '';
         req.on('data', c => { body += c.toString(); });
         req.on('end', async () => {
-            try {
-                const { prompt } = JSON.parse(body);
-
-                if (!generator) {
-                    res.writeHead(503, { 'Content-Type': 'application/json' });
-                    return res.end(JSON.stringify({ error: "Model still loading..." }));
-                }
-
-                console.log(`Query (stream): "${prompt}"`);
-
-                res.writeHead(200, {
-                    'Content-Type':      'text/event-stream',
-                    'Cache-Control':     'no-cache',
-                    'Connection':        'keep-alive',
-                    'X-Accel-Buffering': 'no'   // disable proxy buffering (nginx etc.)
-                });
 
-                const { messages, ragContext } = buildMessages(prompt);
-
-                let fullText = '';
-                let sentAnything = false;
-
-                await generateResponseStream(messages, (token) => {
-                    fullText += token;
-                    sentAnything = true;
-                    res.write(`data: ${JSON.stringify({ token })}\n\n`);
-                });
+            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`);
+            }
 
-                // Fallback if the model produced nothing usable
-                if (!sentAnything || fullText.trim().length < 10) {
-                    const fallback = ragContext
-                        ? `Here is what my knowledge base says:\n\n${ragContext.slice(0, 800)}`
-                        : "I don't have specific information on that topic. Please ask your local agricultural extension officer.";
-                    res.write(`data: ${JSON.stringify({ token: fallback })}\n\n`);
+            // 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 });
                 }
+            }
 
-                res.write('data: [DONE]\n\n');
-                res.end();
+            // 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);
+            }
 
+            // Send queue position if waiting
+            res.write(`data: ${JSON.stringify({ type: 'start', totalTokens, queue: queue.length })}\n\n`);
+
+            try {
+                await slidingWindowStream(res, systemPrompt, prompt.trim(), totalTokens);
             } catch (err) {
-                console.error("Error:", err.message);
-                // If headers already sent (mid-stream), send an SSE error event instead
-                if (!res.headersSent) {
-                    res.writeHead(500, { 'Content-Type': 'application/json' });
-                    return res.end(JSON.stringify({ error: err.message }));
-                }
-                res.write(`data: ${JSON.stringify({ error: err.message })}\n\n`);
+                res.write(`data: ${JSON.stringify({ type: 'error', error: err.message })}\n\n`);
                 res.end();
             }
         });
         return;
     }
 
-    res.setHeader('Content-Type', 'application/json');
-    res.writeHead(404);
-    res.end(JSON.stringify({ error: "Not Found" }));
+    return sendJSON(res, 404, { error: "Not found" });
 });
 
-loadKnowledge();
 loadModel().then(() => {
     server.listen(PORT, '0.0.0.0', () => {
-        console.log(`Mlimi Connect backend on port ${PORT}`);
+        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)`);
     });
 });