Merge pull request #49 from AlexHTW/track-token-usage-per-user

Track usage per user

.env.example

@@ -23,4 +23,13 @@ MAX_HISTORY_SIZE=10
 MAX_CONVERSATION_AGE_MINUTES=180
 
 # Whether to answer to voice messages with the transcript or with a ChatGPT response of the transcript
-VOICE_REPLY_WITH_TRANSCRIPT_ONLY=true
+VOICE_REPLY_WITH_TRANSCRIPT_ONLY=true
+
+# USD-price per 1000 tokens for cost information in usage statistics, see https://openai.com/pricing
+TOKEN_PRICE=0.002
+
+# USD-prices per image for the sizes 256x256,512x512,1024x1024 for cost information in usage statistics
+IMAGE_PRICES="0.016,0.018,0.02"
+
+# USD-price for 1 minute of audio transcription for cost information in usage statistics
+TRANSCRIPTION_PRICE=0.006

.gitignore

@@ -2,3 +2,4 @@
 /.idea
 .env
 .DS_Store
+/usage_logs

README.md

@@ -24,6 +24,8 @@ A [Telegram bot](https://core.telegram.org/bots/api) that integrates with OpenAI
 - [x] (NEW!) Automatic conversation summary to avoid excessive token usage (fixes [#34](https://github.com/n3d1117/chatgpt-telegram-bot/issues/34))
 - [x] (NEW!) Group chat support with inline queries 
   - To use this feature, enable inline queries for your bot in BotFather via the `/setinline` [command](https://core.telegram.org/bots/inline)
+- [x] (NEW!) Track token usage per user (stored in /token_usage/`user_id`.json)
+- [x] (NEW!) Get personal token usage statistics and cost per day/month via the `/stats` command
 
 ## Additional features - help needed!
 - [ ] Add stream support ([#43](https://github.com/n3d1117/chatgpt-telegram-bot/issues/43))
@@ -54,6 +56,9 @@ MAX_TOKENS=2000 # Defaults to 1200
 MAX_HISTORY_SIZE=15 # Defaults to 10
 MAX_CONVERSATION_AGE_MINUTES=120 # Defaults to 180 (3h)
 VOICE_REPLY_WITH_TRANSCRIPT_ONLY=false # Defaults to true
+TOKEN_PRICE=0.002 # Defaults to 0.002, current price: https://openai.com/pricing
+IMAGE_PRICES="0.016,0.018,0.02" # Defaults to OpenAI Dall-E pricing for sizes 256x256,512x512,1024x1024
+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))
@@ -65,6 +70,8 @@ VOICE_REPLY_WITH_TRANSCRIPT_ONLY=false # Defaults to true
 * `MAX_HISTORY_SIZE`: Max number of messages to keep in memory, after which the conversation will be summarised to avoid excessive token usage ([#34](https://github.com/n3d1117/chatgpt-telegram-bot/issues/34))
 * `MAX_CONVERSATION_AGE_MINUTES`: Maximum number of minutes a conversation should live, after which the conversation will be reset to avoid excessive token usage
 * `VOICE_REPLY_WITH_TRANSCRIPT_ONLY`: Whether to answer to voice messages with the transcript only or with a ChatGPT response of the transcript ([#38](https://github.com/n3d1117/chatgpt-telegram-bot/issues/38))
+* `TOKEN_PRICE`: USD-price per 1000 tokens for cost information in usage statistics
+
 
 <details><summary>Additional model parameters can be configured from the `main.py` file.</summary>
 

main.py

@@ -61,8 +61,10 @@ def main():
         'allowed_user_ids': os.environ.get('ALLOWED_TELEGRAM_USER_IDS', '*'),
         'proxy': os.environ.get('PROXY', None),
         'voice_reply_transcript': os.environ.get('VOICE_REPLY_WITH_TRANSCRIPT_ONLY', 'true').lower() == 'true',
+        'token_price': float(os.environ.get('TOKEN_PRICE', 0.002)),
+        'image_prices': [float(i) for i in os.environ.get('IMAGE_PRICES',"0.016,0.018,0.02").split(",")],
+        'transcription_price': float(os.environ.get('TOKEN_PRICE', 0.002)),
     }
-
     # Setup and run ChatGPT and Telegram bot
     openai_helper = OpenAIHelper(config=openai_config)
     telegram_bot = ChatGPT3TelegramBot(config=telegram_config, openai=openai_helper)

openai_helper.py

@@ -75,9 +75,8 @@ class OpenAIHelper:
                     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)"
-
-                return answer
+                              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."
@@ -105,7 +104,7 @@ class OpenAIHelper:
             n=1,
             size=self.config['image_size']
         )
-        return response['data'][0]['url']
+        return response['data'][0]['url'], self.config['image_size']
 
     def transcribe(self, filename):
         """

telegram_bot.py

@@ -8,6 +8,7 @@ from telegram.ext import ApplicationBuilder, ContextTypes, CommandHandler, Messa
 
 from pydub import AudioSegment
 from openai_helper import OpenAIHelper
+from usage_tracker import UsageTracker
 
 
 class ChatGPT3TelegramBot:
@@ -26,10 +27,12 @@ class ChatGPT3TelegramBot:
         self.commands = [
             BotCommand(command='help', description='Show this help message'),
             BotCommand(command='reset', description='Reset the conversation'),
-            BotCommand(command='image', description='Generate image from prompt (e.g. /image cat)')
+            BotCommand(command='image', description='Generate image from prompt (e.g. /image cat)'),
+            BotCommand(command='stats', description='Get your current usage statistics')
         ]
         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.usage = {}
 
     async def help(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
         """
@@ -45,6 +48,40 @@ class ChatGPT3TelegramBot:
                     "Open source at https://github.com/n3d1117/chatgpt-telegram-bot"
         await update.message.reply_text(help_text, disable_web_page_preview=True)
 
+
+    async def stats(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
+        """
+        Returns token usage statistics for current day and month.
+        """
+        if not await self.is_allowed(update):
+            logging.warning(f'User {update.message.from_user.name} is not allowed to request their usage statistics')
+            await self.send_disallowed_message(update, context)
+            return
+
+        logging.info(f'User {update.message.from_user.name} requested their token usage statistics')
+        
+        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)
+
+        tokens_today, tokens_month = self.usage[user_id].get_token_usage()
+        images_today, images_month = self.usage[user_id].get_image_count()
+        transcribe_durations = self.usage[user_id].get_transcription_duration()
+        cost_today, cost_month = self.usage[user_id].get_current_cost()
+        
+        usage_text = f"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}.\n"+\
+                     f"\n----------------------------\n\n"+\
+                     f"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}."
+        await update.message.reply_text(usage_text)
+
     async def reset(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
         """
         Resets the conversation.
@@ -79,12 +116,18 @@ class ChatGPT3TelegramBot:
 
         await context.bot.send_chat_action(chat_id=chat_id, action=constants.ChatAction.UPLOAD_PHOTO)
         try:
-            image_url = self.openai.generate_image(prompt=image_query)
+            image_url, image_size = self.openai.generate_image(prompt=image_query)
             await context.bot.send_photo(
                 chat_id=chat_id,
                 reply_to_message_id=update.message.message_id,
                 photo=image_url
             )
+            # add image request to usage tracker
+            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)
+            self.usage[user_id].add_image_request(image_size, self.config['image_prices'])
+
         except Exception as e:
             logging.exception(e)
             await context.bot.send_message(
@@ -123,6 +166,10 @@ class ChatGPT3TelegramBot:
 
         filename_mp3 = f'{filename}.mp3'
 
+        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)
+
         try:
             if update.message.voice:
                 media_file = await context.bot.get_file(update.message.voice.file_id)
@@ -140,7 +187,8 @@ class ChatGPT3TelegramBot:
 
             # 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'])
             if self.config['voice_reply_transcript']:
                 # Send the transcript
                 await context.bot.send_message(
@@ -151,7 +199,9 @@ class ChatGPT3TelegramBot:
                 )
             else:
                 # Send the response of the transcript
-                response = self.openai.get_chat_response(chat_id=chat_id, query=transcript)
+                response, total_tokens = self.openai.get_chat_response(chat_id=chat_id, query=transcript)
+                # add chat request to usage tracker
+                self.usage[user_id].add_chat_tokens(total_tokens, self.config['token_price'])
                 await context.bot.send_message(
                     chat_id=chat_id,
                     reply_to_message_id=update.message.message_id,
@@ -185,7 +235,14 @@ class ChatGPT3TelegramBot:
         chat_id = update.effective_chat.id
 
         await context.bot.send_chat_action(chat_id=chat_id, action=constants.ChatAction.TYPING)
-        response = self.openai.get_chat_response(chat_id=chat_id, query=update.message.text)
+        response, total_tokens = self.openai.get_chat_response(chat_id=chat_id, query=update.message.text)
+
+        # add chat request to usage tracker
+        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)
+        self.usage[user_id].add_chat_tokens(total_tokens, self.config['token_price'])
+
         await context.bot.send_message(
             chat_id=chat_id,
             reply_to_message_id=update.message.message_id,
@@ -294,6 +351,7 @@ class ChatGPT3TelegramBot:
         application.add_handler(CommandHandler('help', self.help))
         application.add_handler(CommandHandler('image', self.image))
         application.add_handler(CommandHandler('start', self.help))
+        application.add_handler(CommandHandler('stats', self.stats))
         application.add_handler(MessageHandler(filters.VOICE | filters.AUDIO | filters.VIDEO, self.transcribe))
         application.add_handler(MessageHandler(filters.TEXT & (~filters.COMMAND), self.prompt))
         application.add_handler(InlineQueryHandler(self.inline_query, chat_types=[

usage_tracker.py

@@ -0,0 +1,242 @@
+import os.path
+import pathlib
+import json
+from datetime import date
+
+def year_month(date):
+    # extract string of year-month from date, eg: '2023-03'
+    return str(date)[:7]
+
+class UsageTracker:
+    """
+    UsageTracker class
+    Enables tracking of daily/monthly usage per user.
+    User files are stored as JSON in /usage_logs directory.
+    JSON example:
+    {
+        "user_name": "@user_name",
+        "current_cost": {
+            "day": 0.45, 
+            "month": 3.23, 
+            "last_update": "2023-03-14"},
+        "usage_history": {
+            "chat_tokens": {
+                "2023-03-13": 520,
+                "2023-03-14": 1532
+            },
+            "transcription_seconds": {
+                "2023-03-13": 125,
+                "2023-03-14": 64
+            },
+            "number_images": {
+                "2023-03-12": [0, 2, 3],
+                "2023-03-13": [1, 2, 3],
+                "2023-03-14": [0, 1, 2]
+            }
+        }
+    }
+    """
+
+    def __init__(self, user_id, user_name, logs_dir="usage_logs"):
+        """
+        Initializes UsageTracker for a user with current date. 
+        Loads usage data from usage log file.
+        :param user_id: Telegram ID of the user
+        :param user_name: Telegram user name
+        :param logs_dir: path to directory of usage logs, defaults to "usage_logs"
+        """
+        self.user_id = user_id
+        self.logs_dir = logs_dir
+        # path to usage file of given user
+        self.user_file = f"{logs_dir}/{user_id}.json"
+        
+        if os.path.isfile(self.user_file):
+            with open(self.user_file, "r") as file:
+                self.usage = json.load(file)
+        else:
+            # ensure directory exists 
+            pathlib.Path(logs_dir).mkdir(exist_ok=True)
+            # 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())},
+                "usage_history": {"chat_tokens": {}, "transcription_seconds": {}, "number_images": {}}
+            }
+
+    # token usage functions:
+    
+    def add_chat_tokens(self, tokens, tokens_price=0.002):
+        """Adds used tokens from a request 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
+        """
+        today = date.today()
+        last_update = date.fromisoformat(self.usage["current_cost"]["last_update"])
+        # 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"]["month"] += round(tokens * tokens_price / 1000, 6)
+        else:
+            if today.month == last_update.month:
+                self.usage["current_cost"]["month"] += round(tokens * tokens_price / 1000, 6)
+            else:
+                self.usage["current_cost"]["month"] = round(tokens * tokens_price / 1000, 6)
+            self.usage["current_cost"]["day"] = round(tokens * tokens_price / 1000, 6)
+            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
+            self.usage["usage_history"]["chat_tokens"][str(today)] += tokens
+        else:
+            # create new entry for current date
+            self.usage["usage_history"]["chat_tokens"][str(today)] = tokens
+        
+        # write updated token usage to user file
+        with open(self.user_file, "w") as outfile:
+            json.dump(self.usage, outfile)
+
+    def get_token_usage(self, date=date.today()):
+        """Get token amount used on day and month of date-
+
+        :param date: date of interest, defaults to date.today()
+        :return: total number of tokens used per day and per month
+        """        
+        if str(date) in self.usage["usage_history"]["chat_tokens"]:
+            usage_day = self.usage["usage_history"]["chat_tokens"][str(date)]
+        else:
+            usage_day = 0
+        month = str(date)[:7] # year-month as string
+        usage_month = 0
+        for date, tokens in self.usage["usage_history"]["chat_tokens"].items():
+            if date.startswith(month):
+                usage_month += tokens
+        return usage_day, usage_month
+    
+    # image usage functions:
+
+    def add_image_request(self, image_size, image_prices="0.016,0.018,0.02"):
+        """Add image request to users usage history and update current costs.
+
+        :param image_size: requested image size
+        :param image_prices: prices for images of sizes ["256x256", "512x512", "1024x1024"],
+                             defaults to [0.016, 0.018, 0.02]
+        """        
+        sizes = ["256x256", "512x512", "1024x1024"]
+        requested_size = sizes.index(image_size)
+        image_cost = image_prices[requested_size]
+
+        today = date.today()
+        last_update = date.fromisoformat(self.usage["current_cost"]["last_update"])
+        # add current cost, update new day
+        if today == last_update:
+            self.usage["current_cost"]["day"] += image_cost
+            self.usage["current_cost"]["month"] += image_cost
+        else:
+            if today.month == last_update.month:
+                self.usage["current_cost"]["month"] += image_cost
+            else:
+                self.usage["current_cost"]["month"] = image_cost
+            self.usage["current_cost"]["day"] = image_cost
+            self.usage["current_cost"]["last_update"] = str(today)
+
+        # update usage_history
+        if str(today) in self.usage["usage_history"]["number_images"]:
+            # add token usage to existing date
+            self.usage["usage_history"]["number_images"][str(today)][requested_size] += 1
+        else:
+            # create new entry for current date
+            self.usage["usage_history"]["number_images"][str(today)] = [0, 0, 0]
+            self.usage["usage_history"]["number_images"][str(today)][requested_size] += 1
+        
+        # write updated image number to user file
+        with open(self.user_file, "w") as outfile:
+            json.dump(self.usage, outfile)
+
+    def get_image_count(self, date=date.today()):
+        """Get number of images requested on day and month of date.
+
+        :param date: date of interest, defaults to date.today()
+        :return: total number of images requested per day and per month
+        """      
+        if str(date) in self.usage["usage_history"]["number_images"]:
+            usage_day = sum(self.usage["usage_history"]["number_images"][str(date)])
+        else:
+            usage_day = 0
+        month = str(date)[:7] # year-month as string
+        usage_month = 0
+        for date, images in self.usage["usage_history"]["number_images"].items():
+            if date.startswith(month):
+                usage_month += sum(images)
+        return usage_day, usage_month
+
+    # transcription usage functions:
+
+    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
+        """
+        today = date.today()
+        last_update = date.fromisoformat(self.usage["current_cost"]["last_update"])
+        # 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"]["month"] += round(seconds * minute_price / 60, 2)
+        else:
+            if today.month == last_update.month:
+                self.usage["current_cost"]["month"] += round(seconds * minute_price / 60, 2)
+            else:
+                self.usage["current_cost"]["month"] = round(seconds * minute_price / 60, 2)
+            self.usage["current_cost"]["day"] = round(seconds * minute_price / 60, 2)
+            self.usage["current_cost"]["last_update"] = str(today)
+
+        # update usage_history
+        if str(today) in self.usage["usage_history"]["transcription_seconds"]:
+            # add requested seconds to existing date
+            self.usage["usage_history"]["transcription_seconds"][str(today)] += seconds
+        else:
+            # create new entry for current date
+            self.usage["usage_history"]["transcription_seconds"][str(today)] = seconds
+        
+        # write updated token usage to user file
+        with open(self.user_file, "w") as outfile:
+            json.dump(self.usage, outfile)
+
+    def get_transcription_duration(self, date=date.today()):
+        """Get minutes and seconds of audio transcribed on day and month of date.
+
+        :param date: date of interest, defaults to date.today()
+        :return: total amount of time transcribed per day and per month (4 values)
+        """      
+        if str(date) in self.usage["usage_history"]["transcription_seconds"]:
+            seconds_day = self.usage["usage_history"]["transcription_seconds"][str(date)]
+        else:
+            seconds_day = 0
+        month = str(date)[:7] # year-month as string
+        seconds_month = 0
+        for date, seconds in self.usage["usage_history"]["transcription_seconds"].items():
+            if date.startswith(month):
+                seconds_month += seconds
+        minutes_day, seconds_day = divmod(seconds_day, 60)
+        minutes_month, seconds_month = divmod(seconds_month, 60)
+        return int(minutes_day), round(seconds_day, 2), int(minutes_month), round(seconds_month, 2)
+    
+    # general functions
+    def get_current_cost(self):
+        """Get total USD amount of all requests of the current day and month
+
+        :return: cost of current day and month
+        """  
+        today = date.today()
+        last_update = date.fromisoformat(self.usage["current_cost"]["last_update"])
+        if today == last_update:
+            cost_day = self.usage["current_cost"]["day"]
+            cost_month = self.usage["current_cost"]["month"]
+        else:
+            cost_day = 0.0
+            if today.month == last_update.month:
+                cost_month = self.usage["current_cost"]["month"]
+            else:
+                cost_month = 0.0
+        return round(cost_day, 3), round(cost_month, 3)