✨ Feature: Use the /reset command to set the system prompt.

🔌 Plugin: Add travel information plugin. 💻 Code: 1. Upgrade ModelMerge version to 0.4.0 2. Refactor and simplify the code for uploading images. 3. Optimize the plugin button creation process.
yym68686 · May 14, 2024 · 92127a8 · 92127a8
1 parent 6e0b4a8
commit 92127a8
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 89 deletions.
diff --git a/bot.py b/bot.py
@@ -8,7 +8,7 @@
 
 from ModelMerge.models.config import PLUGINS
 from ModelMerge.utils.prompt import translator_en2zh_prompt, translator_prompt, claude3_doc_assistant_prompt
-from ModelMerge.utils.scripts import Document_extract, get_encode_image, claude_replace
+from ModelMerge.utils.scripts import Document_extract, claude_replace
 
 import config
 from config import (
@@ -22,7 +22,8 @@
     update_ENGINE,
     reset_ENGINE,
     update_language,
-    get_robot
+    get_robot,
+    get_image_message
 )
 
 from utils.i18n import strings
@@ -129,16 +130,7 @@ async def command_bot(update, context, language=None, prompt=translator_prompt,
             robot, role = get_robot()
             if "gpt" in config.GPT_ENGINE or (config.CLAUDE_API and "claude-3" in config.GPT_ENGINE):
                 message = [{"type": "text", "text": message}]
-            if image_url and (config.GPT_ENGINE == "gpt-4-turbo-2024-04-09" or "gpt-4o" in config.GPT_ENGINE):
-                base64_image = get_encode_image(image_url)
-                message.append(
-                    {
-                        "type": "image_url",
-                        "image_url": {
-                            "url": base64_image
-                        }
-                    }
-                )
+            message = get_image_message(image_url, message)
             # print("robot", robot)
             await context.bot.send_chat_action(chat_id=chatid, action=ChatAction.TYPING)
             await getChatGPT(update, context, title, robot, message, chatid, messageid)
@@ -153,7 +145,10 @@ async def command_bot(update, context, language=None, prompt=translator_prompt,
 @decorators.GroupAuthorization
 @decorators.Authorization
 async def reset_chat(update, context):
-    reset_ENGINE(update.message.chat_id)
+    message = None
+    if (len(context.args) > 0):
+        message = ' '.join(context.args)
+    reset_ENGINE(update.message.chat_id, message)
 
     await context.bot.send_message(
         chat_id=update.message.chat_id,
@@ -187,25 +182,13 @@ async def getChatGPT(update, context, title, robot, message, chatid, messageid):
             tmpresult = result
             if re.sub(r"```", '', result.split("\n")[-1]).count("`") % 2 != 0:
                 tmpresult = result + "`"
-            # if re.sub(r"```", '', result).count("`") % 2 != 0:
-            #     tmpresult = result + "`"
             if sum([line.strip().startswith("```") for line in result.split('\n')]) % 2 != 0:
                 tmpresult = tmpresult + "\n```"
-            # if result.count("```") % 2 != 0:
-            #     tmpresult = tmpresult + "\n```"
             tmpresult = title + tmpresult
             if "claude" in title:
                 tmpresult = claude_replace(tmpresult)
             if "🌐" in data:
                 tmpresult = data
-            # if "answer:" in result:
-            #     tmpresult = re.sub(r"thought:[\S\s]+?answer:\s", '', tmpresult)
-            #     tmpresult = re.sub(r"action:[\S\s]+?answer:\s", '', tmpresult)
-            #     tmpresult = re.sub(r"answer:\s", '', tmpresult)
-            #     tmpresult = re.sub(r"thought:[\S\s]+", '', tmpresult)
-            #     tmpresult = re.sub(r"action:[\S\s]+", '', tmpresult)
-            # else:
-            #     tmpresult = re.sub(r"thought:[\S\s]+", '', tmpresult)
             modifytime = modifytime + 1
             if (modifytime % Frequency_Modification == 0 and lastresult != tmpresult) or "🌐" in data:
                 await context.bot.edit_message_text(chat_id=chatid, message_id=messageid, text=escape(tmpresult), parse_mode='MarkdownV2', disable_web_page_preview=True, read_timeout=time_out, write_timeout=time_out, pool_timeout=time_out, connect_timeout=time_out)
@@ -217,11 +200,6 @@ async def getChatGPT(update, context, title, robot, message, chatid, messageid):
         print('\033[0m')
         if config.API:
             robot.reset(convo_id=str(chatid), system_prompt=config.systemprompt)
-        if "You exceeded your current quota, please check your plan and billing details." in str(e):
-            print("OpenAI api 已过期！")
-            await context.bot.delete_message(chat_id=chatid, message_id=messageid)
-            messageid = ''
-            config.API = ''
         tmpresult = f"{tmpresult}\n\n`{e}`"
     print(tmpresult)
     if lastresult != tmpresult and messageid:
@@ -400,33 +378,9 @@ async def handle_photo(update, context):
     image_url = photo_file.file_path
 
     robot, role = get_robot()
+    message = get_image_message(image_url, [])
 
-    base64_image = get_encode_image(image_url)
-    if image_url and ("gpt-4" in config.GPT_ENGINE or (config.CLAUDE_API is None and "claude-3" in config.GPT_ENGINE)):
-        message = [
-            {
-                "type": "image_url",
-                "image_url": {
-                    "url": base64_image
-                }
-            }
-        ]
-    if image_url and config.CLAUDE_API and "claude-3" in config.GPT_ENGINE:
-        message = [
-            {
-                "type": "image",
-                "source": {
-                    "type": "base64",
-                    "media_type": "image/jpeg",
-                    "data": base64_image.split(",")[1],
-                }
-            }
-        ]
-
-    # print(message)
     robot.add_to_conversation(message, role, str(chatid))
-    # print(robot.conversation)
-    # print(robot.conversation[str(chatid)])
     # if config.CLAUDE_API and "claude-3" in config.GPT_ENGINE:
     #     robot.add_to_conversation(claude3_doc_assistant_prompt, "assistant", str(update.effective_chat.id))
     message = (
@@ -475,7 +429,6 @@ async def post_init(application: Application) -> None:
     await application.bot.set_my_commands([
         BotCommand('info', 'basic information'),
         BotCommand('pic', 'Generate image'),
-        # BotCommand('copilot', 'Advanced search mode'),
         BotCommand('search', 'search Google or duckduckgo'),
         BotCommand('en2zh', 'translate to Chinese'),
         BotCommand('zh2en', 'translate to English'),
@@ -521,15 +474,12 @@ async def process_update(update):
     application.add_handler(CommandHandler("start", start))
     application.add_handler(CommandHandler("pic", image, block = False))
     application.add_handler(CommandHandler("search", lambda update, context: command_bot(update, context, prompt="search: ", title=f"`🤖️ {config.GPT_ENGINE}`\n\n", robot=config.ChatGPTbot, has_command="search")))
-    # application.add_handler(CommandHandler("search", lambda update, context: search(update, context, title=f"`🤖️ {config.GPT_ENGINE}`\n\n", robot=config.ChatGPTbot)))
     application.add_handler(CallbackQueryHandler(button_press))
     application.add_handler(CommandHandler("reset", reset_chat))
     application.add_handler(CommandHandler("en2zh", lambda update, context: command_bot(update, context, "Simplified Chinese", robot=config.translate_bot)))
     application.add_handler(CommandHandler("zh2en", lambda update, context: command_bot(update, context, "english", robot=config.translate_bot)))
-    # application.add_handler(CommandHandler("copilot", lambda update, context: command_bot(update, context, None, None, title=f"`🤖️ {config.GPT_ENGINE}`\n\n", robot=config.copilot_bot)))
     application.add_handler(CommandHandler("info", info))
     application.add_handler(InlineQueryHandler(inlinequery))
-    # application.add_handler(CommandHandler("qa", qa))
     application.add_handler(MessageHandler(filters.Document.PDF | filters.Document.TXT | filters.Document.DOC, handle_pdf))
     application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, lambda update, context: command_bot(update, context, prompt=None, title=f"`🤖️ {config.GPT_ENGINE}`\n\n", robot=config.ChatGPTbot, has_command=False)))
     application.add_handler(MessageHandler(filters.CAPTION & filters.PHOTO & ~filters.COMMAND, lambda update, context: command_bot(update, context, prompt=None, title=f"`🤖️ {config.GPT_ENGINE}`\n\n", robot=config.ChatGPTbot, has_command=False)))

diff --git a/config.py b/config.py
@@ -5,6 +5,7 @@
 from utils.i18n import strings
 from datetime import datetime
 from ModelMerge.utils import prompt
+from ModelMerge.utils.scripts import get_encode_image
 from ModelMerge.models import chatgpt, claude, groq, claude3, gemini, dalle3
 from ModelMerge.models.config import PLUGINS
 from telegram import InlineKeyboardButton
@@ -82,8 +83,11 @@ def update_ENGINE(data = None):
         gemini_Bot = gemini(api_key=f"{GOOGLE_AI_API_KEY}", engine=GPT_ENGINE, system_prompt=systemprompt, temperature=temperature)
 update_ENGINE()
 
-def reset_ENGINE(chat_id):
-    global ChatGPTbot, translate_bot, claudeBot, claude3Bot, groqBot, gemini_Bot
+def reset_ENGINE(chat_id, message=None):
+    global ChatGPTbot, translate_bot, claudeBot, claude3Bot, groqBot, gemini_Bot, systemprompt, claude_systemprompt
+    if message:
+        systemprompt = message
+        claude_systemprompt = message
     if API and ChatGPTbot:
         ChatGPTbot.reset(convo_id=str(chat_id), system_prompt=systemprompt)
     if CLAUDE_API and claudeBot:
@@ -114,6 +118,30 @@ def get_robot():
 
     return robot, role
 
+def get_image_message(image_url, message):
+    if image_url:
+        base64_image = get_encode_image(image_url)
+        if "gpt-4" in GPT_ENGINE or (CLAUDE_API is None and "claude-3" in GPT_ENGINE):
+            message.append(
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": base64_image
+                    }
+                }
+            )
+        if CLAUDE_API and "claude-3" in GPT_ENGINE:
+            message.append(
+                {
+                    "type": "image",
+                    "source": {
+                        "type": "base64",
+                        "media_type": "image/jpeg",
+                        "data": base64_image.split(",")[1],
+                    }
+                }
+            )
+    return message
 
 whitelist = os.environ.get('whitelist', None)
 if whitelist:
@@ -163,7 +191,7 @@ def delete_model_digit_tail(lst):
             else:
                 return "-".join(lst[:i + 1])
 
-def create_buttons(strings):
+def create_buttons(strings, plugins_status=False, lang="English", button_text=None):
     # 过滤出长度小于15的字符串
     filtered_strings1 = [s for s in strings if len(delete_model_digit_tail(s.split("-"))) <= 14]
     filtered_strings2 = [s for s in strings if len(delete_model_digit_tail(s.split("-"))) > 14]
@@ -172,7 +200,10 @@ def create_buttons(strings):
     temp = []
 
     for string in filtered_strings1:
-        button = InlineKeyboardButton(delete_model_digit_tail(string.split("-")), callback_data=string + "ENGINE")
+        if plugins_status:
+            button = InlineKeyboardButton(f"{get_plugins_status(string)}{button_text[string][lang]}", callback_data=string)
+        else:
+            button = InlineKeyboardButton(delete_model_digit_tail(string.split("-")), callback_data=string + "ENGINE")
         temp.append(button)
 
         # 每两个按钮一组
@@ -185,7 +216,10 @@ def create_buttons(strings):
         buttons.append(temp)
 
     for string in filtered_strings2:
-        button = InlineKeyboardButton(delete_model_digit_tail(string.split("-")), callback_data=string + "ENGINE")
+        if plugins_status:
+            button = InlineKeyboardButton(f"{get_plugins_status(string)}{button_text[string][lang]}", callback_data=string)
+        else:
+            button = InlineKeyboardButton(delete_model_digit_tail(string.split("-")), callback_data=string + "ENGINE")
         buttons.append([button])
 
     return buttons
@@ -255,13 +289,16 @@ def update_first_buttons_message():
             InlineKeyboardButton(strings['button_language'][lang], callback_data="language"),
             InlineKeyboardButton(f"{history} {strings['button_history'][lang]}", callback_data="PASS_HISTORY"),
         ],
-        [
-            InlineKeyboardButton(f"{get_plugins_status('SEARCH')}{strings['button_search'][lang]}", callback_data='SEARCH'),
-            InlineKeyboardButton(f"{get_plugins_status('DATE')}{strings['button_current_time'][lang]}", callback_data='DATE'),
-        ],
-        [
-            InlineKeyboardButton(f"{get_plugins_status('URL')}{strings['button_url'][lang]}", callback_data='URL'),
-            InlineKeyboardButton(f"{get_plugins_status('VERSION')}{strings['button_version'][lang]}", callback_data='VERSION'),
-        ],
     ]
