Merge pull request #178 from AlexHTW/add-budget-types

Add budget periods and improvements

README.md

@@ -57,8 +57,9 @@ The following parameters are optional and can be set in the `.env` file:
 | `ENABLE_QUOTING`                   | Whether to enable message quoting in private chats                                                                                                                                                                               | true                           |
 | `ENABLE_IMAGE_GENERATION`          | Whether to enable image generation via the `/image` command                                                                                                                                                                      | true                           |
 | `ENABLE_TRANSCRIPTION`             | Whether to enable transcriptions of audio and video messages                                                                                                                                                                     | true                           |
-| `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 (`*`)                                   | `*`                            |
+| `BUDGET_PERIOD`                    | Determines the time frame all budgets are applied to. Available periods: `daily` *(resets budget every day)*, `monthly` *(resets budgets on the first of each month)*, `all-time` *(never resets budget)*                        | `monthly`                      |
-| `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`="*") | `100.0`                        |
+| `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. For `*`- user lists the first `USER_BUDGETS` value is given to every user. **Note**: by default, *no limits* for any user (`*`) | `*`                        |
+| `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 (`USER_BUDGETS`="*")         | `100.0`                        |
 | `PROXY`                            | Proxy to be used for OpenAI and Telegram bot (e.g. `http://localhost:8080`)                                                                                                                                                      | -                              |
 | `OPENAI_MODEL`                     | The OpenAI model to use for generating responses                                                                                                                                                                                 | `gpt-3.5-turbo`                |
 | `ASSISTANT_PROMPT`                 | A system message that sets the tone and controls the behavior of the assistant                                                                                                                                                   | `You are a helpful assistant.` |
