diff --git a/README.md b/README.md index dd24d87..84c8a56 100644 --- a/README.md +++ b/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 diff --git a/bot.py b/bot.py index 2cd38b2..a186047 100644 --- a/bot.py +++ b/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 = "✍️ Note: Your current dialog is too long, so your first message was removed from the context.\n Send /new command to start new dialog" + else: + text = f"✍️ Note: Your current dialog is too long, so {n_first_chat_context_messages_removed} first messages 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"{chatgpt.CHAT_MODES[chat_mode_key]['name']} chat mode is set", + f"{chatgpt.CHAT_MODES[chat_mode]['name']} 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 {total_spent_dollars:.03f}$\n" + text += f"You used {total_n_used_tokens} tokens (price: 0.01$ per 1000 tokens)\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) diff --git a/chatgpt.py b/chatgpt.py index 4534b19..e542cea 100644 --- a/chatgpt.py +++ b/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 \ No newline at end of file + return prompt diff --git a/config.example.yml b/config.example.yml index edd8764..744f42c 100644 --- a/config.example.yml +++ b/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) \ No newline at end of file +new_dialog_timeout: 600 # new dialog starts after timeout (in seconds) \ No newline at end of file diff --git a/config.py b/config.py index 744138d..9c6ecdd 100644 --- a/config.py +++ b/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"] diff --git a/utils.py b/utils.py index 05456b7..3090f33 100644 --- a/utils.py +++ b/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()