Add different chat modes

README.md

@@ -9,9 +9,9 @@ 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*: Different chat modes (code writing helper, therapist, movie expert, etc.)
 
 ## Setup
 1. Get your [OpenAI API](https://openai.com/api/) key

bot.py

@@ -5,17 +5,19 @@ import html
 import json
 import time
 
-from telegram import Update
+from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
 from telegram.ext import (
     ApplicationBuilder,
     CallbackContext,
     CommandHandler,
     MessageHandler,
     PicklePersistence,
+    CallbackQueryHandler,
     filters
 )
 from telegram.constants import ParseMode, ChatAction
 
+import chatgpt
 import utils
 import config
 
@@ -89,6 +91,39 @@ async def reset_handle(update: Update, context: CallbackContext):
     context.user_data["chatgpt"].reset_context()
     await update.message.reply_text("Chat context is reset ✅")
 
+    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)
+
+
+async def show_chat_modes_handle(update: Update, context: CallbackContext):
+    utils.init_user(update, context)
+    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}")])
+    reply_markup = InlineKeyboardMarkup(keyboard)
+
+    await update.message.reply_text("Choose 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]
+
+    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",
+        parse_mode=ParseMode.HTML
+    )
+
+    await query.edit_message_text(f"{chatgpt.CHAT_MODES[chat_mode_key]['welcome_message']}", parse_mode=ParseMode.HTML)
+
 
 async def error_handler(update: Update, context: CallbackContext) -> None:
     logger.error(msg="Exception while handling an update:", exc_info=context.error)
@@ -125,9 +160,13 @@ def run_bot() -> None:
 
     application.add_handler(CommandHandler("start", start_handle, filters=user_filter))
     application.add_handler(CommandHandler("help", help_handle, filters=user_filter))
+
     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("mode", show_chat_modes_handle, filters=user_filter))
+    application.add_handler(CallbackQueryHandler(set_chat_mode_handle, pattern="^set_mode"))
     
     application.add_error_handler(error_handler)
     

chatgpt.py

@@ -1,27 +1,36 @@
+import config
+
 import openai
+openai.api_key = config.openai_api_key
 
 
-PROMPT_START = """As an advanced chatbot named ChatGPT, your primary goal is to assist users to the best of your ability.
-This may involve answering questions, providing helpful information, or completing tasks based on user input.
-In order to effectively assist users, it is important to be detailed and thorough in your responses.
-Use examples and evidence to support your points and justify your recommendations or solutions.
-Remember to always prioritize the needs and satisfaction of the user.
-Your ultimate goal is to provide a helpful and enjoyable experience for the user.
-If you have code in your reply, write it inside the <code>, </code> tags
-"""
+CHAT_MODES = {
+    "assistant": {
+        "name": "👩🏼‍🎓 Assistant",
+        "welcome_message": "👩🏼‍🎓 Hi, I'm <b>ChatGPT assistant</b>. How can I help you?",
+        "prompt_start": "As an advanced chatbot named ChatGPT, your primary goal is to assist users to the best of your ability. This may involve answering questions, providing helpful information, or completing tasks based on user input. In order to effectively assist users, it is important to be detailed and thorough in your responses. Use examples and evidence to support your points and justify your recommendations or solutions. Remember to always prioritize the needs and satisfaction of the user. Your ultimate goal is to provide a helpful and enjoyable experience for the user. If you have code in your reply, write it inside <code>, </code> tags."
+    },
 
-# TODO: add different chat modes: code writing helper, therapist, expert in X, etc.
+    "code_assistant": {
+        "name": "👩🏼‍💻 Code Assistant",
+        "welcome_message": "👩🏼‍💻 Hi, I'm <b>ChatGPT code assistant</b>. How can I help you?",
+        "prompt_start": "As an advanced chatbot named ChatGPT, your primary goal is to assist users to write code. This may involve designing/writing/editing/describing code or providing helpful information. Where possible you should provide code examples to support your points and justify your recommendations or solutions. Make sure the code you provide is correct and can be run without errors. Be detailed and thorough in your responses. Your ultimate goal is to provide a helpful and enjoyable experience for the user. Write code inside <code>, </code> tags."
+    },
+
+    "movie_expert": {
+        "name": "🎬 Movie Expert",
+        "welcome_message": "🎬 Hi, I'm <b>ChatGPT movie expert</b>. How can I help you?",
+        "prompt_start": "As an advanced movie expert chatbot named ChatGPT, your primary goal is to assist users to the best of your ability. You can answer questions about movies, actors, directors, and more. You can recommend movies to users based on their preferences. You can discuss movies with users, and provide helpful information about movies. In order to effectively assist users, it is important to be detailed and thorough in your responses. Use examples and evidence to support your points and justify your recommendations or solutions. Remember to always prioritize the needs and satisfaction of the user. Your ultimate goal is to provide a helpful and enjoyable experience for the user."
+    },
+}
 
 
 class ChatGPT:
-    def __init__(self, openai_api_key=""):
-        self.openai_api_key = openai_api_key
+    def __init__(self, chat_mode="assistant"):
+        self.chat_mode = chat_mode
         self.context = []
     
     def send_message(self, message):
-        openai.api_key = self.openai_api_key
-
-        # get answer
         prompt = self._generate_prompt(message)
         r = openai.Completion.create(
             engine="text-davinci-003",
@@ -43,9 +52,15 @@ class ChatGPT:
     def reset_context(self):
         self.context = []
 
+    def set_chat_mode(self, chat_mode):
+        if chat_mode not in CHAT_MODES.keys():
+            raise ValueError(f"Chat mode {chat_mode} is not supported")
+
+        self.chat_mode = chat_mode
+
     def _generate_prompt(self, message):
-        prompt = PROMPT_START
-        prompt += "\n"
+        prompt = CHAT_MODES[self.chat_mode]["prompt_start"]
+        prompt += "\n\n"
 
         # chat history
         if len(self.context) > 0:

utils.py

@@ -10,7 +10,7 @@ import config
 def init_user(update: Update, context: CallbackContext):
     # init chatgpt
     if "chatgpt" not in context.user_data:
-        context.user_data["chatgpt"] = ChatGPT(openai_api_key=config.openai_api_key)
+        context.user_data["chatgpt"] = ChatGPT(chat_mode="assistant")
 
     if "last_interation_timestamp" not in context.user_data:
         context.user_data["last_interation_timestamp"] = time.time()