Update aiEngine.js

aiEngine.js

@@ -3,6 +3,7 @@ import fs from 'fs';
 import path from 'path';
 
 dotenv.config();
+
 const REMOTE_SERVER_URL = process.env.REMOTE_AI_URL || "http://localhost:7860";
 
 let prompts = {};
@@ -19,73 +20,184 @@ const flattenHistory = (history, currentInput, systemPrompt) => {
     return `System: ${systemPrompt}\n\n${context}\nUser: ${currentInput}\nAssistant:`;
 };
 
+// HELPER: STREAM SPLITTER
 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");
     let fullText = "";
 
-    try {
-        while (true) {
-            const { done, value } = await reader.read();
-            if (done) break;
-            const chunk = decoder.decode(value, { stream: true });
-            if (chunk.startsWith("__THINK__")) {
-                if (onThink) onThink(chunk.replace("__THINK__", ""));
-            } else if (chunk.includes("__THINK__")) {
-                const parts = chunk.split("__THINK__");
-                if (parts[0] && onOutput) { onOutput(parts[0]); fullText += parts[0]; }
-                if (parts[1] && onThink) onThink(parts[1]);
-            } else {
-                if (onOutput) onOutput(chunk);
-                fullText += chunk;
+    while (true) {
+        const { done, value } = await reader.read();
+        if (done) break;
+        
+        const chunk = decoder.decode(value, { stream: true });
+        
+        if (chunk.startsWith("__THINK__")) {
+            const thoughtContent = chunk.replace("__THINK__", "");
+            if (onThink) onThink(thoughtContent);
+        } else if (chunk.includes("__THINK__")) {
+            const parts = chunk.split("__THINK__");
+            if (parts[0] && onOutput) {
+                onOutput(parts[0]);
+                fullText += parts[0];
+            }
+            if (parts[1] && onThink) {
+                onThink(parts[1]);
             }
+        } else {
+            if (onOutput) onOutput(chunk);
+            fullText += chunk;
         }
-    } catch (e) { console.error("Stream reader error:", e.message); }
+    }
     return { text: fullText, usage: { totalTokenCount: 0 } };
 };
 
 export const AIEngine = {
+    // --- PM STREAMING ---
     callPMStream: async (history, input, onThink, onOutput) => {
-        const systemPrompt = prompts.pm_system_prompt || "You are a lead lead architect.";
-        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: systemPrompt })
-        });
-        return await handleStreamResponse(response, onThink, onOutput);
+        const systemPrompt = prompts.pm_system_prompt || "You are a pro manager.";
+        const fullPrompt = flattenHistory(history, input, "");
+
+        try {
+            const response = await fetch(`${REMOTE_SERVER_URL}/api/stream`, {
+                method: 'POST',
+                headers: { 'Content-Type': 'application/json' },
+                body: JSON.stringify({
+                    model: "claude",
+                    prompt: fullPrompt,
+                    system_prompt: systemPrompt
+                })
+            });
+
+            return await handleStreamResponse(response, onThink, onOutput);
+
+        } catch (error) {
+            console.log("PM Stream error: ", error);
+            throw error;
+        }
     },
+
+    // --- WORKER STREAMING (WITH IMAGES) ---
     callWorkerStream: async (history, input, onThink, onOutput, images = []) => {
         const systemPrompt = prompts.worker_system_prompt || "You are a worker.";
-        const response = await fetch(`${REMOTE_SERVER_URL}/api/stream`, {
+        const fullPrompt = flattenHistory(history, input, "");
+
+        try {
+            const response = await fetch(`${REMOTE_SERVER_URL}/api/stream`, {
+                method: 'POST',
+                headers: { 'Content-Type': 'application/json' },
+                body: JSON.stringify({
+                    model: "gpt",
+                    prompt: fullPrompt,
+                    system_prompt: systemPrompt,
+                    images: images // Pass images to Main Agent
+                })
+            });
+
+            return await handleStreamResponse(response, onThink, onOutput);
+
+        } catch (error) {
+            console.log("Worker Stream error: ", error);
+            throw error;
+        }
+    },
+
+    // --- LEGACY BLOCKING CALLS ---
+    callPM: async (history, input) => {
+        const systemPrompt = prompts.pm_system_prompt || "You are a pro manager.";
+        const fullPrompt = flattenHistory(history, input, "" );
+
+      try {
+        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: systemPrompt, images })
+            body: JSON.stringify({
+                model: "claude",
+                prompt: fullPrompt,
+                system_prompt: systemPrompt
+            })
         });
