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()