Update aiEngine.js

aiEngine.js

@@ -6,113 +6,174 @@ dotenv.config();
 
 const REMOTE_SERVER_URL = process.env.REMOTE_AI_URL || "http://localhost:7860";
 
+// --- PROMPT LOADING ---
 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); }
+    if (fs.existsSync(promptsPath)) {
+        prompts = JSON.parse(fs.readFileSync(promptsPath, 'utf8'));
+    }
+} catch (e) { 
+    console.error("Prompt Load Error:", e); 
+}
 
+// --- HISTORY FLATTENER (WITH ECONOMIC LIMITS) ---
 const flattenHistory = (history, currentInput, systemPrompt, limit = 10, gdd = null) => {
+    // ECONOMIC CAP: Slice history to the last 'limit' messages to save tokens
     const recentHistory = history.slice(-limit); 
+    
     let context = recentHistory.map(m => {
         const roleName = m.role === 'model' ? 'Assistant' : 'User';
         return `${roleName}: ${m.parts[0].text}`;
     }).join('\n');
-    const projectAnchor = gdd ? `[GDD REFERENCE]:\n${gdd}\n\n` : "";
+
+    // Inject GDD only if provided (usually for PM)
+    const projectAnchor = gdd ? `[PROJECT GDD REFERENCE]:\n${gdd}\n\n` : "";
+    
     return `System: ${systemPrompt}\n\n${projectAnchor}${context}\nUser: ${currentInput}\nAssistant:`;
 };
 
+// --- STREAM HANDLER & USAGE PARSER ---
 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 fullStreamData = "";
 
     while (true) {
         const { done, value } = await reader.read();
         if (done) break;
+        
         const chunk = decoder.decode(value, { stream: true });
         fullStreamData += chunk;
 
-        if (chunk.includes("__THINK__")) {
-            const parts = chunk.split("__THINK__");
-            if (parts[0]) onOutput?.(parts[0]);
-            if (parts[1]) onThink?.(parts[1]);
-        } else if (!chunk.includes("__USAGE__")) {
-            onOutput?.(chunk);
+        // Streaming Logic: Don't show Usage to frontend, parse thoughts
+        if (!chunk.includes("__USAGE__")) {
+            if (chunk.includes("__THINK__")) {
+                const parts = chunk.split("__THINK__");
+                if (parts[0] && onOutput) onOutput(parts[0]);
+                if (parts[1] && onThink) onThink(parts[1]);
+            } else {
+                if (onOutput) onOutput(chunk);
+            }
         }
     }
     
-    let usage = { inputTokens: 0, outputTokens: 0, totalTokenCount: 0 };
-    let finalOutput = fullStreamData;
+    // --- USAGE EXTRACTION FOR BILLING ---
+    let usage = { totalTokenCount: 0, inputTokens: 0, outputTokens: 0 };
+    let finalCleanText = fullStreamData;
 
     if (fullStreamData.includes("__USAGE__")) {
         const parts = fullStreamData.split("__USAGE__");
-        finalOutput = parts[0];
+        finalCleanText = parts[0]; // The actual text content
+        const usageRaw = parts[1];
+        
         try { 
-            const rawUsage = JSON.parse(parts[1]);
-            // Standardize usage object
-            usage.totalTokenCount = rawUsage.totalTokenCount || 0;
-            // If the remote server provides a breakdown, use it; otherwise estimate 80/20 split
-            usage.inputTokens = rawUsage.inputTokens || Math.floor(usage.totalTokenCount * 0.8);
-            usage.outputTokens = rawUsage.outputTokens || Math.ceil(usage.totalTokenCount * 0.2);
-        } catch (e) { console.warn("Usage Parse Failed"); }
+            const parsedUsage = JSON.parse(usageRaw);
+            usage.totalTokenCount = parsedUsage.totalTokenCount || 0;
+            usage.inputTokens = parsedUsage.inputTokens || 0;
+            usage.outputTokens = parsedUsage.outputTokens || 0;
+        } catch (e) { 
+            console.warn("Usage Parse Failed in Engine:", e); 
+        }
     }
 
