modified_ai_server

Paused

Update aiEngine.js

8d5582b verified 8 months ago

8.21 kB

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";

// --- 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); 
}

// --- 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');

    // 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;

        // 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);
            }
        }
    }
    
    // --- USAGE EXTRACTION FOR BILLING ---
    let usage = { totalTokenCount: 0, inputTokens: 0, outputTokens: 0 };
    let finalCleanText = fullStreamData;

    if (fullStreamData.includes("__USAGE__")) {
        const parts = fullStreamData.split("__USAGE__");
        finalCleanText = parts[0]; // The actual text content
        const usageRaw = parts[1];
        
        try { 
            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); 
        }
    }

    // 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 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,
                                      10,
                                      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: systemPrompt 
            })
        });
        return await handleStreamResponse(response, onThink, onOutput);
    },

    callWorkerStream: async (history, input, onThink, onOutput, images = []) => {
        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: systemPrompt,
                images: images
            })
        });
        return await handleStreamResponse(response, onThink, onOutput);
    },

    // --- BLOCKING CALLS (Background Initialization) ---

    callPM: async (history, input, gdd = null) => {
        const systemPrompt = prompts.pm_system_prompt || "You are a Project Manager.";
        const prompt = flattenHistory(history, input, systemPrompt, 
                                     // 15,
                                      10,
                                      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: systemPrompt 
            })
        });
        const result = await response.json();
        return { text: result.data, usage: result.usage || { totalTokenCount: 0 } };
    },

    callWorker: async (history, input) => {
        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: 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 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 };
    },

    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 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; // Attach usage for billing
        return parsed;
    },
    
    generateImage: async (prompt) => {
         try {
            const response = await fetch(`${REMOTE_SERVER_URL}/api/image`, {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify({ prompt })
            });
            const result = await response.json();
            return result; // Expected { image: "base64..." }
         } catch (e) {
             console.error("Image Gen Error:", e);
             return null;
         }
    }
};

	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";

	// --- 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);
	}

	// --- 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');

	// 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;

	// 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);
	}
	}
	}

	// --- USAGE EXTRACTION FOR BILLING ---
	let usage = { totalTokenCount: 0, inputTokens: 0, outputTokens: 0 };
	let finalCleanText = fullStreamData;

	if (fullStreamData.includes("__USAGE__")) {
	const parts = fullStreamData.split("__USAGE__");
	finalCleanText = parts[0]; // The actual text content
	const usageRaw = parts[1];

	try {
	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);
	}
	}

	// 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 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,
	10,
	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: systemPrompt
	})
	});
	return await handleStreamResponse(response, onThink, onOutput);
	},

	callWorkerStream: async (history, input, onThink, onOutput, images = []) => {
	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: systemPrompt,
	images: images
	})
	});
	return await handleStreamResponse(response, onThink, onOutput);
	},

	// --- BLOCKING CALLS (Background Initialization) ---

	callPM: async (history, input, gdd = null) => {
	const systemPrompt = prompts.pm_system_prompt \|\| "You are a Project Manager.";
	const prompt = flattenHistory(history, input, systemPrompt,
	// 15,
	10,
	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: systemPrompt
	})
	});
	const result = await response.json();
	return { text: result.data, usage: result.usage \|\| { totalTokenCount: 0 } };
	},

	callWorker: async (history, input) => {
	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: 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 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 };
	},

	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 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; // Attach usage for billing
	return parsed;
	},

	generateImage: async (prompt) => {
	try {
	const response = await fetch(`${REMOTE_SERVER_URL}/api/image`, {
	method: 'POST',
	headers: { 'Content-Type': 'application/json' },
	body: JSON.stringify({ prompt })
	});
	const result = await response.json();
	return result; // Expected { image: "base64..." }
	} catch (e) {
	console.error("Image Gen Error:", e);
	return null;
	}
	}
	};