Update server.js

server.js

@@ -9,7 +9,7 @@ const KNOWLEDGE_DIR = './knowledge';
 let generator;
 let knowledgeBase = [];
 
-// ── LOAD KNOWLEDGE FILES ──────────────────────────────────────────────────────
+// ── KNOWLEDGE FILES ───────────────────────────────────────────────────────────
 function loadKnowledge() {
     if (!fs.existsSync(KNOWLEDGE_DIR)) { fs.mkdirSync(KNOWLEDGE_DIR); return; }
     const files = fs.readdirSync(KNOWLEDGE_DIR).filter(f => f.endsWith('.txt'));
@@ -19,7 +19,7 @@ function loadKnowledge() {
         const chunks  = splitChunks(content, 600, 60);
         chunks.forEach(c => knowledgeBase.push({ source: file, text: c }));
     }
-    console.log(`Knowledge: ${files.length} files, ${knowledgeBase.length} chunks`);
+    console.log(`Knowledge loaded: ${files.length} files, ${knowledgeBase.length} chunks`);
 }
 
 function splitChunks(text, size, overlap) {
@@ -47,7 +47,7 @@ function retrieveContext(prompt, topK = 4) {
         .join('\n\n---\n\n');
 }
 
-// ── LOAD MODEL ────────────────────────────────────────────────────────────────
+// ── MODEL ─────────────────────────────────────────────────────────────────────
 async function loadModel() {
     console.log("Loading model...");
     generator = await pipeline('text-generation', MODEL_NAME, { dtype: 'q4' });
@@ -61,7 +61,12 @@ async function generateResponse(messages) {
         repetition_penalty: 1.2,
         do_sample:          false
     });
-    return output[0].generated_text.at(-1).content || '';
+    // Extract only the assistant reply content
+    const generated = output[0].generated_text;
+    if (Array.isArray(generated)) {
+        return generated.at(-1)?.content || '';
+    }
+    return String(generated || '');
 }
 
 // ── SERVER ────────────────────────────────────────────────────────────────────
@@ -74,12 +79,18 @@ const server = http.createServer(async (req, res) => {
 
     const pathname = req.url.split('?')[0];
 
+    // Status check
     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 }));
+        return res.end(JSON.stringify({
+            status: "running",
+            model: MODEL_NAME,
+            knowledge_chunks: knowledgeBase.length
+        }));
     }
 
+    // Reload knowledge
     if (pathname === '/reload-knowledge' && req.method === 'POST') {
         loadKnowledge();
         res.setHeader('Content-Type', 'application/json');
@@ -87,66 +98,57 @@ const server = http.createServer(async (req, res) => {
         return res.end(JSON.stringify({ message: `Reloaded: ${knowledgeBase.length} chunks` }));
     }
 
+    // Main generate endpoint — returns plain JSON (no SSE, no streaming delays)
     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, system } = JSON.parse(body);
+
                 if (!generator) {
-                    res.setHeader('Content-Type', 'application/json');
                     res.writeHead(503);
-                    return res.end(JSON.stringify({ error: "Model loading..." }));
+                    return res.end(JSON.stringify({ error: "Model still loading, please wait..." }));
                 }
 
-                // Retrieve from knowledge files
+                // RAG: retrieve relevant knowledge chunks
                 const ragContext = retrieveContext(prompt, 4);
                 const ragSection = ragContext
-                    ? `\n\nKNOWLEDGE BASE (use this for your answer):\n${ragContext}\n`
+                    ? `\n\nKNOWLEDGE BASE — use ONLY this information to answer:\n${ragContext}\n`
                     : '';
 
-                // Use system prompt from frontend if provided, else default
-                const finalSystem = system
-                    ? system + ragSection
-                    : `You are Mlimi Connect AI, a free agricultural advisor for Malawian farmers. Only answer agriculture questions.${ragSection}`;
+                // Build final system prompt
+                const finalSystem = (system || `You are Mlimi Connect AI, a free agricultural advisor for Malawian farmers. Only answer agriculture questions.`) + ragSection;
 
                 const messages = [
                     { role: 'system', content: finalSystem },
                     { role: 'user',   content: prompt       }
                 ];
 
-                // SSE streaming
-                res.writeHead(200, {
-                    'Content-Type':    'text/event-stream',
-                    'Cache-Control':   'no-cache',
-                    'Connection':      'keep-alive',
-                    'X-Accel-Buffering': 'no'
-                });
-
-                res.write(`data: ${JSON.stringify({ status: "generating" })}\n\n`);
+                console.log(`Generating response for: "${prompt.slice(0, 60)}..."`);
 
                 let result = await generateResponse(messages);
 
-                if (!result || result.trim().length < 3) {
-                    result = ragContext
-                        ? `Based on available knowledge:\n\n${ragContext.slice(0, 400)}`
-                        : "I don't have specific information on that. Please consult your local agricultural extension officer.";
+                // Fallback if model returns empty
+                if (!result || result.trim().length < 5) {
+                    if (ragContext) {
+                        result = `Here is what I know about this topic:\n\n${ragContext.slice(0, 600)}`;
+                    } else {
+                        result = "I don't have specific information on that topic. Please consult your local agricultural extension officer (AEO) for advice.";
+                    }
                 }
 
-                // Stream word by word
-                const words = result.split(' ');
-                for (let i = 0; i < words.length; i++) {
-                    const chunk = (i === 0 ? '' : ' ') + words[i];
-                    res.write(`data: ${JSON.stringify({ chunk })}\n\n`);
-                    await new Promise(r => setTimeout(r, 12));
-                }
+                console.log(`Response ready: ${result.length} chars`);
 
-                res.write(`data: ${JSON.stringify({ done: true, result })}\n\n`);
-                res.end();
+                res.writeHead(200);
+                res.end(JSON.stringify({ result }));
 
             } catch (err) {
-                res.write(`data: ${JSON.stringify({ error: err.message })}\n\n`);
-                res.end();
+                console.error("Generation error:", err.message);
+                res.writeHead(500);
+                res.end(JSON.stringify({ error: err.message || "Generation failed" }));
             }
         });
         return;