From 14fed8142c47a9286e41c0558f688447b81207b9 Mon Sep 17 00:00:00 2001 From: yym68686 Date: Tue, 11 Jun 2024 19:24:08 +0800 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Feature:=20Add=20the=20function=20o?= =?UTF-8?q?f=20uploading=20original=20image=20Q&A.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 💻 Code: Upgrade ModelMerge version to 0.6.6 --- bot.py | 92 ++++++++++++++++++++---------------------------- config.py | 26 -------------- requirements.txt | 2 +- 3 files changed, 40 insertions(+), 80 deletions(-) diff --git a/bot.py b/bot.py index c1fdfc52..51179a38 100644 --- a/bot.py +++ b/bot.py @@ -7,7 +7,7 @@ from md2tgmd import escape from ModelMerge.utils.prompt import translator_en2zh_prompt, translator_prompt, claude3_doc_assistant_prompt -from ModelMerge.utils.scripts import Document_extract, claude_replace +from ModelMerge.utils.scripts import Document_extract, claude_replace, get_image_message import config from config import ( @@ -24,7 +24,6 @@ update_ENGINE, reset_ENGINE, get_robot, - get_image_message, get_ENGINE, update_language_status, update_models_buttons, @@ -89,16 +88,26 @@ def CutNICK(update_text, update_message): else: return None +async def get_file_url(file, context): + file_id = file.file_id + new_file = await context.bot.get_file(file_id) + file_url = new_file.file_path + return file_url + async def GetMesage(update_message, context): image_url = None + file_url = None reply_to_message_text = None + chatid = update_message.chat_id message_thread_id = update_message.message_thread_id if message_thread_id: convo_id = str(chatid) + "_" + str(message_thread_id) else: convo_id = chatid + messageid = update_message.message_id + if update_message.text: message = CutNICK(update_message.text, update_message) rawtext = update_message.text @@ -108,24 +117,31 @@ async def GetMesage(update_message, context): if update_message.photo: photo = update_message.photo[-1] - file_id = photo.file_id - photo_file = await context.bot.getFile(file_id) - image_url = photo_file.file_path + + image_url = await get_file_url(photo, context) message = rawtext = CutNICK(update_message.caption, update_message) - return message, rawtext, image_url, chatid, messageid, reply_to_message_text, message_thread_id, convo_id + + if update_message.document: + file = update_message.document + + file_url = await get_file_url(file, context) + + message = rawtext = CutNICK(update_message.caption, update_message) + + return message, rawtext, image_url, chatid, messageid, reply_to_message_text, message_thread_id, convo_id, file_url async def GetMesageInfo(update, context): if update.edited_message: - message, rawtext, image_url, chatid, messageid, reply_to_message_text, message_thread_id, convo_id = await GetMesage(update.edited_message, context) + message, rawtext, image_url, chatid, messageid, reply_to_message_text, message_thread_id, convo_id, file_url = await GetMesage(update.edited_message, context) update_message = update.edited_message elif update.callback_query: - message, rawtext, image_url, chatid, messageid, reply_to_message_text, message_thread_id, convo_id = await GetMesage(update.callback_query.message, context) + message, rawtext, image_url, chatid, messageid, reply_to_message_text, message_thread_id, convo_id, file_url = await GetMesage(update.callback_query.message, context) update_message = update.callback_query.message else: - message, rawtext, image_url, chatid, messageid, reply_to_message_text, message_thread_id, convo_id = await GetMesage(update.message, context) + message, rawtext, image_url, chatid, messageid, reply_to_message_text, message_thread_id, convo_id, file_url = await GetMesage(update.message, context) update_message = update.message - return message, rawtext, image_url, chatid, messageid, reply_to_message_text, update_message, message_thread_id, convo_id + return message, rawtext, image_url, chatid, messageid, reply_to_message_text, update_message, message_thread_id, convo_id, file_url # 定义一个缓存来存储消息 message_cache = defaultdict(lambda: []) @@ -136,7 +152,7 @@ async def GetMesageInfo(update, context): async def command_bot(update, context, language=None, prompt=translator_prompt, title="", robot=None, has_command=True): stop_event.clear() print("update", update) - message, rawtext, image_url, chatid, messageid, reply_to_message_text, update_message, message_thread_id, convo_id = await GetMesageInfo(update, context) + message, rawtext, image_url, chatid, messageid, reply_to_message_text, update_message, message_thread_id, convo_id, file_url = await GetMesageInfo(update, context) print("\033[32m", update.effective_user.username, update.effective_user.id, rawtext, "\033[0m") if has_command == False or len(context.args) > 0: @@ -188,7 +204,7 @@ async def command_bot(update, context, language=None, prompt=translator_prompt, if "gpt" in engine or (config.CLAUDE_API and "claude-3" in engine): message = [{"type": "text", "text": message}] - message = get_image_message(image_url, message, convo_id) + message = get_image_message(image_url, message, engine) await context.bot.send_chat_action(chat_id=chatid, message_thread_id=message_thread_id, action=ChatAction.TYPING) if Users.get_config(convo_id, "TITLE"): title = f"`🤖️ {engine}`\n\n" @@ -219,7 +235,7 @@ async def delete_message(update, context, messageid, delay=60): @decorators.Authorization async def reset_chat(update, context): global target_convo_id - _, _, _, chatid, _, _, _, message_thread_id, convo_id = await GetMesageInfo(update, context) + _, _, _, chatid, _, _, _, message_thread_id, convo_id, file_url = await GetMesageInfo(update, context) target_convo_id = convo_id print("target_chatid", convo_id) stop_event.set() @@ -322,7 +338,7 @@ async def getChatGPT(update, context, title, robot, message, chatid, messageid, @decorators.Authorization async def button_press(update, context): """Function to handle the button press""" - message, rawtext, image_url, chatid, messageid, reply_to_message_text, update_message, message_thread_id, convo_id = await GetMesageInfo(update, context) + message, rawtext, image_url, chatid, messageid, reply_to_message_text, update_message, message_thread_id, convo_id, file_url = await GetMesageInfo(update, context) callback_query = update.callback_query print("callback_query chatid", chatid) info_message = update_info_message(convo_id) @@ -429,55 +445,25 @@ async def button_press(update, context): @decorators.GroupAuthorization @decorators.Authorization async def info(update, context): - _, _, _, chatid, _, _, _, message_thread_id, convo_id = await GetMesageInfo(update, context) + _, _, _, chatid, _, _, _, message_thread_id, convo_id, file_url = await GetMesageInfo(update, context) info_message = update_info_message(convo_id) message = await context.bot.send_message(chat_id=chatid, message_thread_id=message_thread_id, text=escape(info_message), reply_markup=InlineKeyboardMarkup(update_first_buttons_message(convo_id)), parse_mode='MarkdownV2', disable_web_page_preview=True) await delete_message(update, context, message.message_id) @decorators.GroupAuthorization @decorators.Authorization -async def handle_pdf(update, context): - # 获取接收到的文件 - pdf_file = update.message.document - # 得到文件的url - file_id = pdf_file.file_id - new_file = await context.bot.get_file(file_id) - file_url = new_file.file_path - extracted_text_with_prompt = Document_extract(file_url) +async def handle_file(update, context): robot, role = get_robot() + _, _, image_url, chatid, _, messageid, update_message, message_thread_id, convo_id, file_url = await GetMesageInfo(update, context) + engine = get_ENGINE(chatid) - _, _, _, chatid, _, _, _, message_thread_id, convo_id = await GetMesageInfo(update, context) - robot.add_to_conversation(extracted_text_with_prompt, role, str(convo_id)) - engine = get_ENGINE(convo_id) - if config.CLAUDE_API and "claude-3" in engine: - robot.add_to_conversation(claude3_doc_assistant_prompt, "assistant", str(convo_id)) - message = ( - f"文档上传成功!\n\n" - ) - message = await context.bot.send_message(chat_id=chatid, message_thread_id=message_thread_id, text=escape(message), parse_mode='MarkdownV2', disable_web_page_preview=True) - await delete_message(update, context, message.message_id) - - -@decorators.GroupAuthorization -@decorators.Authorization -async def handle_photo(update, context): - _, _, _, chatid, _, messageid, update_message, message_thread_id, convo_id = await GetMesageInfo(update, context) - - photo = update_message.photo[-1] - file_id = photo.file_id - photo_file = await context.bot.getFile(file_id) - image_url = photo_file.file_path - - robot, role = get_robot() - message = get_image_message(image_url, [], convo_id) + message = Document_extract(file_url, None, engine) robot.add_to_conversation(message, role, str(convo_id)) - # 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 = ( - f"图片上传成功!\n\n" + f"上传成功!\n\n" ) - message = await context.bot.send_message(chat_id=chatid, text=escape(message), parse_mode='MarkdownV2', disable_web_page_preview=True) + message = await context.bot.send_message(chat_id=chatid, message_thread_id=message_thread_id, text=escape(message), parse_mode='MarkdownV2', disable_web_page_preview=True) await delete_message(update, context, message.message_id) # DEBOUNCE_TIME = 4 @@ -596,10 +582,10 @@ async def process_update(update): application.add_handler(CommandHandler("zh2en", lambda update, context: command_bot(update, context, "english", robot=config.translate_bot))) application.add_handler(CommandHandler("info", info)) application.add_handler(InlineQueryHandler(inlinequery)) - application.add_handler(MessageHandler(filters.Document.PDF | filters.Document.TXT | filters.Document.DOC, handle_pdf)) + application.add_handler(MessageHandler(filters.Document.PDF | filters.Document.TXT | filters.Document.DOC | filters.Document.FileExtension("jpg") | filters.Document.FileExtension("jpeg"), handle_file)) application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, lambda update, context: command_bot(update, context, prompt=None, has_command=False), block = False)) application.add_handler(MessageHandler(filters.CAPTION & filters.PHOTO & ~filters.COMMAND, lambda update, context: command_bot(update, context, prompt=None, has_command=False))) - application.add_handler(MessageHandler(~filters.CAPTION & filters.PHOTO & ~filters.COMMAND, handle_photo)) + application.add_handler(MessageHandler(~filters.CAPTION & filters.PHOTO & ~filters.COMMAND, handle_file)) application.add_handler(MessageHandler(filters.COMMAND, unknown)) application.add_error_handler(error) diff --git a/config.py b/config.py index 27ec8652..2ed411ae 100644 --- a/config.py +++ b/config.py @@ -221,32 +221,6 @@ def get_robot(chat_id = None): return robot, role -def get_image_message(image_url, message, chatid = None): - engine = get_ENGINE(chatid) - if image_url: - base64_image = get_encode_image(image_url) - if "gpt-4" in engine or (CLAUDE_API is None and "claude-3" in engine): - message.append( - { - "type": "image_url", - "image_url": { - "url": base64_image - } - } - ) - if CLAUDE_API and "claude-3" in 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: whitelist = [int(id) for id in whitelist.split(",")] diff --git a/requirements.txt b/requirements.txt index 51e114da..62e75a5f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ --index-url https://pypi.python.org/simple/ python-dotenv md2tgmd==0.2.2 -ModelMerge==0.6.3 +ModelMerge==0.6.6 python-telegram-bot[webhooks,rate-limiter]==21.0.1 \ No newline at end of file