-    const cleanText = finalOutput.split("__THINK__")[0].trim();
-    return { text: cleanText, usage };
+    // Clean any remaining tags
+    finalCleanText = finalCleanText.split("__THINK__")[0].trim();
+
+    return { text: finalCleanText, usage };
 };
 
 export const AIEngine = {
+    // --- STREAMING METHODS (Main Loop) ---
+
     callPMStream: async (history, input, onThink, onOutput, gdd = null) => {
-        const prompt = flattenHistory(history, input, prompts.pm_system_prompt, 15, gdd); // Reduced to 15 for cost
+        const systemPrompt = prompts.pm_system_prompt || "You are a Project Manager.";
+        // ECONOMIC CAP: 15 messages max for PM to maintain context but control costs
+        const prompt = flattenHistory(history, input, systemPrompt, 15, gdd); 
+        
         const response = await fetch(`${REMOTE_SERVER_URL}/api/stream`, {
             method: 'POST',
             headers: { 'Content-Type': 'application/json' },
-            body: JSON.stringify({ model: "claude", prompt: prompt, system_prompt: prompts.pm_system_prompt })
+            body: JSON.stringify({ 
+                model: "claude", 
+                prompt: prompt, 
+                system_prompt: systemPrompt 
+            })
         });
         return await handleStreamResponse(response, onThink, onOutput);
     },
 
     callWorkerStream: async (history, input, onThink, onOutput, images = []) => {
-        const prompt = flattenHistory(history, input, prompts.worker_system_prompt, 8, null); // Reduced to 8 for cost
+        const systemPrompt = prompts.worker_system_prompt || "You are a Senior Engineer.";
+        // ECONOMIC CAP: 8 messages max for Worker (they only need recent context)
+        const prompt = flattenHistory(history, input, systemPrompt, 8, null); 
+        
         const response = await fetch(`${REMOTE_SERVER_URL}/api/stream`, {
             method: 'POST',
             headers: { 'Content-Type': 'application/json' },
-            body: JSON.stringify({ model: "gpt", prompt: prompt, system_prompt: prompts.worker_system_prompt, images })
+            body: JSON.stringify({ 
+                model: "gpt", 
+                prompt: prompt, 
+                system_prompt: systemPrompt,
+                images: images
+            })
         });
         return await handleStreamResponse(response, onThink, onOutput);
     },
 
+    // --- BLOCKING CALLS (Background Initialization) ---
+
     callPM: async (history, input, gdd = null) => {
-        const prompt = flattenHistory(history, input, prompts.pm_system_prompt, 15, gdd);
+        const systemPrompt = prompts.pm_system_prompt || "You are a Project Manager.";
+        const prompt = flattenHistory(history, input, systemPrompt, 15, gdd); // Limit 15
+
         const response = await fetch(`${REMOTE_SERVER_URL}/api/generate`, {
             method: 'POST',
             headers: { 'Content-Type': 'application/json' },
-            body: JSON.stringify({ model: "claude", prompt: prompt, system_prompt: prompts.pm_system_prompt })
+            body: JSON.stringify({ 
+                model: "claude", 
+                prompt: prompt, 
+                system_prompt: systemPrompt 
+            })
         });
         const result = await response.json();
         return { text: result.data, usage: result.usage || { totalTokenCount: 0 } };
     },
 
     callWorker: async (history, input) => {
-        const prompt = flattenHistory(history, input, prompts.worker_system_prompt, 8, null);
+        const systemPrompt = prompts.worker_system_prompt || "You are a Senior Engineer.";
+        const prompt = flattenHistory(history, input, systemPrompt, 8, null); // Limit 8
+
         const response = await fetch(`${REMOTE_SERVER_URL}/api/generate`, {
             method: 'POST',
             headers: { 'Content-Type': 'application/json' },
-            body: JSON.stringify({ model: "gpt", prompt: prompt, system_prompt: prompts.worker_system_prompt })
+            body: JSON.stringify({ 
+                model: "gpt", 
+                prompt: prompt, 
+                system_prompt: systemPrompt 
+            })
         });
         const result = await response.json();
         return { text: result.data, usage: result.usage || { totalTokenCount: 0 } };
     },
 
+    // --- UTILITIES (One-off calls) ---
+
     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 })
+            body: JSON.stringify({ 
+                model: "gpt", 
+                prompt: `Analyze this project idea: ${desc}`, 
+                system_prompt: prompts.analyst_system_prompt || "Output JSON only."
+            })
         });
         const result = await response.json();
+        // Return parsed data AND usage for billing
         return { ...JSON.parse(result.data), usage: result.usage };
     },
 
@@ -120,293 +181,30 @@ export const AIEngine = {
         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 })
+            body: JSON.stringify({ 
+                model: "gpt", 
+                prompt: `Grade this project. Description: ${desc} Answers: ${JSON.stringify(ans)}`, 
+                system_prompt: prompts.analyst_system_prompt || "Output JSON only."
+            })
         });
         const result = await response.json();
         const parsed = JSON.parse(result.data);
