minor improvements

README.md

@@ -45,7 +45,7 @@ TELEGRAM_BOT_TOKEN="YOUR_TELEGRAM_BOT_TOKEN"
 
 # Optional parameters
 ALLOWED_TELEGRAM_USER_IDS="USER_ID_1,USER_ID_2,..." # Defaults to "*" (everyone)
-MONTHLY_USER_BUDGETS="BUDGET_USER_ID_1, BUDGET_USER_ID_2,..." # Defaults to "*" (no restrictions)
+MONTHLY_USER_BUDGETS="BUDGET_USER_ID_1,BUDGET_USER_ID_2,..." # Defaults to "*" (no restrictions)
 MONTHLY_GUEST_BUDGET="20.0" # Defaults to 100.0
 PROXY="YOUR_PROXY" # e.g. "http://localhost:8080", defaults to none
 OPENAI_MODEL="gpt-3.5-turbo" # Defaults to gpt-3.5-turbo
@@ -68,8 +68,8 @@ 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 (`*`)
-* `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="*").
+* `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
 * `OPENAI_MODEL`: Define which OpenAI model to use (default is `gpt-3.5-turbo`)
 * `PRESENCE_PENALTY`: Number between -2.0 and 2.0. Positive values penalize new tokens based on whether they appear in the text so far, increasing the model's likelihood to talk about new topics.
