From f5f5b2bbdb02f24e36e098dc9a810de35ec45524 Mon Sep 17 00:00:00 2001
From: Danny Avila <110412045+danny-avila@users.noreply.github.com>
Date: Sun, 26 Nov 2023 18:12:27 -0500
Subject: [PATCH] =?UTF-8?q?fix:=20Resolve=20Token=20Credit=20Balance=20Iss?=
 =?UTF-8?q?ues=20for=20Instruct=20Models=20=F0=9F=9B=A0=EF=B8=8F=20(#1232)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Fix: balance update error and add environment variable check

* fix(ChatGPTClient): return promptTokens for instruct/davinci models

* chore: remove unnecessary comments
---
 api/app/clients/ChatGPTClient.js | 2 +-
 config/add-balance.js            | 9 ++++++++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/api/app/clients/ChatGPTClient.js b/api/app/clients/ChatGPTClient.js
index 7ff1433f2..58483bb7f 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 a6069733c..ed2e61289 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);