+    PLUGINS_LIST = list(PLUGINS.keys())
+    buttons = create_buttons(PLUGINS_LIST, plugins_status=True, lang=lang, button_text=strings)
+    first_buttons.extend(buttons)
+        # [
+        #     InlineKeyboardButton(f"{get_plugins_status('SEARCH')}{strings['SEARCH'][lang]}", callback_data='SEARCH'),
+        #     InlineKeyboardButton(f"{get_plugins_status('DATE')}{strings['DATE'][lang]}", callback_data='DATE'),
+        # ],
+        # [
+        #     InlineKeyboardButton(f"{get_plugins_status('URL')}{strings['URL'][lang]}", callback_data='URL'),
+        #     InlineKeyboardButton(f"{get_plugins_status('VERSION')}{strings['VERSION'][lang]}", callback_data='VERSION'),
+        # ],
     return first_buttons
diff --git a/requirements.txt b/requirements.txt
@@ -3,6 +3,6 @@ pytz
 python-dotenv
 md2tgmd==0.1.9
 fake_useragent
-ModelMerge==0.3.11
+ModelMerge==0.4.0
 oauth2client==3.0.0
 python-telegram-bot[webhooks,rate-limiter]==21.0.1
diff --git a/utils/i18n.py b/utils/i18n.py
@@ -7,26 +7,10 @@
         "zh": "切换模型",
         "en": "Change model",
     },
-    "button_search": {
-        "zh": "搜索",
-        "en": "Search",
-    },
-    "button_url": {
-        "zh": "网址总结",
-        "en": "URL summary",
-    },
     "button_history": {
         "zh": "历史记录",
         "en": "Chat history",
     },
-    "button_current_time": {
-        "zh": "当前时间",
-        "en": "Current time",
-    },
-    "button_version": {
-        "zh": "版本信息",
-        "en": "Version info",
-    },
     "button_back": {
         "zh": "⬅️ 返回",
         "en": "⬅️ Back",
@@ -43,4 +27,24 @@
         "zh": "👇 下面可以随时更改默认模型：",
         "en": "👇 Change model below:",
     },
+    "SEARCH": {
+        "zh": "搜索",
+        "en": "Search",
+    },
+    "URL": {
+        "zh": "网址总结",
+        "en": "URL summary",
+    },
+    "DATE": {
+        "zh": "当前时间",
+        "en": "Current time",
+    },
+    "VERSION": {
+        "zh": "版本信息",
+        "en": "Version info",
+    },
+    "TARVEL": {
+        "zh": "旅行信息",
+        "en": "Travel info",
+    },
 }