@@ -85,7 +85,7 @@ TRANSCRIPTION_PRICE=0.006 # Defaults to minute price of OpenAI Whisper of 0.006
 * `IMAGE_SIZE`: The DALL·E generated image size. Allowed values: "256x256", "512x512", or "1024x1024"
 * `GROUP_TRIGGER_KEYWORD`: If set, the bot will only respond to messages that start with this keyword. This is useful for bots added to groups with privacy mode disabled. **Note**: by default, *no keyword* is required (`""`)
 * `TOKEN_PRICE`: USD-price per 1000 tokens for cost information in usage statistics. Defaults to [OpenAI price](https://openai.com/pricing) for gpt-3.5-turbo.
-* `IMAGE_PRICES`: A comma-separated list with 3 elements of prices for the different image sizes 256x256,512x512,1024x1024. Defaults to [OpenAI prices](https://openai.com/pricing) for Dall-E.
+* `IMAGE_PRICES`: A comma-separated list with 3 elements of prices for the different image sizes 256x256, 512x512 and 1024x1024. Defaults to [OpenAI prices](https://openai.com/pricing) for Dall-E.
 * `TRANSCRIPTION_PRICE`: USD-price for one minute of audio transcription. Defaults to [OpenAI price](https://openai.com/pricing) for Whisper.
 
 Check out the [official API reference](https://platform.openai.com/docs/api-reference/chat) for more details.

openai_helper.py

@@ -1,5 +1,7 @@
 import datetime
 import logging
+from typing import Tuple, Any
+
 import openai
 
 
@@ -19,12 +21,12 @@ def __init__(self, config: dict):
         self.conversations: dict[int: list] = {}  # {chat_id: history}
         self.last_updated: dict[int: datetime] = {}  # {chat_id: last_update_timestamp}
 
-    def get_chat_response(self, chat_id: int, query: str) -> str:
+    def get_chat_response(self, chat_id: int, query: str) -> tuple[str, str] | str:
         """
         Gets a response from the GPT-3 model.
         :param chat_id: The chat ID
         :param query: The query to send to the model
-        :return: The answer from the model
+        :return: The answer from the model and the number of tokens used
         """
         try:
             if chat_id not in self.conversations or self.__max_age_reached(chat_id):
@@ -75,11 +77,12 @@ def get_chat_response(self, chat_id: int, query: str) -> str:
                     answer += "\n\n---\n" \
                               f"💰 Tokens used: {str(response.usage['total_tokens'])}" \
                               f" ({str(response.usage['prompt_tokens'])} prompt," \
-                              f" {str(response.usage['completion_tokens'])} completion)"            
+                              f" {str(response.usage['completion_tokens'])} completion)"
+
                 return answer, response.usage['total_tokens']
-            else:
-                logging.error('No response from GPT-3')
-                return "⚠️ _An error has occurred_ ⚠️\nPlease try again in a while."
+
+            logging.error('No response from GPT-3')
+            return "⚠️ _An error has occurred_ ⚠️\nPlease try again in a while."
 
         except openai.error.RateLimitError as e:
             logging.exception(e)
@@ -93,11 +96,11 @@ def get_chat_response(self, chat_id: int, query: str) -> str:
             logging.exception(e)
             return f"⚠️ _An error has occurred_ ⚠️\n{str(e)}"
 
-    def generate_image(self, prompt: str) -> str:
+    def generate_image(self, prompt: str) -> tuple[str, str]:
         """
         Generates an image from the given prompt using DALL·E model.
         :param prompt: The prompt to send to the model
-        :return: The image URL
+        :return: The image URL and the image size
         """
         response = openai.Image.create(
             prompt=prompt,

telegram_bot.py

@@ -129,10 +129,10 @@ async def image(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
                 photo=image_url
             )
             # add image request to users usage tracker
-            self.usage[update.message.from_user.id].add_image_request(image_size, self.config['image_prices'])
+            user_id = update.message.from_user.id
+            self.usage[user_id].add_image_request(image_size, self.config['image_prices'])
             # add guest chat request to guest usage tracker
-            allowed_user_ids = self.config['allowed_user_ids'].split(',')
-            if str(update.message.from_user.id) not in allowed_user_ids:
+            if str(user_id) not in self.config['allowed_user_ids'].split(','):
                 self.usage["guests"].add_image_request(image_size, self.config['image_prices'])
 
         except Exception as e:
@@ -194,7 +194,7 @@ async def transcribe(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
 
             # Transcribe the audio file
             transcript = self.openai.transcribe(filename_mp3)
-            
+
             # add transcription seconds to usage tracker
             self.usage[user_id].add_transcription_seconds(audio_track.duration_seconds, self.config['transcription_price'])
             # add guest chat request to guest usage tracker
@@ -217,7 +217,6 @@ async def transcribe(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
                 # add chat request to users usage tracker
                 self.usage[user_id].add_chat_tokens(total_tokens, self.config['token_price'])
                 # add guest chat request to guest usage tracker
-                allowed_user_ids = self.config['allowed_user_ids'].split(',')
                 if str(user_id) not in allowed_user_ids:
                     self.usage["guests"].add_chat_tokens(total_tokens, self.config['token_price'])
 
@@ -258,6 +257,7 @@ async def prompt(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
 
         logging.info(f'New message received from user {update.message.from_user.name}')
         chat_id = update.effective_chat.id
+        user_id = update.message.from_user.id
         prompt = update.message.text
 
         if self.is_group_chat(update):
@@ -272,10 +272,10 @@ async def prompt(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
         response, total_tokens = self.openai.get_chat_response(chat_id=chat_id, query=prompt)
 
         # add chat request to users usage tracker
-        self.usage[update.message.from_user.id].add_chat_tokens(total_tokens, self.config['token_price'])
+        self.usage[user_id].add_chat_tokens(total_tokens, self.config['token_price'])
         # add guest chat request to guest usage tracker
         allowed_user_ids = self.config['allowed_user_ids'].split(',')
-        if str(update.message.from_user.id) not in allowed_user_ids:
+        if str(user_id) not in allowed_user_ids:
             self.usage["guests"].add_chat_tokens(total_tokens, self.config['token_price'])
 
         await context.bot.send_message(
@@ -322,8 +322,7 @@ async def send_budget_reached_message(self, update: Update, context: ContextType
         """
         await context.bot.send_message(
             chat_id=update.effective_chat.id,
-            text=self.budget_limit_message,
-            disable_web_page_preview=True
+            text=self.budget_limit_message
         )
 
     async def error_handler(self, update: object, context: ContextTypes.DEFAULT_TYPE) -> None:
@@ -397,12 +396,9 @@ async def is_within_budget(self, update: Update) -> bool:
                 logging.warning(f'No budget set for user: {update.message.from_user.name} ({user_id}).')
                 return False
             user_budget = float(user_budgets[user_index])
-            cost_month = self.usage[user_id].get_current_cost()[1] 
+            cost_month = self.usage[user_id].get_current_cost()[1]
             # Check if allowed user is within budget
-            if user_budget > cost_month:
-                return True
-            else:
-                return False
+            return user_budget > cost_month
 
         # Check if group member is within budget
         if self.is_group_chat(update):
@@ -412,9 +408,8 @@ async def is_within_budget(self, update: Update) -> bool:
                         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
-                    else:
-                        logging.warning(f'Monthly guest budget for group chats used up.')
-                        return False
+                    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} are not allowed')
         return False
 
@@ -441,7 +436,7 @@ def run(self):
         application.add_handler(CommandHandler('start', self.help))
         application.add_handler(CommandHandler('stats', self.stats))
         application.add_handler(MessageHandler(
-            filters.AUDIO | filters.VOICE | filters.Document.AUDIO | 
+            filters.AUDIO | filters.VOICE | filters.Document.AUDIO |
             filters.VIDEO | filters.VIDEO_NOTE | filters.Document.VIDEO,
             self.transcribe))
         application.add_handler(MessageHandler(filters.TEXT & (~filters.COMMAND), self.prompt))