-        return await handleStreamResponse(response, onThink, onOutput);
+
+        const result = await response.json();
+        if (!result.success) throw new Error(result.error);
+       
+        return { 
+            text: result.data, 
+            usage: { totalTokenCount: 0 } 
+        };
+      } catch (error) {
+        console.log("PM error: ",error);
+        return { text: "", error };
+      }
     },
-    callPM: async (history, input) => {
-        const res = 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 })
+
+    callWorker: async (history, input, images = []) => {
+        const systemPrompt = prompts.worker_system_prompt || "You are a worker.";
+        const fullPrompt = flattenHistory(history, input, "" );
+
+      try {
+        const response = await fetch(`${REMOTE_SERVER_URL}/api/generate`, {
+            method: 'POST',
+            headers: { 'Content-Type': 'application/json' },
+            body: JSON.stringify({
+                model: "gpt",
+                prompt: fullPrompt,
+              system_prompt: systemPrompt
+            })
         });
-        const result = await res.json();
-        return { text: result.data, usage: { totalTokenCount: 0 } };
+
+        const result = await response.json();
+        if (!result.success) throw new Error(result.error);
+       
+        return { 
+            text: result.data, 
+            usage: { totalTokenCount: 0 } 
+        };
+      } catch (error) {
+        console.log("Worker error: ",error);
+        return { text: "", error };
+      }
     },
+
     generateEntryQuestions: async (desc) => {
-        const res = await fetch(`${REMOTE_SERVER_URL}/api/generate`, {
-            method: 'POST', headers: { 'Content-Type': 'application/json' },
-            body: JSON.stringify({ model: "gpt", prompt: `[OUTPUT ONLY JSON]\n Questions: ${desc}`, system_prompt: prompts.analyst_system_prompt })
+           try {
+        const response = await fetch(`${REMOTE_SERVER_URL}/api/generate`, {
+            method: 'POST',
+            headers: { 'Content-Type': 'application/json' },
+            body: JSON.stringify({
+                model: "gpt",
+                prompt: `[OUTPUT ONLY JSON]\n Generate entry questions for this idea: ${desc}`,
+                system_prompt: `Goal: ${prompts.analyst_system_prompt}`
+            })
         });
-        const result = await res.json();
-        return JSON.parse(result.data);
+
+        const result = await response.json();
+             return JSON.parse(result.data);
+
+      } catch (error) {
+        console.log("GenerateQ error: ",error)
+      }
     },
+
     gradeProject: async (desc, ans) => {
-        const res = await fetch(`${REMOTE_SERVER_URL}/api/generate`, {
-            method: 'POST', headers: { 'Content-Type': 'application/json' },
-            body: JSON.stringify({ model: "gpt", prompt: `[OUTPUT ONLY JSON]\n Grade: ${desc}\nAns: ${JSON.stringify(ans)}`, system_prompt: prompts.analyst_system_prompt })
+           try {
+        const response = await fetch(`${REMOTE_SERVER_URL}/api/generate`, {
+            method: 'POST',
+            headers: { 'Content-Type': 'application/json' },
+            body: JSON.stringify({
+                model: "gpt",
+                prompt: `[OUTPUT ONLY JSON]\n. Grade this project. Desc: ${desc}\nAnswers: ${JSON.stringify(ans)}`,
+                system_prompt: prompts.analyst_system_prompt
+            })
         });
-        const result = await res.json();
-        return JSON.parse(result.data);
+
+        const result = await response.json();
+             return JSON.parse(result.data);
+
+      } catch (error) {
+        console.log("GenerateQ error: ",error)
+      }
     }
 };