@@ -80,6 +81,7 @@ The following parameters are optional and can be set in the `.env` file:
 | `TRANSCRIPTION_PRICE`              | USD-price for one minute of audio transcription                                                                                                                                                                                  | 0.006                          |
 
 Check out the [official API reference](https://platform.openai.com/docs/api-reference/chat) for more details.  
+Check out our [Budget Manual](https://github.com/n3d1117/chatgpt-telegram-bot/discussions/184) for possible budget configurations.
 
 ### Installing
 Clone the repository and navigate to the project directory:

bot/main.py

@@ -44,6 +44,16 @@ def main():
         'frequency_penalty': float(os.environ.get('FREQUENCY_PENALTY', 0.0)),
     }
 
+    # log deprecation warning for old budget variable names
+    # old variables are caught in the telegram_config definition for now
+    # remove support for old budget names at some point in the future
+    if os.environ.get('MONTHLY_USER_BUDGETS') is not None:
+        logging.warning('The environment variable MONTHLY_USER_BUDGETS is deprecated. '
+                     'Please use USER_BUDGETS with BUDGET_PERIOD instead.')
+    if os.environ.get('MONTHLY_GUEST_BUDGET') is not None:
+        logging.warning('The environment variable MONTHLY_GUEST_BUDGET is deprecated. '
+                     'Please use GUEST_BUDGET with BUDGET_PERIOD instead.')
+
     telegram_config = {
         'token': os.environ['TELEGRAM_BOT_TOKEN'],
         'admin_user_ids': os.environ.get('ADMIN_USER_IDS', '-'),
@@ -51,8 +61,9 @@ def main():
         'enable_quoting': os.environ.get('ENABLE_QUOTING', 'true').lower() == 'true',
         'enable_image_generation': os.environ.get('ENABLE_IMAGE_GENERATION', 'true').lower() == 'true',
         'enable_transcription': os.environ.get('ENABLE_TRANSCRIPTION', 'true').lower() == 'true',
-        'monthly_user_budgets': os.environ.get('MONTHLY_USER_BUDGETS', '*'),
+        'budget_period': os.environ.get('BUDGET_PERIOD', 'monthly').lower(),
-        'monthly_guest_budget': float(os.environ.get('MONTHLY_GUEST_BUDGET', '100.0')),
+        'user_budgets': os.environ.get('USER_BUDGETS', os.environ.get('MONTHLY_USER_BUDGETS', '*')),
+        'guest_budget': float(os.environ.get('GUEST_BUDGET', os.environ.get('MONTHLY_GUEST_BUDGET', '100.0'))),
         'stream': os.environ.get('STREAM', 'true').lower() == 'true',
         'proxy': os.environ.get('PROXY', None),
         'voice_reply_transcript': os.environ.get('VOICE_REPLY_WITH_TRANSCRIPT_ONLY', 'true').lower() == 'true',

bot/openai_helper.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import datetime
 import logging
 

bot/telegram_bot.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import logging
 import os
 import itertools
@@ -34,6 +35,18 @@ class ChatGPTTelegramBot:
     """
     Class representing a ChatGPT Telegram Bot.
     """
+    # Mapping of budget period to cost period
+    budget_cost_map = {
+            "monthly":"cost_month",
+            "daily":"cost_today",
+            "all-time":"cost_all_time"
+        }
+    # Mapping of budget period to a print output
+    budget_print_map = {
+        "monthly": " this month",
+        "daily": " today",
+        "all-time": ""
+    }
 
     def __init__(self, config: dict, openai: OpenAIHelper):
         """
@@ -56,7 +69,7 @@ class ChatGPTTelegramBot:
         ] + self.commands
         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"
-        self.budget_limit_message = "Sorry, you have reached your monthly usage limit."
+        self.budget_limit_message = f"Sorry, you have reached your usage limit{self.budget_print_map[config['budget_period']]}."
         self.usage = {}
         self.last_message = {}
 
@@ -95,12 +108,13 @@ class ChatGPTTelegramBot:
 
         tokens_today, tokens_month = self.usage[user_id].get_current_token_usage()
         images_today, images_month = self.usage[user_id].get_current_image_count()
-        transcribe_durations = self.usage[user_id].get_current_transcription_duration()
+        (transcribe_minutes_today, transcribe_seconds_today, transcribe_minutes_month, 
-        cost_today, cost_month = self.usage[user_id].get_current_cost()
+            transcribe_seconds_month) = self.usage[user_id].get_current_transcription_duration()
+        current_cost = self.usage[user_id].get_current_cost()
         
         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)
+        remaining_budget = self.get_remaining_budget(update)
 
         text_current_conversation = f"*Current conversation:*\n"+\
                      f"{chat_messages} chat messages in history.\n"+\
@@ -109,18 +123,19 @@ class ChatGPTTelegramBot:
         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"{transcribe_minutes_today} minutes and {transcribe_seconds_today} seconds transcribed.\n"+\
-                     f"💰 For a total amount of ${cost_today:.2f}\n"+\
+                     f"💰 For a total amount of ${current_cost['cost_today']:.2f}\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"{transcribe_minutes_month} minutes and {transcribe_seconds_month} seconds transcribed.\n"+\
-                     f"💰 For a total amount of ${cost_month:.2f}"
+                     f"💰 For a total amount of ${current_cost['cost_month']:.2f}"
         # text_budget filled with conditional content
         text_budget = "\n\n"
-        if budget < float('inf'):
+        budget_period =self.config['budget_period']
-            text_budget += f"You have a remaining budget of ${budget:.2f} this month.\n"
+        if remaining_budget < float('inf'):
+            text_budget += f"You have a remaining budget of ${remaining_budget:.2f}{self.budget_print_map[budget_period]}.\n"
         # add OpenAI account information for admin request
         if self.is_admin(update):
             text_budget += f"Your OpenAI account was billed ${self.openai.get_billing_current_month():.2f} this month."
@@ -631,13 +646,14 @@ class ChatGPTTelegramBot:
 
         return False
 
-    def is_admin(self, update: Update) -> bool:
+    def is_admin(self, update: Update, log_no_admin=False) -> 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.')
+            if log_no_admin:
+                logging.info('No admin user defined.')
             return False
 
         admin_user_ids = self.config['admin_user_ids'].split(',')
@@ -648,78 +664,72 @@ class ChatGPTTelegramBot:
 
         return False
 
-    async def get_remaining_budget(self, update: Update) -> float:
+    def get_user_budget(self, update: Update) -> float | None:
+        """
+        Get the user's budget based on their user ID and the bot configuration.
+        :param update: Telegram update object
+        :return: The user's budget as a float, or None if the user is not found in the allowed user list
+        """
+        
+        # no budget restrictions for admins and '*'-budget lists
+        if self.is_admin(update) or self.config['user_budgets'] == '*':
+            return float('inf')
+        
+        user_budgets = self.config['user_budgets'].split(',')
+        if self.config['allowed_user_ids'] == '*':
+            # same budget for all users, use value in first position of budget list
+            if len(user_budgets) > 1:
+                logging.warning('multiple values for budgets set with unrestricted user list '
+                                'only the first value is used as budget for everyone.')
+            return float(user_budgets[0])
+
         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 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}).')
+                logging.warning(f'No budget set for user id: {user_id}. Budget list shorter than user list.')
                 return 0.0
-            user_budget = float(user_budgets[user_index])
+            return float(user_budgets[user_index])
-            cost_month = self.usage[user_id].get_current_cost()[1]
+        return None
-            remaining_budget = user_budget - cost_month
-            return remaining_budget
-        else:
-            return 0.0
 
-    async def is_within_budget(self, update: Update, context: CallbackContext) -> bool:
+    def get_remaining_budget(self, update: Update) -> float:
         """
-        Checks if the user reached their monthly usage limit.
+        Calculate the remaining budget for a user based on their current usage.
-        Initializes UsageTracker for user and guest when needed.
+        :param update: Telegram update object
+        :return: The remaining budget for the user as a 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 self.is_admin(update):
+        # Get budget for users
-            return True
+        user_budget = self.get_user_budget(update)
+        budget_period = self.config['budget_period']
+        if user_budget is not None:
+            cost = self.usage[user_id].get_current_cost()[self.budget_cost_map[budget_period]]
+            return user_budget - cost
 
-        if self.config['monthly_user_budgets'] == '*':
+        # Get budget for guests
-            return True
+        if 'guests' not in self.usage:
+            self.usage['guests'] = UsageTracker('guests', 'all guest users in group chats')
+        cost = self.usage['guests'].get_current_cost()[self.budget_cost_map[budget_period]]
+        return self.config['guest_budget'] - cost
 
-        allowed_user_ids = self.config['allowed_user_ids'].split(',')
+    def is_within_budget(self, update: Update) -> bool:
-        if str(user_id) in allowed_user_ids:
+        """
-            # find budget for allowed user
+        Checks if the user reached their usage limit.
-            user_index = allowed_user_ids.index(str(user_id))
+        Initializes UsageTracker for user and guest when needed.
-            user_budgets = self.config['monthly_user_budgets'].split(',')
+        :param update: Telegram update object
-            # check if user is included in budgets list
+        :return: Boolean indicating if the user has a positive budget
-            if len(user_budgets) <= user_index:
+        """
-                logging.warning(f'No budget set for user: {update.message.from_user.name} ({user_id}).')
+        user_id = update.message.from_user.id
-                return False
+        if user_id not in self.usage:
-            user_budget = float(user_budgets[user_index])
+            self.usage[user_id] = UsageTracker(user_id, update.message.from_user.name)
-            cost_month = self.usage[user_id].get_current_cost()[1]
-            # Check if allowed user is within budget
-            return user_budget > cost_month
 
-        # Check if group member is within budget
+        remaining_budget = self.get_remaining_budget(update)
-        if self.is_group_chat(update):
+
-            admin_user_ids = self.config['admin_user_ids'].split(',')
+        return remaining_budget > 0
-            for user in itertools.chain(allowed_user_ids, admin_user_ids):
-                if not user.strip():
-                    continue
-                if await self.is_user_in_group(update, context, user):
-                    if 'guests' not in self.usage:
-                        self.usage['guests'] = UsageTracker('guests', 'all guest users in group chats')
-                    if self.config['monthly_guest_budget'] >= self.usage['guests'].get_current_cost()[1]:
-                        return True
-                    logging.warning('Monthly guest budget for group chats used up.')
-                    return False
-            logging.info(f'Group chat messages from user {update.message.from_user.name} '
-                f'(id: {update.message.from_user.id}) are not allowed')
-        return False
 
     async def check_allowed_and_within_budget(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> bool:
         """
@@ -734,7 +744,7 @@ class ChatGPTTelegramBot:
             await self.send_disallowed_message(update, context)
             return False
 
-        if not await self.is_within_budget(update, context):
+        if not self.is_within_budget(update):
             logging.warning(f'User {update.message.from_user.name} (id: {update.message.from_user.id}) '
                 f'reached their usage limit')
             await self.send_budget_reached_message(update, context)

bot/usage_tracker.py

@@ -18,6 +18,7 @@ class UsageTracker:
         "current_cost": {
             "day": 0.45, 
             "month": 3.23,
+            "all_time": 3.23,
             "last_update": "2023-03-14"},
         "usage_history": {
             "chat_tokens": {
@@ -59,7 +60,7 @@ class UsageTracker:
             # create new dictionary for this user
             self.usage = {
                 "user_name": user_name,
-                "current_cost": {"day": 0.0, "month": 0.0, "last_update": str(date.today())},
+                "current_cost": {"day": 0.0, "month": 0.0, "all_time": 0.0, "last_update": str(date.today())},
                 "usage_history": {"chat_tokens": {}, "transcription_seconds": {}, "number_images": {}}
             }
 
@@ -72,18 +73,20 @@ class UsageTracker:
         """
         today = date.today()
         last_update = date.fromisoformat(self.usage["current_cost"]["last_update"])
+        token_cost = round(tokens * tokens_price / 1000, 6)
+        # add to all_time cost, initialize with calculation of total_cost if key doesn't exist
+        self.usage["current_cost"]["all_time"] = self.usage["current_cost"].get("all_time", self.initialize_all_time_cost()) + token_cost
         # add current cost, update new day
         if today == last_update:
-            self.usage["current_cost"]["day"] += round(tokens * tokens_price / 1000, 6)
+            self.usage["current_cost"]["day"] += token_cost
-            self.usage["current_cost"]["month"] += round(tokens * tokens_price / 1000, 6)
+            self.usage["current_cost"]["month"] += token_cost
         else:
             if today.month == last_update.month:
-                self.usage["current_cost"]["month"] += round(tokens * tokens_price / 1000, 6)
+                self.usage["current_cost"]["month"] += token_cost
             else:
-                self.usage["current_cost"]["month"] = round(tokens * tokens_price / 1000, 6)
+                self.usage["current_cost"]["month"] = token_cost
-            self.usage["current_cost"]["day"] = round(tokens * tokens_price / 1000, 6)
+            self.usage["current_cost"]["day"] = token_cost
             self.usage["current_cost"]["last_update"] = str(today)
-
         # update usage_history
         if str(today) in self.usage["usage_history"]["chat_tokens"]:
             # add token usage to existing date
@@ -128,6 +131,8 @@ class UsageTracker:
 
         today = date.today()
         last_update = date.fromisoformat(self.usage["current_cost"]["last_update"])
+        # add to all_time cost, initialize with calculation of total_cost if key doesn't exist
+        self.usage["current_cost"]["all_time"] = self.usage["current_cost"].get("all_time", self.initialize_all_time_cost()) + image_cost
         # add current cost, update new day
         if today == last_update:
             self.usage["current_cost"]["day"] += image_cost
@@ -175,20 +180,23 @@ class UsageTracker:
     def add_transcription_seconds(self, seconds, minute_price=0.006):
         """Adds requested transcription seconds to a users usage history and updates current cost.
         :param tokens: total tokens used in last request
-        :param tokens_price: price per 1000 tokens, defaults to 0.002
+        :param tokens_price: price per minute transcription, defaults to 0.006
         """
         today = date.today()
         last_update = date.fromisoformat(self.usage["current_cost"]["last_update"])
+        transcription_price = round(seconds * minute_price / 60, 2)
+        # add to all_time cost, initialize with calculation of total_cost if key doesn't exist
+        self.usage["current_cost"]["all_time"] = self.usage["current_cost"].get("all_time", self.initialize_all_time_cost()) + transcription_price
         # add current cost, update new day
         if today == last_update:
-            self.usage["current_cost"]["day"] += round(seconds * minute_price / 60, 2)
+            self.usage["current_cost"]["day"] += transcription_price
-            self.usage["current_cost"]["month"] += round(seconds * minute_price / 60, 2)
+            self.usage["current_cost"]["month"] += transcription_price
         else:
             if today.month == last_update.month:
-                self.usage["current_cost"]["month"] += round(seconds * minute_price / 60, 2)
+                self.usage["current_cost"]["month"] += transcription_price
             else:
-                self.usage["current_cost"]["month"] = round(seconds * minute_price / 60, 2)
+                self.usage["current_cost"]["month"] = transcription_price
-            self.usage["current_cost"]["day"] = round(seconds * minute_price / 60, 2)
+            self.usage["current_cost"]["day"] = transcription_price
             self.usage["current_cost"]["last_update"] = str(today)
 
         # update usage_history
@@ -239,4 +247,28 @@ class UsageTracker:
                 cost_month = self.usage["current_cost"]["month"]
             else:
                 cost_month = 0.0
-        return round(cost_day, 3), round(cost_month, 3)
+        # add to all_time cost, initialize with calculation of total_cost if key doesn't exist
+        cost_all_time = self.usage["current_cost"].get("all_time", self.initialize_all_time_cost())
+        return {"cost_today": cost_day, "cost_month": cost_month, "cost_all_time": cost_all_time}
+
+    def initialize_all_time_cost(self, tokens_price=0.002, image_prices="0.016,0.018,0.02", minute_price=0.006):
+        """Get total USD amount of all requests in history
+        
+        :param tokens_price: price per 1000 tokens, defaults to 0.002
+        :param image_prices: prices for images of sizes ["256x256", "512x512", "1024x1024"],
+            defaults to [0.016, 0.018, 0.02]
+        :param tokens_price: price per minute transcription, defaults to 0.006
+        :return: total cost of all requests
+        """
+        total_tokens = sum(self.usage['usage_history']['chat_tokens'].values())
+        token_cost = round(total_tokens * tokens_price / 1000, 6)
+        
+        total_images = [sum(values) for values in zip(*self.usage['usage_history']['number_images'].values())]
+        image_prices_list = [float(x) for x in image_prices.split(',')]
+        image_cost = sum([count * price for count, price in zip(total_images, image_prices_list)])
+        
+        total_transcription_seconds = sum(self.usage['usage_history']['transcription_seconds'].values())
+        transcription_cost = round(total_transcription_seconds * minute_price / 60, 2)
+
+        all_time_cost = token_cost + transcription_cost + image_cost
+        return all_time_cost