Update aiEngine.js

aiEngine.js

@@ -20,6 +20,145 @@ const flattenHistory = (history, currentInput, systemPrompt) => {
     return `System: ${systemPrompt}\n\n${context}\nUser: ${currentInput}\nAssistant:`;
 };
 
+// --- ROBUST STREAM HANDLER ---
+const handleStreamResponse = async (response, onThink, onOutput) => {
+    if (!response.ok) throw new Error(`Stream Error: ${response.statusText}`);
+
+    const reader = response.body.getReader();
+    const decoder = new TextDecoder("utf-8");
+    
+    // Buffer the entire stream data to prevent split-parsing issues
+    let fullStreamData = "";
+
+    while (true) {
+        const { done, value } = await reader.read();
+        if (done) break;
+        
+        const chunk = decoder.decode(value, { stream: true });
+        fullStreamData += chunk;
+
+        // Still send live feedback to frontend, but we will parse from the buffer later
+        if (chunk.startsWith("__THINK__")) {
+            if (onThink) onThink(chunk.replace("__THINK__", ""));
+        } else if (chunk.includes("__THINK__")) {
+            const parts = chunk.split("__THINK__");
+            if (parts[0]) onOutput?.(parts[0]);
+            if (parts[1]) onThink?.(parts[1]);
+        } else {
+            if (chunk) onOutput?.(chunk);
+        }
+    }
+    
+    let usage = { totalTokenCount: 0 };
+    let fullText = fullStreamData;
+
+    // --- PARSE FROM THE COMPLETE BUFFER ---
+    // This is safer because the tag and JSON are guaranteed to be complete.
+    if (fullStreamData.includes("__USAGE__")) {
+        const parts = fullStreamData.split("__USAGE__");
+        fullText = parts[0]; // The actual text response
+        try {
+            if (parts[1]) {
+                usage = JSON.parse(parts[1]);
+            }
+        } catch (e) {
+            console.warn("Failed to parse final usage footer from stream buffer.", e);
+        }
+    }
+
+    // Now, we must strip any thought blocks from the final text to be saved
+    fullText = fullText.split("__THINK__")[0];
+    
+    return { text: fullText.trim(), usage };
+};
+
+export const AIEngine = {
+    callPMStream: async (history, input, onThink, onOutput) => {
+        const response = await fetch(`${REMOTE_SERVER_URL}/api/stream`, {
+            method: 'POST',
+            headers: { 'Content-Type': 'application/json' },
+            body: JSON.stringify({ model: "claude", prompt: flattenHistory(history, input, ""), system_prompt: prompts.pm_system_prompt })
+        });
+        return await handleStreamResponse(response, onThink, onOutput);
+    },
+
+    callWorkerStream: async (history, input, onThink, onOutput, images = []) => {
+        const response = await fetch(`${REMOTE_SERVER_URL}/api/stream`, {
+            method: 'POST',
+            headers: { 'Content-Type': 'application/json' },
+            body: JSON.stringify({ model: "gpt", prompt: flattenHistory(history, input, ""), system_prompt: prompts.worker_system_prompt, images })
+        });
+        return await handleStreamResponse(response, onThink, onOutput);
+    },
+
+    callPM: async (history, input) => {
+        const response = await fetch(`${REMOTE_SERVER_URL}/api/generate`, {
+            method: 'POST',
+            headers: { 'Content-Type': 'application/json' },
+            body: JSON.stringify({ model: "claude", prompt: flattenHistory(history, input, ""), system_prompt: prompts.pm_system_prompt })
+        });
+        const result = await response.json();
+        return { text: result.data, usage: result.usage || { totalTokenCount: 0 } };
+    },
+
+    callWorker: async (history, input) => {
+        const response = await fetch(`${REMOTE_SERVER_URL}/api/generate`, {
+            method: 'POST',
+            headers: { 'Content-Type': 'application/json' },
+            body: JSON.stringify({ model: "gpt", prompt: flattenHistory(history, input, ""), system_prompt: prompts.worker_system_prompt })
+        });
+        const result = await response.json();
+        return { text: result.data, usage: result.usage || { totalTokenCount: 0 } };
+    },
+
+    generateEntryQuestions: async (desc) => {
+        const response = await fetch(`${REMOTE_SERVER_URL}/api/generate`, {
+            method: 'POST',
+            headers: { 'Content-Type': 'application/json' },
+            body: JSON.stringify({ model: "gpt", prompt: `Analyze: ${desc}`, system_prompt: prompts.analyst_system_prompt })
+        });
+        const result = await response.json();
+        const parsed = JSON.parse(result.data);
+        parsed.usage = result.usage;
+        return parsed;
+    },
+
+    gradeProject: async (desc, ans) => {
+        const response = await fetch(`${REMOTE_SERVER_URL}/api/generate`, {
+            method: 'POST',
+            headers: { 'Content-Type': 'application/json' },
+            body: JSON.stringify({ model: "gpt", prompt: `Grade: ${desc} ${JSON.stringify(ans)}`, system_prompt: prompts.analyst_system_prompt })
+        });
+        const result = await response.json();
+        const parsed = JSON.parse(result.data);
+        parsed.usage = result.usage;
+        return parsed;
+    }
+};
+
+
+/*import dotenv from 'dotenv';
+import fs from 'fs';
+import path from 'path';
+
+dotenv.config();
+
+const REMOTE_SERVER_URL = process.env.REMOTE_AI_URL || "http://localhost:7860";
+
+let prompts = {};
+try {
+    const promptsPath = path.resolve('./prompts.json');
+    if (fs.existsSync(promptsPath)) prompts = JSON.parse(fs.readFileSync(promptsPath, 'utf8'));
+} catch (e) { console.error("Prompt Load Error:", e); }
+
+const flattenHistory = (history, currentInput, systemPrompt) => {
+    const context = history.map(m => {
+        const roleName = m.role === 'model' ? 'Assistant' : 'User';
+        return `${roleName}: ${m.parts[0].text}`;
+    }).join('\n');
+    return `System: ${systemPrompt}\n\n${context}\nUser: ${currentInput}\nAssistant:`;
+};
+
 /*
 // HELPER: STREAM SPLITTER & USAGE PARSER
 const handleStreamResponse = async (response, onThink, onOutput) => {
@@ -77,6 +216,7 @@ const handleStreamResponse = async (response, onThink, onOutput) => {
     };
 };
 */
+/*
 const handleStreamResponse = async (response, onThink, onOutput) => {
     if (!response.ok) throw new Error(`Stream Error: ${response.statusText}`);
 
@@ -277,4 +417,4 @@ export const AIEngine = {
         console.log("GenerateQ error: ",error)
       }
     }
-};
+}; */