Add balance inquiry function

bot/openai_helper.py

@@ -6,12 +6,14 @@ import tiktoken
 
 import openai
 
+
+import requests
+import json
 # Models can be found here: https://platform.openai.com/docs/models/overview
 GPT_3_MODELS = ("gpt-3.5-turbo", "gpt-3.5-turbo-0301")
 GPT_4_MODELS = ("gpt-4", "gpt-4-0314")
 GPT_4_32K_MODELS = ("gpt-4-32k", "gpt-4-32k-0314")
 GPT_ALL_MODELS = GPT_3_MODELS + GPT_4_MODELS + GPT_4_32K_MODELS
-
 class OpenAIHelper:
     """
     ChatGPT helper class.
@@ -221,3 +223,13 @@ class OpenAIHelper:
             return num_tokens
         else:
             raise NotImplementedError(f"__count_tokens() is not presently implemented for model {model}")
+
+    def get_balance(self):
+
+        headers = {
+            "Authorization": f"Bearer {openai.api_key}"
+        }
+        response = requests.get("https://api.openai.com/dashboard/billing/credit_grants", headers=headers)
+        billing_data = json.loads(response.text)
+        balance = billing_data["total_available"]
+        return balance

bot/telegram_bot.py

@@ -28,7 +28,8 @@ class ChatGPT3TelegramBot:
             BotCommand(command='help', description='Show help message'),
             BotCommand(command='reset', description='Reset the conversation. Optionally pass high-level instructions for the conversation (e.g. /reset You are a helpful assistant)'),
             BotCommand(command='image', description='Generate image from prompt (e.g. /image cat)'),
-            BotCommand(command='stats', description='Get your current usage statistics')
+            BotCommand(command='stats', description='Get your current usage statistics'),
+            BotCommand(command='balance', description='Get your OpenAI account balance')
         ]
         self.disallowed_message = "Sorry, you are not allowed to use this bot. You can check out the source code at " \
                                   "https://github.com/n3d1117/chatgpt-telegram-bot"
@@ -90,6 +91,20 @@ class ChatGPT3TelegramBot:
                      f"{chat_token_length} chat tokens in history.\n"
         await update.message.reply_text(usage_text)
 
+    async def balance(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
+        """
+        Returns OpenAI account balance.
+        """
+        if not await self.is_allowed(update):
+            logging.warning(f'User {update.message.from_user.name} is not allowed to request OpenAI account balance')
+            await self.send_disallowed_message(update, context)
+            return
+
+        logging.info(f'User {update.message.from_user.name} requested OpenAI account balance')
+        
+        balance = self.openai.get_balance()
+        await update.message.reply_text(f"The balance of your OpenAI account is: ${balance:.3f}")
+
     async def reset(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
         """
         Resets the conversation.
@@ -496,6 +511,7 @@ class ChatGPT3TelegramBot:
         application.add_handler(CommandHandler('image', self.image))
         application.add_handler(CommandHandler('start', self.help))
         application.add_handler(CommandHandler('stats', self.stats))
+        application.add_handler(CommandHandler('balance', self.balance))
         application.add_handler(MessageHandler(
             filters.AUDIO | filters.VOICE | filters.Document.AUDIO |
             filters.VIDEO | filters.VIDEO_NOTE | filters.Document.VIDEO,