Update aiEngine.js

aiEngine.js

@@ -20,50 +20,60 @@ const flattenHistory = (history, currentInput, systemPrompt) => {
     return `System: ${systemPrompt}\n\n${context}\nUser: ${currentInput}\nAssistant:`;
 };
 
-// HELPER: STREAM SPLITTER
+// 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");
     
-    // fullText is what gets saved to the DB. It should NOT contain thoughts.
     let fullText = ""; 
+    let usage = { totalTokenCount: 0 }; // Default
 
     while (true) {
         const { done, value } = await reader.read();
         if (done) break;
         
-        const chunk = decoder.decode(value, { stream: true });
+        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);
+            }
+        }
         
-        // CASE 1: Pure Thought Chunk
+        // STANDARD CHUNK PROCESSING
         if (chunk.startsWith("__THINK__")) {
             const thoughtContent = chunk.replace("__THINK__", "");
             if (onThink) onThink(thoughtContent);
         } 
-        // CASE 2: Mixed Chunk (Output then Thought)
         else if (chunk.includes("__THINK__")) {
             const parts = chunk.split("__THINK__");
-            
-            // part[0] is real output
             if (parts[0]) {
                 if (onOutput) onOutput(parts[0]);
                 fullText += parts[0];
             }
-            
-            // part[1] is thought
             if (parts[1] && onThink) {
                 onThink(parts[1]);
             }
         } 
-        // CASE 3: Pure Output Chunk
         else {
             if (onOutput) onOutput(chunk);
             fullText += chunk;
         }
     }
     
-    return { text: fullText, usage: { totalTokenCount: 0 } };
+    return { 
+        text: fullText, 
+        usage: usage 
+    };
 };
 
 export const AIEngine = {
@@ -104,7 +114,7 @@ export const AIEngine = {
                     model: "gpt",
                     prompt: fullPrompt,
                     system_prompt: systemPrompt,
-                    images: images // Pass images to Main Agent
+                    images: images 
                 })
             });
 
@@ -137,7 +147,7 @@ export const AIEngine = {
        
         return { 
             text: result.data, 
-            usage: { totalTokenCount: 0 } 
+            usage: result.usage || { totalTokenCount: 0 } 
         };
       } catch (error) {
         console.log("PM error: ",error);
@@ -165,7 +175,7 @@ export const AIEngine = {
        
         return { 
             text: result.data, 
-            usage: { totalTokenCount: 0 } 
+            usage: result.usage || { totalTokenCount: 0 } 
         };
       } catch (error) {
         console.log("Worker error: ",error);
@@ -186,7 +196,11 @@ export const AIEngine = {
         });
 
         const result = await response.json();
-             return JSON.parse(result.data);
+             
+        return {
+            ...JSON.parse(result.data),
+            usage: result.usage || { totalTokenCount: 0 }
+        };
 
       } catch (error) {
         console.log("GenerateQ error: ",error)
@@ -206,7 +220,9 @@ export const AIEngine = {
         });
 
         const result = await response.json();
-             return JSON.parse(result.data);
+        const parsed = JSON.parse(result.data);
+        parsed.usage = result.usage || { totalTokenCount: 0 };
+        return parsed;
 
       } catch (error) {
         console.log("GenerateQ error: ",error)