Add support of ChatGPT API (#70)

Co-authored-by: Karim Iskakov <[email protected]>
father-bot · Mar 2, 2023 · 0779808 · 0779808
1 parent 7264f7c
commit 0779808
Show file tree

Hide file tree

Showing 6 changed files with 54 additions and 21 deletions.
diff --git a/README.md b/README.md
@@ -14,11 +14,15 @@ This repo is ChatGPT re-created with GPT-3.5 LLM as Telegram Bot. **And it works
 
 You can deploy your own bot, or use mine: [@chatgpt_karfly_bot](https://t.me/chatgpt_karfly_bot)
 
+## News
+- *2 Mar 2023*: Added support of [ChatGPT API](https://platform.openai.com/docs/guides/chat/introduction). It's enabled by default and can be disabled with `use_chatgpt_api` option in config. Don't forget to **rebuild** you docker image (`--build`).
+
 ## Features
 - Low latency replies (it usually takes about 3-5 seconds) 
 - No request limits
 - Code highlighting
 - Special chat modes: 👩🏼‍🎓 Assistant, 👩🏼‍💻 Code Assistant, 🎬 Movie Expert. More soon
+- Support of [ChatGPT API](https://platform.openai.com/docs/guides/chat/introduction)
 - List of allowed Telegram users
 - Track $ balance spent on OpenAI API
 

diff --git a/bot/bot.py b/bot/bot.py
@@ -105,7 +105,8 @@ async def message_handle(update: Update, context: CallbackContext, message=None,
     try:
         message = message or update.message.text
 
-        answer, prompt, n_used_tokens, n_first_dialog_messages_removed = chatgpt.ChatGPT().send_message(
+        chatgpt_instance = chatgpt.ChatGPT(use_chatgpt_api=config.use_chatgpt_api)
+        answer, n_used_tokens, n_first_dialog_messages_removed = chatgpt_instance.send_message(
             message,
             dialog_messages=db.get_dialog_messages(user_id, dialog_id=None),
             chat_mode=db.get_user_attribute(user_id, "current_chat_mode"),
@@ -194,10 +195,12 @@ async def show_balance_handle(update: Update, context: CallbackContext):
     db.set_user_attribute(user_id, "last_interaction", datetime.now())
 
     n_used_tokens = db.get_user_attribute(user_id, "n_used_tokens")
-    n_spent_dollars = n_used_tokens * (0.02 / 1000)
+
+    price = 0.002 if config.use_chatgpt_api else 0.02
+    n_spent_dollars = n_used_tokens * (price / 1000)
 
     text = f"You spent <b>{n_spent_dollars:.03f}$</b>\n"
-    text += f"You used <b>{n_used_tokens}</b> tokens <i>(price: 0.02$ per 1000 tokens)</i>\n"
+    text += f"You used <b>{n_used_tokens}</b> tokens <i>(price: {price}$ per 1000 tokens)</i>\n"
 
     await update.message.reply_text(text, parse_mode=ParseMode.HTML)
 

diff --git a/bot/chatgpt.py b/bot/chatgpt.py
@@ -30,10 +30,18 @@
     },
 }
 
+OPENAI_COMPLETION_OPTIONS = {
+    "temperature": 0.7,
+    "max_tokens": 1000,
+    "top_p": 1,
+    "frequency_penalty": 0,
+    "presence_penalty": 0
+}
+
 
 class ChatGPT:
-    def __init__(self):
-        pass
+    def __init__(self, use_chatgpt_api=True):
+        self.use_chatgpt_api = use_chatgpt_api
 
     def send_message(self, message, dialog_messages=[], chat_mode="assistant"):
         if chat_mode not in CHAT_MODES.keys():
@@ -42,22 +50,27 @@ def send_message(self, message, dialog_messages=[], chat_mode="assistant"):
         n_dialog_messages_before = len(dialog_messages)
         answer = None
         while answer is None:
-            prompt = self._generate_prompt(message, dialog_messages, 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 = self._postprocess_answer(answer)
+                if self.use_chatgpt_api:
+                    messages = self._generate_prompt_messages_for_chatgpt_api(message, dialog_messages, chat_mode)
+                    r = openai.ChatCompletion.create(
+                        model="gpt-3.5-turbo",
+                        messages=messages,
+                        **OPENAI_COMPLETION_OPTIONS
+                    )
+                    answer = r.choices[0].message["content"]
+                else:
+                    prompt = self._generate_prompt(message, dialog_messages, chat_mode)
+                    r = openai.Completion.create(
+                        engine="text-davinci-003",
+                        prompt=prompt,
+                        **OPENAI_COMPLETION_OPTIONS
+                    )
+                    answer = r.choices[0].text
 
+                answer = self._postprocess_answer(answer)
                 n_used_tokens = r.usage.total_tokens
-
+                
             except openai.error.InvalidRequestError as e:  # too many tokens
                 if len(dialog_messages) == 0:
                     raise ValueError("Dialog messages is reduced to zero, but still has too many tokens to make completion") from e
@@ -67,7 +80,7 @@ def send_message(self, message, dialog_messages=[], chat_mode="assistant"):
 
         n_first_dialog_messages_removed = n_dialog_messages_before - len(dialog_messages)
 
-        return answer, prompt, n_used_tokens, n_first_dialog_messages_removed
+        return answer, n_used_tokens, n_first_dialog_messages_removed
 
     def _generate_prompt(self, message, dialog_messages, chat_mode):
         prompt = CHAT_MODES[chat_mode]["prompt_start"]
@@ -86,6 +99,17 @@ def _generate_prompt(self, message, dialog_messages, chat_mode):
 
         return prompt
 
+    def _generate_prompt_messages_for_chatgpt_api(self, message, dialog_messages, chat_mode):
+        prompt = CHAT_MODES[chat_mode]["prompt_start"]
+
+        messages = [{"role": "system", "content": prompt}]
+        for dialog_message in dialog_messages:
+            messages.append({"role": "user", "content": dialog_message["user"]})
+            messages.append({"role": "assistant", "content": dialog_message["bot"]})
+        messages.append({"role": "user", "content": message})
+
+        return messages
+
     def _postprocess_answer(self, answer):
         answer = answer.strip()
         return answer
diff --git a/bot/config.py b/bot/config.py
@@ -14,6 +14,7 @@
 # config parameters
 telegram_token = config_yaml["telegram_token"]
 openai_api_key = config_yaml["openai_api_key"]
+use_chatgpt_api = config_yaml.get("use_chatgpt_api", True)
 allowed_telegram_usernames = config_yaml["allowed_telegram_usernames"]
 new_dialog_timeout = config_yaml["new_dialog_timeout"]
 mongodb_uri = f"mongodb://mongo:{config_env['MONGODB_PORT']}"
diff --git a/config/config.example.yml b/config/config.example.yml
@@ -1,4 +1,5 @@
 telegram_token: ""
 openai_api_key: ""
+use_chatgpt_api: true
 allowed_telegram_usernames: []  # if empty, the bot is available to anyone
 new_dialog_timeout: 600  # new dialog starts after timeout (in seconds)
diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,5 @@
-python-telegram-bot==20.0a0
-openai>=0.26.1
+python-telegram-bot==20.1
+openai>=0.27.0
 PyYAML==6.0
 pymongo==4.3.3
 python-dotenv==0.21.0