Merge pull request #103 from AlexHTW/add-balance-function

Add balance functions to stats command

.env.example

@@ -4,5 +4,8 @@ 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 everyone
+# Telegram user ID of admins, or - to assign no admin.
+ADMIN_USER_IDS="ADMIN_1_USER_ID,ADMIN_2_USER_ID"
+
+# Comma separated list of telegram user IDs, or * to allow all.
 ALLOWED_TELEGRAM_USER_IDS="USER_ID_1,USER_ID_2"

README.md

@@ -47,6 +47,7 @@ Customize the configuration by copying `.env.example` and renaming it to `.env`,
 |-----------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
 | `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))                                            |
+| `ADMIN_USER_IDS`            | Telegram user IDs of admins. These users have access to special admin commands and information and no budget restrictions. Admin IDs don't have to be added to ALLOWED_TELEGRAM_USER_IDS. **Note**: by default, no admin ('-') |
 | `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 (`*`) |
 
 ### Optional configuration

bot/main.py

@@ -44,6 +44,7 @@ def main():
 
     telegram_config = {
         'token': os.environ['TELEGRAM_BOT_TOKEN'],
+        'admin_user_ids': os.environ.get('ADMIN_USER_IDS', '-'),
         'allowed_user_ids': os.environ.get('ALLOWED_TELEGRAM_USER_IDS', '*'),
         'monthly_user_budgets': os.environ.get('MONTHLY_USER_BUDGETS', '*'),
         'monthly_guest_budget': float(os.environ.get('MONTHLY_GUEST_BUDGET', '100.0')),

bot/openai_helper.py

@@ -5,6 +5,11 @@ import tiktoken
 
 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")
@@ -266,3 +271,38 @@ class OpenAIHelper:
                     num_tokens += tokens_per_name
         num_tokens += 2  # every reply is primed with <im_start>assistant
         return num_tokens
+
+    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}"
+        }
+        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
+    
+    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

@@ -64,7 +64,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:
@@ -77,23 +77,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"
+        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"
-                     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 filled with conditional content
-                     f"Current conversation:\n"+\
+        text_budget = "\n\n"
-                     f"{chat_messages} chat messages in history.\n"+\
+        if budget < float('inf'):
-                     f"{chat_token_length} chat tokens in history.\n"
+            text_budget += f"You have a remaining budget of ${budget:.2f} this month.\n"
-        await update.message.reply_text(usage_text)
+        # 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, parse_mode=constants.ParseMode.MARKDOWN)
 
     async def reset(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
         """
@@ -176,10 +189,8 @@ class ChatGPT3TelegramBot:
 
         chat_id = update.effective_chat.id
         await context.bot.send_chat_action(chat_id=chat_id, action=constants.ChatAction.TYPING)
-
         filename = update.message.effective_attachment.file_unique_id
         filename_mp3 = f'{filename}.mp3'
-
         try:
             media_file = await context.bot.get_file(update.message.effective_attachment.file_id)
             await media_file.download_to_drive(filename)
@@ -475,8 +486,10 @@ class ChatGPT3TelegramBot:
         if self.config['allowed_user_ids'] == '*':
             return True
         
-        allowed_user_ids = self.config['allowed_user_ids'].split(',')
+        if await self.is_admin(update):
+            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
@@ -491,6 +504,50 @@ class ChatGPT3TelegramBot:
 
         return False
 
+    async def is_admin(self, update: Update) -> bool:
+        """
+        Checks if the user is the admin of the bot.
+        The first user in the user list is the admin.
+        """
+        if self.config['admin_user_ids'] == '-':
+            logging.info('No admin user defined.')
+            return False
+
+        admin_user_ids = self.config['admin_user_ids'].split(',')
+
+        # Check if user is in the admin user list
+        if str(update.message.from_user.id) in admin_user_ids:
+            return True
+
+        return False
+
+    async def get_remaining_budget(self, update: Update) -> float:
+        user_id = update.message.from_user.id
+        if user_id not in self.usage:
+            self.usage[user_id] = UsageTracker(user_id, update.message.from_user.name)
+
+        if await self.is_admin(update):
+            return float('inf')
+
+        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
+            user_index = allowed_user_ids.index(str(user_id))
+            user_budgets = self.config['monthly_user_budgets'].split(',')
+            # check if user is included in budgets list
+            if len(user_budgets) <= user_index:
+                logging.warning(f'No budget set for user: {update.message.from_user.name} ({user_id}).')
+                return 0.0
+            user_budget = float(user_budgets[user_index])
+            cost_month = self.usage[user_id].get_current_cost()[1]
+            remaining_budget = user_budget - cost_month
+            return remaining_budget
+        else:
+            return 0.0
+
     async def is_within_budget(self, update: Update) -> bool:
         """
         Checks if the user reached their monthly usage limit.
@@ -500,6 +557,9 @@ class ChatGPT3TelegramBot:
         if user_id not in self.usage:
             self.usage[user_id] = UsageTracker(user_id, update.message.from_user.name)
 
+        if await self.is_admin(update):
+            return True
+
         if self.config['monthly_user_budgets'] == '*':
             return True