/balance command + limit of dialog context + user state move to user data

README.md

@@ -9,10 +9,16 @@ This repo is ChatGPT re-created with GPT-3.5 LLM as Telegram Bot. **And it works
 - Low latency replies (it usually takes about 3-5 seconds) 
 - No request limits
 - Code highlighting
-- Different chat modes (👩🏼‍🎓 Assistant, 👩🏼‍💻 Code Assistant, 🎬 Movie Expert)
-- `/retry` command to regenerate last bot answer
-- Control of allowed Telegram users
-- *Next up*: warn user that the size of the context is close to the maximum
+- Different chat modes: 👩🏼‍🎓 Assistant, 👩🏼‍💻 Code Assistant, 🎬 Movie Expert. More soon
+- List of allowed Telegram users
+- Track $ balance spent on OpenAI API
+
+## Bot commands
+- `/retry` – Regenerate last bot answer
+- `/new` – Start new dialog
+- `/mode` – Select chat mode
+- `/balance` – Show balance
+- `/help` – Show help
 
 ## Setup
 1. Get your [OpenAI API](https://openai.com/api/) key

bot.py

@@ -26,10 +26,11 @@ import config
 logger = logging.getLogger(__name__)
 
 HELP_MESSAGE = """Commands:
-⚪ /retry – regenerate last bot answer
-⚪ /reset – reset chat context
-⚪ /mode – select chat mode
-⚪ /help – show help
+⚪ /retry – Regenerate last bot answer
+⚪ /new – Start new dialog
+⚪ /mode – Select chat mode
+⚪ /balance – Show balance
+⚪ /help – Show help
 """
 
 
@@ -54,45 +55,64 @@ async def help_handle(update: Update, context: CallbackContext):
 async def retry_handle(update: Update, context: CallbackContext):
     context.user_data["last_interation_timestamp"] = time.time()
 
-    if len(context.user_data["chatgpt"].context) == 0:
+    if len(context.user_data["chat_context"]) == 0:
         await update.message.reply_text("No message to retry 🤷‍♂️")
         return
 
-    last_message = context.user_data["chatgpt"].context.pop()
-    await message_handle(update, context, message=last_message["user"], use_reset_timeout=False)
+    last_chat_context_item = context.user_data["chat_context"].pop()
+    await message_handle(update, context, message=last_chat_context_item["user"], use_new_dialog_timeout=False)
 
 
-async def message_handle(update: Update, context: CallbackContext, message=None, use_reset_timeout=True):
+async def message_handle(update: Update, context: CallbackContext, message=None, use_new_dialog_timeout=True):
     utils.init_user(update, context)
 
-    # reset timeout
-    if use_reset_timeout:
-        if time.time() - context.user_data["last_interation_timestamp"] > config.reset_timeout:
-            context.user_data["chatgpt"].reset_context()
-            await update.message.reply_text("Chat context is reset due to timeout ✅")
-        context.user_data["last_interation_timestamp"] = time.time()
+    # new dialog timeout
+    if use_new_dialog_timeout:
+        if time.time() - context.user_data["last_interation_timestamp"] > config.new_dialog_timeout:
+            context.user_data["chat_context"] = []
+            await update.message.reply_text("Starting new dialog due to timeout ✅")
+    context.user_data["last_interation_timestamp"] = time.time()
 
+    # send typing action
     await update.message.chat.send_action(action="typing")
 
     try:
         message = message or update.message.text
-        answer, prompt = context.user_data["chatgpt"].send_message(message)
-        await update.message.reply_text(answer, parse_mode=ParseMode.HTML)
+        answer, prompt, chat_context, n_used_tokens, n_first_chat_context_messages_removed = chatgpt.ChatGPT().send_message(
+            message,
+            chat_context=context.user_data["chat_context"],
+            chat_mode=context.user_data["chat_mode"]
+        )
     except Exception as e:
         error_text = f"Something went wrong during completion. Reason: {e}"
         logger.error(error_text)
         await update.message.reply_text(error_text)
+        return
+
+    # update user data
+    context.user_data["chat_context"] = chat_context
+    context.user_data["total_n_used_tokens"] += n_used_tokens
+
+    # send message if some messages were removed from the context
+    if n_first_chat_context_messages_removed > 0:
+        if n_first_chat_context_messages_removed == 1:
+            text = "✍️ <i>Note:</i> Your current dialog is too long, so your <b>first message</b> was removed from the context.\n Send /new command to start new dialog"
+        else:
+            text = f"✍️ <i>Note:</i> Your current dialog is too long, so <b>{n_first_chat_context_messages_removed} first messages</b> were removed from the context.\n Send /new command to start new dialog"
+        await update.message.reply_text(text, parse_mode=ParseMode.HTML)
+
+    await update.message.reply_text(answer, parse_mode=ParseMode.HTML)
 
 
-async def reset_handle(update: Update, context: CallbackContext):
+async def new_dialog_handle(update: Update, context: CallbackContext):
     utils.init_user(update, context)
     context.user_data["last_interation_timestamp"] = time.time()
 
-    context.user_data["chatgpt"].reset_context()
-    await update.message.reply_text("Chat context is reset ✅")
+    context.user_data["chat_context"] = []
+    await update.message.reply_text("Starting new dialog ✅")
 
-    chat_mode_key = context.user_data["chatgpt"].chat_mode
-    await update.message.reply_text(f"{chatgpt.CHAT_MODES[chat_mode_key]['welcome_message']}", parse_mode=ParseMode.HTML)
+    chat_mode = context.user_data["chat_mode"]
+    await update.message.reply_text(f"{chatgpt.CHAT_MODES[chat_mode]['welcome_message']}", parse_mode=ParseMode.HTML)
 
 
 async def show_chat_modes_handle(update: Update, context: CallbackContext):
@@ -100,29 +120,41 @@ async def show_chat_modes_handle(update: Update, context: CallbackContext):
     context.user_data["last_interation_timestamp"] = time.time()
 
     keyboard = []
-    for mode_key, mode_dict in chatgpt.CHAT_MODES.items():
-        keyboard.append([InlineKeyboardButton(mode_dict["name"], callback_data=f"set_mode|{mode_key}")])
+    for chat_mode, chat_mode_dict in chatgpt.CHAT_MODES.items():
+        keyboard.append([InlineKeyboardButton(chat_mode_dict["name"], callback_data=f"set_chat_mode|{chat_mode}")])
     reply_markup = InlineKeyboardMarkup(keyboard)
 
-    await update.message.reply_text("Choose chat mode:", reply_markup=reply_markup)
+    await update.message.reply_text("Select chat mode:", reply_markup=reply_markup)
 
 
 async def set_chat_mode_handle(update: Update, context: CallbackContext):
     query = update.callback_query
-
     await query.answer()
-    await query.edit_message_text(text="See you next time!")
 
-    chat_mode_key = query.data.split("|")[1]
+    chat_mode = query.data.split("|")[1]
+
+    context.user_data["chat_mode"] = chat_mode
+    context.user_data["chat_context"] = []
 
-    context.user_data["chatgpt"].set_chat_mode(chat_mode_key)
-    context.user_data["chatgpt"].reset_context()
     await query.edit_message_text(
-        f"<b>{chatgpt.CHAT_MODES[chat_mode_key]['name']}</b> chat mode is set",
+        f"<b>{chatgpt.CHAT_MODES[chat_mode]['name']}</b> chat mode is set",
         parse_mode=ParseMode.HTML
     )
 
-    await query.edit_message_text(f"{chatgpt.CHAT_MODES[chat_mode_key]['welcome_message']}", parse_mode=ParseMode.HTML)
+    await query.edit_message_text(f"{chatgpt.CHAT_MODES[chat_mode]['welcome_message']}", parse_mode=ParseMode.HTML)
+
+
+async def show_balance_handle(update: Update, context: CallbackContext):
+    utils.init_user(update, context)
+    context.user_data["last_interation_timestamp"] = time.time()
+
+    total_n_used_tokens = context.user_data['total_n_used_tokens']
+    total_spent_dollars = total_n_used_tokens * (0.01 / 1000)
+
+    text = f"You spent <b>{total_spent_dollars:.03f}$</b>\n"
+    text += f"You used <b>{total_n_used_tokens}</b> tokens <i>(price: 0.01$ per 1000 tokens)</i>\n"
+
+    await update.message.reply_text(text, parse_mode=ParseMode.HTML)
 
 
 async def error_handler(update: Update, context: CallbackContext) -> None:
@@ -163,10 +195,12 @@ def run_bot() -> None:
 
     application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND & user_filter, message_handle))
     application.add_handler(CommandHandler("retry", retry_handle, filters=user_filter))
