mirror of
https://github.com/aljazceru/chatgpt-telegram-bot.git
synced 2026-01-31 02:35:57 +01:00
added tracking of image requests and cost
This commit is contained in:
AlexHTW
117
usage_tracker.py
117
usage_tracker.py
@@ -10,11 +10,11 @@ def year_month(date):
|
||||
class UsageTracker:
|
||||
"""
|
||||
UsageTracker class
|
||||
Enables tracking of daily/monthly token usage per user.
|
||||
User files are stored as JSON in /token_usage directory.
|
||||
Enables tracking of daily/monthly usage per user.
|
||||
User files are stored as JSON in /usage_logs directory.
|
||||
JSON example:
|
||||
{
|
||||
"user_name": "user_name",
|
||||
"user_name": "@user_name",
|
||||
"current_cost": {
|
||||
"day": 0.45,
|
||||
"month": 3.23,
|
||||
@@ -43,7 +43,7 @@ class UsageTracker:
|
||||
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, default "usage_logs"
|
||||
:param logs_dir: path to directory of usage logs, defaults to "usage_logs"
|
||||
"""
|
||||
self.user_id = user_id
|
||||
self.logs_dir = logs_dir
|
||||
@@ -66,24 +66,22 @@ class UsageTracker:
|
||||
# token usage functions:
|
||||
|
||||
def add_chat_tokens(self, tokens, tokens_price=0.002):
|
||||
"""
|
||||
Adds used tokens from a request to a users usage history.
|
||||
Updates current cost
|
||||
"""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
|
||||
: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"] += tokens * (tokens_price * 0.001)
|
||||
self.usage["current_cost"]["month"] += tokens * (tokens_price * 0.001)
|
||||
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"] += tokens * (tokens_price * 0.001)
|
||||
self.usage["current_cost"]["month"] += round(tokens * tokens_price / 1000, 6)
|
||||
else:
|
||||
self.usage["current_cost"]["month"] = tokens * (tokens_price * 0.001)
|
||||
self.usage["current_cost"]["day"] = tokens * (tokens_price * 0.001)
|
||||
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
|
||||
@@ -99,6 +97,11 @@ class UsageTracker:
|
||||
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:
|
||||
@@ -109,6 +112,63 @@ class UsageTracker:
|
||||
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:
|
||||
|
||||
@@ -127,29 +187,6 @@ class UsageTracker:
|
||||
second_price = minute_price/60
|
||||
return seconds * second_price
|
||||
|
||||
def get_transcription_seconds_and_cost(self, date=date.today(), minute_price=0.006):
|
||||
# TODO: implement
|
||||
pass
|
||||
|
||||
# image usage functions:
|
||||
|
||||
def add_image_request(self, seconds):
|
||||
# TODO: implement
|
||||
pass
|
||||
|
||||
def get_image_count(self, date=date.today()):
|
||||
# TODO: implement
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def cost_images(image_counts, image_prices=[0.016, 0.018, 0.02]):
|
||||
# TODO: implement
|
||||
pass
|
||||
|
||||
def get_image_counts_and_costs(self, date=date.today(), image_prices=[0.016, 0.018, 0.02]):
|
||||
# TODO: implement
|
||||
pass
|
||||
|
||||
# general functions
|
||||
def get_current_cost(self):
|
||||
pass
|
||||
@@ -158,11 +195,3 @@ class UsageTracker:
|
||||
image_prices=[0.016, 0.018, 0.02]):
|
||||
# TODO: implement
|
||||
pass
|
||||
|
||||
"""
|
||||
testing
|
||||
user = UsageTracker("hi", "my_name")
|
||||
user.add_chat_tokens(12)
|
||||
print(user.get_token_usage(date.fromisoformat('2023-05-03')))
|
||||
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user