added prompt assist to token balance system

MichielMAnalytics · MichielMAnalytics · commit 93dfbc364064 · 2025-06-28T10:27:09.000+02:00
diff --git a/api/server/routes/promptAssist.js b/api/server/routes/promptAssist.js
@@ -1,7 +1,12 @@
 const express = require('express');
 const OpenAI = require('openai');
-const { getCustomConfig } = require('~/server/services/Config');
+const { promptTokensEstimate } = require('openai-chat-tokens');
+const { EModelEndpoint, supportsBalanceCheck } = require('librechat-data-provider');
+const { getCustomConfig, getBalanceConfig } = require('~/server/services/Config');
 const { requireJwtAuth } = require('~/server/middleware');
+const { checkBalance } = require('~/models/balanceMethods');
+const { spendTokens } = require('~/models/spendTokens');
+const { getValueKey } = require('~/models/tx');
 const { logger } = require('~/config');
 
 const router = express.Router();
@@ -35,21 +40,17 @@ router.post('/', requireJwtAuth, async (req, res) => {
       apiKey: process.env.OPENAI_API_KEY,
     });
     
-    // Construct the prompt for the LLM
-    const variablesSection = availableVariables && availableVariables.length > 0 
-      ? `\n\nAvailable Variables (use these when appropriate):
-${availableVariables.map(v => `- ${v.syntax}: ${v.description}`).join('\n')}
-
-When enhancing the prompt, consider incorporating these variables where they would be useful. For example, if the assistant needs current information, suggest using {{current_date}} or {{current_datetime}}. If it should personalize responses, mention {{current_user}}.`
-      : '';
-
+    // Prepare messages for token estimation
     const systemPrompt = `You are an expert at writing clear, effective prompts for AI assistants. 
 Your task is to enhance the given instructions to make them more clear, comprehensive, and effective.
 
-IMPORTANT FORMATTING RULES:
+CRITICAL FORMATTING RULES - FOLLOW EXACTLY:
 - Do NOT include the title or description in your response (they are already set by the user)
-- Do NOT use markdown formatting like **bold** or ##headers
-- Use plain text with clear structure using dashes, numbers, or bullet points
+- NEVER use asterisks (*) or double asterisks (**) for ANY reason
+- NEVER use markdown formatting like **bold**, *italic*, ##headers, or backticks
+- NEVER write **Target Audience** or **Value Proposition** - write Target Audience or VALUE PROPOSITION instead
+- Use ONLY plain text with numbers (1. 2. 3.) or simple dashes (-) for structure
+- For emphasis use CAPITAL LETTERS, not asterisks
 - Start directly with the enhanced instructions content
 
 The enhanced prompt should:
@@ -58,7 +59,12 @@ The enhanced prompt should:
 - Specify the desired output format when applicable
 - Include examples if helpful
 - Be well-structured and easy to follow
-- Incorporate dynamic variables when they would enhance the functionality${variablesSection}
+- Incorporate dynamic variables when they would enhance the functionality${availableVariables && availableVariables.length > 0 
+      ? `\n\nAvailable Variables (use these when appropriate):
+${availableVariables.map(v => `- ${v.syntax}: ${v.description}`).join('\n')}
+
+When enhancing the prompt, consider incorporating these variables where they would be useful. For example, if the assistant needs current information, suggest using {{current_date}} or {{current_datetime}}. If it should personalize responses, mention {{current_user}}.`
+      : ''}
 
 Return ONLY the enhanced instructions content without any title, description, headers, or markdown formatting.`;
 
@@ -70,12 +76,41 @@ Current Instructions: ${instructions || 'No instructions provided yet'}
 
 Create enhanced instructions that will help the AI assistant perform its intended function effectively. Consider incorporating the available variables where they would be beneficial. Remember to return only the instructions content without repeating the title or description.`;
 
-    // Use OpenAI client directly
+    // Estimate token usage for balance check
     const messages = [
       { role: 'system', content: systemPrompt },
       { role: 'user', content: userPrompt }
     ];
     
+    const promptTokens = promptTokensEstimate({ messages, model });
+    const estimatedCompletionTokens = 1000; // Max tokens we're requesting
+    
+    // Check balance if enabled
+    const balanceConfig = await getBalanceConfig();
+    if (balanceConfig?.enabled && supportsBalanceCheck[EModelEndpoint.openAI]) {
+      try {
+        await checkBalance({
+          req,
+          res,
+          txData: {
+            user: req.user.id,
+            tokenType: 'prompt',
+            amount: promptTokens + estimatedCompletionTokens,
+            model,
+            endpoint: EModelEndpoint.openAI,
+            context: 'prompt-assist',
+          },
+        });
+      } catch (err) {
+        logger.error('[/api/prompt-assist] Balance check failed:', err);
+        return res.status(402).json({ 
+          error: 'Insufficient balance', 
+          details: err.message 
+        });
+      }
+    }
+
+    // Use OpenAI client directly
     const response = await openai.chat.completions.create({
       model,
       messages,
@@ -85,6 +120,31 @@ Create enhanced instructions that will help the AI assistant perform its intende
 
     const enhancedPrompt = response.choices[0]?.message?.content || '';
     
+    // Log token usage
+    if (response.usage) {
+      logger.info(`[/api/prompt-assist] Token usage for instruction enhancement - Total: ${response.usage.total_tokens} tokens (prompt: ${response.usage.prompt_tokens}, completion: ${response.usage.completion_tokens}) for user: ${req.user.id} using model: ${model}`);
+    }
+    
+    // Record token usage
+    if (balanceConfig?.enabled && response.usage) {
+      const { prompt_tokens, completion_tokens } = response.usage;
+      await spendTokens(
+        {
+          user: req.user.id,
+          conversationId: 'prompt-assist',
+          model,
+          context: 'prompt-assist',
+          endpoint: EModelEndpoint.openAI,
+          endpointTokenConfig: customConfig?.endpoints?.[EModelEndpoint.openAI]?.tokenConfig,
+          valueKey: getValueKey(model, EModelEndpoint.openAI),
+        },
+        {
+          promptTokens: prompt_tokens,
+          completionTokens: completion_tokens,
+        }
+      );
+    }
+    
     res.json({ enhancedPrompt });
   } catch (error) {
     logger.error('[/api/prompt-assist] Error enhancing prompt:', error);
@@ -125,6 +185,7 @@ router.post('/enhance-message', requireJwtAuth, async (req, res) => {
       apiKey: process.env.OPENAI_API_KEY,
     });
     
+    // Prepare system prompt
     const systemPrompt = `You are an expert at improving user messages to AI assistants to make them much clearer, more specific, and more effective for getting high-quality responses.
 
 Your task is to significantly enhance the user's message while preserving their original intent. The enhancement should:
@@ -147,24 +208,92 @@ Examples:
 - "help me write code" → "Help me write clean, well-documented Python code for [specific functionality]. Include error handling, follow PEP 8 standards, and add inline comments explaining the logic. Provide the complete code with example usage."
 - "make a sales email" → "Create a professional sales email template for B2B cold outreach targeting [specific industry/role]. The email should be personalized, include a clear value proposition, have a compelling subject line, and end with a specific call-to-action. Keep it under 150 words and maintain a consultative tone."
 
-Return the enhanced message as plain text.`;
+Return the enhanced message as plain text.
+
+CRITICAL FORMATTING RULES - YOU MUST FOLLOW THESE:
+- NEVER use asterisks (*) or double asterisks (**) for ANY reason
+- NEVER use markdown formatting of any kind
+- NEVER write **bold text** - write BOLD TEXT or Bold Text instead
+- NEVER use backticks for code or formatting
+- Use ONLY plain text with clear structure
+- Use ONLY numbers (1. 2. 3.) or simple dashes (-) for lists
+- For emphasis, use CAPITAL LETTERS instead of asterisks
+- Write headings as plain text (e.g., "Target Audience:" not "**Target Audience**")`;
 
     const userPrompt = `Transform this user message into a detailed, comprehensive prompt that will help an AI assistant provide the best possible response:
 
 ${message}`;
 
+    // Prepare messages for API call
+    const messages = [
+      { role: 'system', content: systemPrompt },
+      { role: 'user', content: userPrompt }
+    ];
+    
+    // Estimate token usage for balance check
+    const promptTokens = promptTokensEstimate({ messages, model });
+    const estimatedCompletionTokens = 500; // Max tokens we're requesting
+    
+    // Check balance if enabled
+    const balanceConfig = await getBalanceConfig();
+    if (balanceConfig?.enabled && supportsBalanceCheck[EModelEndpoint.openAI]) {
+      try {
+        await checkBalance({
+          req,
+          res,
+          txData: {
+            user: req.user.id,
+            tokenType: 'prompt',
+            amount: promptTokens + estimatedCompletionTokens,
+            model,
+            endpoint: EModelEndpoint.openAI,
+            context: 'prompt-assist-message',
+          },
+        });
+      } catch (err) {
+        logger.error('[/api/prompt-assist/enhance-message] Balance check failed:', err);
+        return res.status(402).json({ 
+          error: 'Insufficient balance', 
+          details: err.message,
+          originalMessage: message
+        });
+      }
+    }
+
     const response = await openai.chat.completions.create({
       model,
-      messages: [
-        { role: 'system', content: systemPrompt },
-        { role: 'user', content: userPrompt }
-      ],
+      messages,
       temperature: 0.3,
       max_tokens: 500,
     });
 
     const enhancedMessage = response.choices[0]?.message?.content || message;
     
+    // Log token usage
+    if (response.usage) {
+      logger.info(`[/api/prompt-assist/enhance-message] Token usage for message enhancement - Total: ${response.usage.total_tokens} tokens (prompt: ${response.usage.prompt_tokens}, completion: ${response.usage.completion_tokens}) for user: ${req.user.id} using model: ${model}`);
+    }
+    
+    // Record token usage
+    if (balanceConfig?.enabled && response.usage) {
+      const { prompt_tokens, completion_tokens } = response.usage;
+      await spendTokens(
+        {
+          user: req.user.id,
+          conversationId: 'prompt-assist-message',
+          model,
+          context: 'prompt-assist-message',
+          endpoint: EModelEndpoint.openAI,
+          endpointTokenConfig: customConfig?.endpoints?.[EModelEndpoint.openAI]?.tokenConfig,
+          valueKey: getValueKey(model, EModelEndpoint.openAI),
+        },
+        {
+          promptTokens: prompt_tokens,
+          completionTokens: completion_tokens,
+        }
+      );
+    }
+    
     res.json({ enhancedMessage });
   } catch (error) {
     logger.error('[/api/prompt-assist/enhance-message] Error enhancing message:', error);
