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

Add balance functions to stats command

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"This month:\n"+\
+                     f"----------------------------\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"\n----------------------------\n\n"+\
-                     f"Current conversation:\n"+\
-                     f"{chat_messages} chat messages in history.\n"+\
-                     f"{chat_token_length} chat tokens in history.\n"
-        await update.message.reply_text(usage_text)
+                     f"💰 For a total amount of ${cost_month:.2f}"
+        # text_budget filled with conditional content
+        text_budget = "\n\n"
+        if budget < float('inf'):
+            text_budget += f"You have a remaining budget of ${budget:.2f} this month.\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, 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)
@@ -474,9 +485,11 @@ class ChatGPT3TelegramBot:
         """
         if self.config['allowed_user_ids'] == '*':
             return True
-
+        
+        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