Update server.js

server.js

@@ -1,4 +1,4 @@
-import { pipeline } from '@huggingface/transformers';
+import { pipeline, TextStreamer } from '@huggingface/transformers';
 import http from 'http';
 import fs from 'fs';
 import path from 'path';
@@ -43,7 +43,6 @@ function retrieveContext(prompt, topK = 5) {
         const lower = chunk.text.toLowerCase();
         let score   = 0;
         for (const w of words) {
-            // Count every occurrence not just presence — better scoring
             const matches = (lower.match(new RegExp(w, 'g')) || []).length;
             score += matches;
         }
@@ -60,6 +59,39 @@ function retrieveContext(prompt, topK = 5) {
     return top.map(c => c.text).join('\n\n---\n\n');
 }
 
+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
+    };
+}
+
 // ── MODEL ─────────────────────────────────────────────────────────────────────
 async function loadModel() {
     console.log("Loading model...");
@@ -67,6 +99,7 @@ async function loadModel() {
     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,
@@ -79,6 +112,31 @@ async function generateResponse(messages) {
     return String(generated || '');
 }
 
+// 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);
+        }
+    });
+
+    const output = await generator(messages, {
+        max_new_tokens:     600,
+        temperature:        0.2,
+        repetition_penalty: 1.15,
+        do_sample:          false,
+        streamer:           streamer
+    });
+
+    const generated = output[0].generated_text;
+    if (Array.isArray(generated)) return generated.at(-1)?.content || '';
+    return String(generated || '');
+}
+
 // ── SERVER ────────────────────────────────────────────────────────────────────
 const server = http.createServer(async (req, res) => {
     res.setHeader('Access-Control-Allow-Origin',  '*');
@@ -106,6 +164,7 @@ const server = http.createServer(async (req, res) => {
         return res.end(JSON.stringify({ message: `Reloaded: ${knowledgeBase.length} chunks` }));
     }
 
+    // ── Existing non-streaming endpoint (unchanged behavior) ──────────────
     if (pathname === '/generate' && req.method === 'POST') {
         let body = '';
         req.on('data', c => { body += c.toString(); });
@@ -121,43 +180,9 @@ const server = http.createServer(async (req, res) => {
 
                 console.log(`Query: "${prompt}"`);
 
-                // ── RAG: find relevant knowledge chunks ───────────────────
-                const ragContext = retrieveContext(prompt, 5);
-                console.log(`RAG chunks found: ${ragContext.length} chars`);
-
-                // ── Build system prompt with knowledge injected ────────────
-                // IMPORTANT: knowledge comes FIRST, before any other instruction
-                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.`;
-
-                const messages = [
-                    { role: 'system', content: systemPrompt },
-                    { role: 'user',   content: prompt        }
-                ];
-
+                const { messages, ragContext } = buildMessages(prompt);
                 let result = await generateResponse(messages);
 
-                // If still empty, return the raw knowledge chunk directly
                 if (!result || result.trim().length < 10) {
                     result = ragContext
                         ? `Here is what my knowledge base says:\n\n${ragContext.slice(0, 800)}`
@@ -177,6 +202,64 @@ ONLY answer agriculture questions.`;
         return;
     }
 
+    // ── NEW: streaming endpoint (SSE) ──────────────────────────────────────
+    if (pathname === '/generate-stream' && 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`);
+                });
+
+                // 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`);
+                }
+
+                res.write('data: [DONE]\n\n');
+                res.end();
+
+            } 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.end();
+            }
+        });
+        return;
+    }
+
     res.setHeader('Content-Type', 'application/json');
     res.writeHead(404);
     res.end(JSON.stringify({ error: "Not Found" }));