added grant balance and openAI usage to stats

.env.example

@@ -4,7 +4,7 @@ OPENAI_API_KEY="XXX"
 # Your Telegram bot token obtained using @BotFather
 TELEGRAM_BOT_TOKEN="XXX"
 
-# Comma separated list of telegram user IDs, or * to allow all
+# Comma separated list of telegram user IDs, or * to allow all. First USER_ID has access to admin commands.
 ALLOWED_TELEGRAM_USER_IDS="USER_ID_1,USER_ID_2"
 
 # Comma separated list of user budgets, USD limit per month, or * to allow all users unlimited usage

README.md

@@ -71,7 +71,7 @@ TRANSCRIPTION_PRICE=0.006 # Defaults to minute price of OpenAI Whisper of 0.006
 ```
 * `OPENAI_API_KEY`: Your OpenAI API key, you can get it from [here](https://platform.openai.com/account/api-keys)
 * `TELEGRAM_BOT_TOKEN`: Your Telegram bot's token, obtained using [BotFather](http://t.me/botfather) (see [tutorial](https://core.telegram.org/bots/tutorial#obtain-your-bot-token))
-* `ALLOWED_TELEGRAM_USER_IDS`: A comma-separated list of Telegram user IDs that are allowed to interact with the bot (use [getidsbot](https://t.me/getidsbot) to find your user ID). **Note**: by default, *everyone* is allowed (`*`)
+* `ALLOWED_TELEGRAM_USER_IDS`: A comma-separated list of Telegram user IDs that are allowed to interact with the bot (use [getidsbot](https://t.me/getidsbot) to find your user ID). First USER_ID has access to admin commands. **Note**: by default, *everyone* is allowed (`*`) and has access to admin commands.
 * `MONTHLY_USER_BUDGETS`: A comma-separated list of $-amounts per user from list `ALLOWED_TELEGRAM_USER_IDS` to set custom usage limit of OpenAI API costs for each. **Note**: by default, *no limits* for anyone (`*`)
 * `MONTHLY_GUEST_BUDGET`: $-amount as usage limit for all guest users. Guest users are users in group chats that are not in the `ALLOWED_TELEGRAM_USER_IDS` list. Value is ignored if no usage limits are set in user budgets (`MONTHLY_USER_BUDGETS`="*")
 * `PROXY`: Proxy to be used for OpenAI and Telegram bot

bot/openai_helper.py

@@ -9,11 +9,15 @@ import openai
 
 import requests
 import json
+from datetime import date
+from calendar import monthrange
+
 # 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.
@@ -224,8 +228,11 @@ class OpenAIHelper:
         else:
             raise NotImplementedError(f"__count_tokens() is not presently implemented for model {model}")
 
-    def get_balance(self):
+    def get_grant_balance(self):
+        """Gets remaining grant balance for new users from OpenAI API.
 
+        :return: remaining grant balance
+        """
         headers = {
             "Authorization": f"Bearer {openai.api_key}"
         }
@@ -233,3 +240,25 @@ class OpenAIHelper:
         billing_data = json.loads(response.text)
         balance = billing_data["total_available"]
         return balance
+    
+    def get_billing_current_month(self):
+        """Gets billed usage for current month from OpenAI API.
+
+        :return: dollar amount of usage this month
+        """
+        headers = {
+            "Authorization": f"Bearer {openai.api_key}"
+        }
+        # calculate first and last day of current month
+        today = date.today()
+        first_day = date(today.year, today.month, 1)
+        _, last_day_of_month = monthrange(today.year, today.month)
+        last_day = date(today.year, today.month, last_day_of_month)
+        params = {
+            "start_date": first_day,
+            "end_date": last_day
+        }
+        response = requests.get("https://api.openai.com/dashboard/billing/usage", headers=headers, params=params)
+        billing_data = json.loads(response.text)
+        usage_month = billing_data["total_usage"] / 100 # convert cent amount to dollars
+        return usage_month

bot/telegram_bot.py

@@ -60,7 +60,7 @@ class ChatGPT3TelegramBot:
             await self.send_disallowed_message(update, context)
             return
 
-        logging.info(f'User {update.message.from_user.name} requested their token usage statistics')
+        logging.info(f'User {update.message.from_user.name} requested their usage statistics')
         
         user_id = update.message.from_user.id
         if user_id not in self.usage:
@@ -73,38 +73,36 @@ class ChatGPT3TelegramBot:
         
         chat_id = update.effective_chat.id
         chat_messages, chat_token_length = self.openai.get_conversation_stats(chat_id)
+        budget = await self.get_remaining_budget(update)
 
-        usage_text = f"Today:\n"+\
+        text_current_conversation = f"Current conversation:\n"+\
+                     f"{chat_messages} chat messages in history.\n"+\
+                     f"{chat_token_length} chat tokens in history.\n"+\
+                     f"\n----------------------------\n\n"
+        text_today = f"Usage today:\n"+\
                      f"{tokens_today} chat tokens used.\n"+\
                      f"{images_today} images generated.\n"+\
                      f"{transcribe_durations[0]} minutes and {transcribe_durations[1]} seconds transcribed.\n"+\
                      f"💰 For a total amount of ${cost_today:.2f}\n"+\
-                     f"\n----------------------------\n\n"+\
+                     f"\n----------------------------\n\n"
-                     f"This month:\n"+\
+        text_month = f"Usage this month:\n"+\
                      f"{tokens_month} chat tokens used.\n"+\
                      f"{images_month} images generated.\n"+\
                      f"{transcribe_durations[2]} minutes and {transcribe_durations[3]} seconds transcribed.\n"+\
-                     f"💰 For a total amount of ${cost_month:.2f}"+\
+                     f"💰 For a total amount of ${cost_month:.2f}"
-                     f"\n----------------------------\n\n"+\
+        text_budget = "\n\n"
-                     f"Current conversation:\n"+\
+        if budget < float('inf'):
-                     f"{chat_messages} chat messages in history.\n"+\
+            text_budget += f"You have a remaining budget of ${budget:.2f} this month.\n"
-                     f"{chat_token_length} chat tokens in history.\n"
+        # add OpenAI account information for admin request
+        if await self.is_admin(update):
+            grant_balance = self.openai.get_grant_balance()
+            if grant_balance > 0.0:
+                text_budget += f"Your remaining OpenAI grant balance is ${grant_balance:.2f}.\n"
+            text_budget += f"Your OpenAI account was billed ${self.openai.get_billing_current_month():.2f} this month."
+        
+        usage_text = text_current_conversation + text_today + text_month + text_budget
         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.
@@ -429,7 +427,6 @@ class ChatGPT3TelegramBot:
             return True
 
         allowed_user_ids = self.config['allowed_user_ids'].split(',')
-
         # Check if user is allowed
         if str(update.message.from_user.id) in allowed_user_ids:
             return True
@@ -461,16 +458,13 @@ class ChatGPT3TelegramBot:
         return False
 
     async def get_remaining_budget(self, update: Update) -> float:
-        logging.info(f'Getting remaining budget for user: {update.message.from_user.name} ({user_id}).')
-
         user_id = update.message.from_user.id
-        
-        if self.config['monthly_user_budgets'] == '*':
-            return self.openai.get_balance()
-
         if user_id not in self.usage:
             self.usage[user_id] = UsageTracker(user_id, update.message.from_user.name)
 
+        if self.config['monthly_user_budgets'] == '*':
+            return float('inf')
+
         allowed_user_ids = self.config['allowed_user_ids'].split(',')
         if str(user_id) in allowed_user_ids:
             # find budget for allowed user
@@ -554,7 +548,6 @@ 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,