-        parsed.usage = result.usage;
+        parsed.usage = result.usage; // Attach usage for billing
         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) => {
-    if (!response.ok) throw new Error(`Stream Error: ${response.statusText}`);
-
-    const reader = response.body.getReader();
-    const decoder = new TextDecoder("utf-8");
-    
-    let fullText = ""; 
-    let usage = { totalTokenCount: 0 }; // Default
-
-    while (true) {
-        const { done, value } = await reader.read();
-        if (done) break;
-        
-        let chunk = decoder.decode(value, { stream: true });
-
-        // CHECK FOR USAGE FOOTER
-        if (chunk.includes("__USAGE__")) {
-            const parts = chunk.split("__USAGE__");
-            chunk = parts[0]; // The text part
-            try {
-                if (parts[1]) {
-                    usage = JSON.parse(parts[1]);
-                }
-            } catch (e) {
-                console.warn("Failed to parse usage footer:", e);
-            }
-        }
-        
-        // STANDARD CHUNK PROCESSING
-        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]) {
-                if (onOutput) onOutput(parts[0]);
-                fullText += parts[0];
-            }
-            if (parts[1] && onThink) {
-                onThink(parts[1]);
-            }
-        } 
-        else {
-            if (onOutput) onOutput(chunk);
-            fullText += chunk;
-        }
-    }
-    
-    return { 
-        text: fullText, 
-        usage: usage 
-    };
-};
-*/
-/*
-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 = ""; 
-    let usage = { totalTokenCount: 0 }; // Initialize usage object
-
-    while (true) {
-        const { done, value } = await reader.read();
-        if (done) break;
-        
-        let chunk = decoder.decode(value, { stream: true });
-
-        // --- CREDIT LOGIC: Parse Usage Tag ---
-        if (chunk.includes("__USAGE__")) {
-            const parts = chunk.split("__USAGE__");
-            chunk = parts[0]; // Process the text before the tag
-            try {
-                if (parts[1]) {
-                    usage = JSON.parse(parts[1]); // Capture the usage from the remote server
-                }
-            } catch (e) {
-                console.warn("Failed to parse usage footer");
-            }
-        }
-        
-        // --- THOUGHT/OUTPUT SPLITTING (Your existing logic) ---
-        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]) {
-                if (onOutput) onOutput(parts[0]);
-                fullText += parts[0];
-            }
-            if (parts[1] && onThink) onThink(parts[1]);
-        } else {
-            if (onOutput) onOutput(chunk);
-            fullText += chunk;
-        }
-    }
-    
-    // Return both the text and the usage data
-    return { text: fullText, usage: usage }; 
-};
-export const AIEngine = {
-    // --- PM STREAMING ---
-    callPMStream: async (history, input, 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 fullPrompt = flattenHistory(history, input, "");
-
-        try {
-            const response = await fetch(`${REMOTE_SERVER_URL}/api/stream`, {
+    
+    generateImage: async (prompt) => {
+         try {
+            const response = await fetch(`${REMOTE_SERVER_URL}/api/image`, {
                 method: 'POST',
                 headers: { 'Content-Type': 'application/json' },
-                body: JSON.stringify({
-                    model: "gpt",
-                    prompt: fullPrompt,
-                    system_prompt: systemPrompt,
-                    images: images 
-                })
+                body: JSON.stringify({ prompt })
             });
-
-            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: "claude",
-                prompt: fullPrompt,
-                system_prompt: systemPrompt
-            })
-        });
-
-        const result = await response.json();
-        if (!result.success) throw new Error(result.error);
-       
-        return { 
-            text: result.data, 
-            usage: result.usage || { totalTokenCount: 0 } 
-        };
-      } catch (error) {
-        console.log("PM error: ",error);
-        return { text: "", error };
-      }
-    },
-
-    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 response.json();
-        if (!result.success) throw new Error(result.error);
-       
-        return { 
-            text: result.data, 
-            usage: result.usage || { totalTokenCount: 0 } 
-        };
-      } catch (error) {
-        console.log("Worker error: ",error);
-        return { text: "", error };
-      }
-    },
-
-    generateEntryQuestions: async (desc) => {
-           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 response.json();
-             
-        return {
-            ...JSON.parse(result.data),
-            usage: result.usage || { totalTokenCount: 0 }
-        };
-
-      } catch (error) {
-        console.log("GenerateQ error: ",error)
-      }
-    },
-
-    gradeProject: async (desc, ans) => {
-           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 response.json();
-        const parsed = JSON.parse(result.data);
-        parsed.usage = result.usage || { totalTokenCount: 0 };
-        return parsed;
-
-      } catch (error) {
-        console.log("GenerateQ error: ",error)
-      }
+            const result = await response.json();
+            return result; // Expected { image: "base64..." }
+         } catch (e) {
+             console.error("Image Gen Error:", e);
+             return null;
+         }
     }
-}; */
+};