diff --git a/api/app/clients/ChatGPTClient.js b/api/app/clients/ChatGPTClient.js
index 7ff1433f23..58483bb7f7 100644
--- a/api/app/clients/ChatGPTClient.js
+++ b/api/app/clients/ChatGPTClient.js
@@ -548,7 +548,7 @@ ${botMessage.message}
     if (isChatGptModel) {
       return { prompt: [instructionsPayload, messagePayload], context };
     }
-    return { prompt, context };
+    return { prompt, context, promptTokens: currentTokenCount };
   }
 
   getTokenCount(text) {
diff --git a/config/add-balance.js b/config/add-balance.js
index a6069733c6..ed2e612898 100644
--- a/config/add-balance.js
+++ b/config/add-balance.js
@@ -52,6 +52,13 @@ const Transaction = require('../api/models/Transaction');
     // console.purple(`[DEBUG] Args Length: ${process.argv.length}`);
   }
 
+  if (!process.env.CHECK_BALANCE) {
+    console.red(
+      'Error: CHECK_BALANCE environment variable is not set! Configure it to use it: `CHECK_BALANCE=true`',
+    );
+    silentExit(1);
+  }
+
   /**
    * If we don't have the right number of arguments, lets prompt the user for them
    */
@@ -99,7 +106,7 @@ const Transaction = require('../api/models/Transaction');
   }
 
   // Check the result
-  if (!result.tokenCredits) {
+  if (!result?.tokenCredits) {
     console.red('Error: Something went wrong while updating the balance!');
     console.error(result);
     silentExit(1);