-    application.add_handler(CommandHandler("reset", reset_handle, filters=user_filter))
-
+    application.add_handler(CommandHandler("new", new_dialog_handle, filters=user_filter))
+    
     application.add_handler(CommandHandler("mode", show_chat_modes_handle, filters=user_filter))
-    application.add_handler(CallbackQueryHandler(set_chat_mode_handle, pattern="^set_mode"))
+    application.add_handler(CallbackQueryHandler(set_chat_mode_handle, pattern="^set_chat_mode"))
+
+    application.add_handler(CommandHandler("balance", show_balance_handle, filters=user_filter))
     
     application.add_error_handler(error_handler)
     

chatgpt.py

@@ -26,51 +26,59 @@ CHAT_MODES = {
 
 
 class ChatGPT:
-    def __init__(self, chat_mode="assistant"):
-        self.chat_mode = chat_mode
-        self.context = []
+    def __init__(self):
+        pass
     
-    def send_message(self, message):
-        prompt = self._generate_prompt(message)
-        r = openai.Completion.create(
-            engine="text-davinci-003",
-            prompt=prompt,
-            temperature=0.7,
-            max_tokens=1000,
-            top_p=1,
-            frequency_penalty=0,
-            presence_penalty=0,
-        )
-        answer = r.choices[0].text
-        answer = answer.strip()
-
-        # update context
-        self.context.append({"user": message, "chatgpt": answer})
-
-        return answer, prompt
-
-    def reset_context(self):
-        self.context = []
-
-    def set_chat_mode(self, chat_mode):
+    def send_message(self, message, chat_context=[], chat_mode="assistant"):
         if chat_mode not in CHAT_MODES.keys():
             raise ValueError(f"Chat mode {chat_mode} is not supported")
 
-        self.chat_mode = chat_mode
+        chat_context_len_before = len(chat_context)
+        answer = None
+        while answer is None:
+            prompt = self._generate_prompt(message, chat_context, chat_mode)
+            try:
+                r = openai.Completion.create(
+                    engine="text-davinci-003",
+                    prompt=prompt,
+                    temperature=0.7,
+                    max_tokens=1000,
+                    top_p=1,
+                    frequency_penalty=0,
+                    presence_penalty=0,
+                )
+                answer = r.choices[0].text
+                answer = answer.strip()
 
-    def _generate_prompt(self, message):
-        prompt = CHAT_MODES[self.chat_mode]["prompt_start"]
+                n_used_tokens = r.usage.total_tokens
+
+            except openai.error.InvalidRequestError as e:  # too many tokens
+                if len(chat_context) == 0:
+                    raise ValueError("chat_context is reduced to zero, but still has too many tokens to make completion") from e
+
+                # forget first message in chat_context
+                chat_context = chat_context[1:]
+
+        n_first_chat_context_messages_removed = chat_context_len_before - len(chat_context)
+
+        # update chat_context
+        chat_context.append({"user": message, "chatgpt": answer})
+
+        return answer, prompt, chat_context, n_used_tokens, n_first_chat_context_messages_removed
+
+    def _generate_prompt(self, message, chat_context, chat_mode):
+        prompt = CHAT_MODES[chat_mode]["prompt_start"]
         prompt += "\n\n"
 
-        # chat history
-        if len(self.context) > 0:
+        # add chat context
+        if len(chat_context) > 0:
             prompt += "Chat:\n"
-            for context_item in self.context:
-                prompt += f"User: {context_item['user']}\n"
-                prompt += f"ChatGPT: {context_item['chatgpt']}\n"
+            for chat_context_item in chat_context:
+                prompt += f"User: {chat_context_item['user']}\n"
+                prompt += f"ChatGPT: {chat_context_item['chatgpt']}\n"
 
         # current message
         prompt += f"User: {message}\n"
         prompt += "ChatGPT: "
 
-        return prompt
+        return prompt

config.example.yml

@@ -2,4 +2,4 @@ telegram_token: ""
 openai_api_key: ""
 allowed_telegram_usernames: []  # if empty, the bot is available to anyone
 persistence_path: "./persistence.pkl"  # path where to store user data
-reset_timeout: 600  # chat context is reset after timeout (in seconds)
+new_dialog_timeout: 600  # new dialog starts after timeout (in seconds)

config.py

@@ -8,4 +8,4 @@ telegram_token = config["telegram_token"]
 openai_api_key = config["openai_api_key"]
 allowed_telegram_usernames = config["allowed_telegram_usernames"]
 persistence_path = config["persistence_path"]
-reset_timeout = config["reset_timeout"]
+new_dialog_timeout = config["new_dialog_timeout"]

utils.py

@@ -9,8 +9,14 @@ import config
 
 def init_user(update: Update, context: CallbackContext):
     # init chatgpt
-    if "chatgpt" not in context.user_data:
-        context.user_data["chatgpt"] = ChatGPT(chat_mode="assistant")
+    if "chat_context" not in context.user_data:
+        context.user_data["chat_context"] = []
+
+    if "chat_mode" not in context.user_data:
+        context.user_data["chat_mode"] = "assistant"
+
+    if "total_n_used_tokens" not in context.user_data:
+        context.user_data["total_n_used_tokens"] = 0
 
     if "last_interation_timestamp" not in context.user_data:
         context.user_data["last_interation_timestamp"] = time.time()