diff --git a/src/message_manager.py b/src/message_manager.py index d76dff0..fac897e 100644 --- a/src/message_manager.py +++ b/src/message_manager.py @@ -13,7 +13,7 @@ class MessageManager: config_dict = {} openai_parser = None access_manager = None - + def __init__(self, access_manager): self.openai_parser = OpenAIParser() self.access_manager = access_manager @@ -34,7 +34,7 @@ def get_response(self, id, user, message): self.userDict[id] = ChatSession(t, message) else: self.userDict[id].update(t, message, "user") - + # send user info for statistics (answer, usage) = self.__sendMessage( user, self.userDict[id].messageList) @@ -78,11 +78,17 @@ def get_transcript(self, user, audio_file): print(e) return "" + def get_usage(self): + try: + return f"Usage in the last 30 days: {self.openai_parser.get_usage():.2f} $" + except Exception as e: + return f"Error: {e}" + def __sendMessage(self, user, messageList): ans = self.openai_parser.get_response(user, messageList) return ans - - + + if __name__ == "__main__": acm = AccessManager() msm = MessageManager(acm) diff --git a/src/openai_parser.py b/src/openai_parser.py index 572eec5..52dc2d9 100644 --- a/src/openai_parser.py +++ b/src/openai_parser.py @@ -14,6 +14,7 @@ import openai, json, os import datetime +import requests class OpenAIParser: @@ -34,7 +35,7 @@ def _get_single_response(self, message): {"role": "user", "content": message} ]) return response["choices"][0]["message"]["content"] - + def get_response(self, userid, context_messages): context_messages.insert(0, {"role": "system", "content": "You are a helpful assistant"}) try: @@ -57,7 +58,27 @@ def image_generation(self, userid, prompt): usage = 1 # reserve for future use return (image_url, usage) + def get_usage(self): + """Get the usage of the API for the last 30 days. + + Example API call: + https://api.openai.com/dashboard/billing/usage?end_date=2023-04-01&start_date=2023-03-01 + + The API call return a json with a key total_usage with contains the cents of dollar used in the specified period. + + Return the usage in dollars. + """ + today = datetime.date.today() + last_month = today - datetime.timedelta(days=30) + url = "https://api.openai.com/dashboard/billing/usage?end_date={}&start_date={}".format(today, last_month) + headers = {"Authorization": "Bearer {}".format(self.config_dict["openai_api_key"])} + + req = requests.get(url, headers=headers) + usage = req.json()["total_usage"]/100 + + return usage + if __name__ == "__main__": openai_parser = OpenAIParser() # print(openai_parser._get_single_response("Tell me a joke.")) - print(openai_parser.get_response("123", [{"role": "user", "content": "Tell me a joke."}])) \ No newline at end of file + print(openai_parser.get_response("123", [{"role": "user", "content": "Tell me a joke."}])) diff --git a/src/telegram_message_parser.py b/src/telegram_message_parser.py index 19aaa88..4109e94 100644 --- a/src/telegram_message_parser.py +++ b/src/telegram_message_parser.py @@ -54,6 +54,7 @@ def add_handlers(self): self.bot.add_handler(CommandHandler("start", self.start)) self.bot.add_handler(CommandHandler("clear", self.clear_context)) self.bot.add_handler(CommandHandler("getid", self.get_user_id)) + self.bot.add_handler(CommandHandler("usage", self.usage)) # special message handlers if self.config_dict["enable_voice"]: @@ -104,8 +105,8 @@ async def chat_text(self, update: Update, context: ContextTypes.DEFAULT_TYPE): # send message to openai response = self.message_manager.get_response( - str(update.effective_chat.id), - str(update.effective_user.id), + str(update.effective_chat.id), + str(update.effective_user.id), message ) # reply response to user @@ -134,8 +135,8 @@ async def chat_text_command(self, update: Update, context: ContextTypes.DEFAULT_ # send message to openai response = self.message_manager.get_response( - str(update.effective_chat.id), - str(update.effective_user.id), + str(update.effective_chat.id), + str(update.effective_user.id), message ) @@ -177,13 +178,13 @@ async def chat_voice(self, update: Update, context: ContextTypes.DEFAULT_TYPE): subprocess.call( ['ffmpeg', '-i', file_id + '.ogg', file_id + '.wav'], - stdout=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL ) with open(file_id + ".wav", "rb") as audio_file: transcript = self.message_manager.get_transcript( - str(update.effective_user.id), + str(update.effective_user.id), audio_file ) os.remove(file_id + ".ogg") @@ -194,8 +195,8 @@ async def chat_voice(self, update: Update, context: ContextTypes.DEFAULT_TYPE): return response = self.message_manager.get_response( - str(update.effective_chat.id), - str(update.effective_user.id), + str(update.effective_chat.id), + str(update.effective_user.id), transcript ) await update.message.reply_text("\"" + transcript + "\"\n\n" + response) @@ -207,7 +208,7 @@ async def image_generation(self, update: Update, context: ContextTypes.DEFAULT_T # send prompt to openai image generation and get image url image_url, prompt = self.message_manager.get_generated_image_url( - str(update.effective_user.id), + str(update.effective_user.id), message ) @@ -233,7 +234,7 @@ async def image_generation(self, update: Update, context: ContextTypes.DEFAULT_T # inline text messages async def inline_query(self, update: Update, context: ContextTypes.DEFAULT_TYPE): # get query message - query = update.inline_query.query + query = update.inline_query.query if query == "": return @@ -266,7 +267,7 @@ async def inline_query(self, update: Update, context: ContextTypes.DEFAULT_TYPE) # await update.inline_query.answer(results, cache_time=0, is_personal=True, switch_pm_text="Chat Privately 🤫", switch_pm_parameter="start") await update.inline_query.answer(results, cache_time=0, is_personal=True) - + async def inline_query_result_chosen(self, update: Update, context: ContextTypes.DEFAULT_TYPE): # invalid user won't get a response try: @@ -276,7 +277,7 @@ async def inline_query_result_chosen(self, update: Update, context: ContextTypes inline_message_id = update.chosen_inline_result.inline_message_id query = update.chosen_inline_result.query # query_id = query[query.find("My_Memory_ID: ")+14:query.find("\n=======")] - + # if query_id == "": # if no query_id, generate one # query_id = str(uuid4()) # else: # if query_id, remove it from query @@ -298,7 +299,7 @@ async def inline_query_result_chosen(self, update: Update, context: ContextTypes ) except Exception as e: pass - + # file and photo messages @@ -355,6 +356,14 @@ async def get_user_id(self, update: Update, context: ContextTypes.DEFAULT_TYPE): text=str(update.effective_user.id) ) + async def usage(self, update: Update, context: ContextTypes.DEFAULT_TYPE): + # Get OpenAI API usage + usage = self.message_manager.get_usage() + await context.bot.send_message( + chat_id=update.effective_chat.id, + text=str(usage) + ) + # unknown command async def unknown(self, update: Update, context: ContextTypes.DEFAULT_TYPE): await context.bot.send_message(