diff --git a/handlers/__init__.py b/handlers/__init__.py index ed6801a..0892ade 100644 --- a/handlers/__init__.py +++ b/handlers/__init__.py @@ -5,6 +5,7 @@ import importlib import re import traceback from functools import update_wrapper +from mimetypes import guess_type from pathlib import Path from typing import Any, Callable, TypeVar @@ -114,6 +115,29 @@ def extract_prompt(message: str, bot_name: str) -> str: return message.strip() +def remove_prompt_prefix(message: str) -> str: + """ + Remove "/cmd" or "/cmd@bot_name" or "cmd:" + """ + message += " " + # Explanation of the regex pattern: + # ^ - Match the start of the string + # ( - Start of the group + # / - Literal forward slash + # [a-zA-Z] - Any letter (start of the command) + # [a-zA-Z0-9_]* - Any number of letters, digits, or underscores + # (@\w+)? - Optionally match @ followed by one or more word characters (for bot name) + # \s - A single whitespace character (space or newline) + # | - OR + # [a-zA-Z] - Any letter (start of the command) + # [a-zA-Z0-9_]* - Any number of letters, digits, or underscores + # :\s - Colon followed by a single whitespace character + # ) - End of the group + pattern = r"^(/[a-zA-Z][a-zA-Z0-9_]*(@\w+)?\s|[a-zA-Z][a-zA-Z0-9_]*:\s)" + + return re.sub(pattern, "", message).strip() + + def wrap_handler(handler: T, bot: TeleBot) -> T: def wrapper(message: Message, *args: Any, **kwargs: Any) -> None: try: @@ -219,9 +243,10 @@ def enrich_text_with_urls(text: str) -> str: def image_to_data_uri(file_path): + content_type = guess_type(file_path)[0] with open(file_path, "rb") as image_file: encoded_image = base64.b64encode(image_file.read()).decode("utf-8") - return f"data:image/png;base64,{encoded_image}" + return f"data:{content_type};base64,{encoded_image}" import json @@ -452,11 +477,31 @@ class TelegraphAPI: return new_dom + def upload_image(self, file_name: str) -> str: + base_url = "https://telegra.ph" + upload_url = f"{base_url}/upload" + + try: + content_type = guess_type(file_name)[0] + with open(file_name, "rb") as f: + response = requests.post( + upload_url, files={"file": ("blob", f, content_type)} + ) + response.raise_for_status() + # [{'src': '/file/xx.jpg'}] + response = response.json() + image_url = f"{base_url}{response[0]['src']}" + return image_url + except Exception as e: + print(f"upload image: {e}") + return "https://telegra.ph/api" + # `import *` will give you these __all__ = [ "bot_reply_first", "bot_reply_markdown", + "remove_prompt_prefix", "enrich_text_with_urls", "image_to_data_uri", "TelegraphAPI", diff --git a/handlers/useful.py b/handlers/useful.py index 7cbeb37..22cb4d7 100644 --- a/handlers/useful.py +++ b/handlers/useful.py @@ -192,8 +192,12 @@ def latest_handle_messages(message: Message, bot: TeleBot): """ignore""" chat_id = message.chat.id chat_user_id = message.from_user.id + # if not text, ignore + if message.text is None: + return + # if is bot command, ignore - if message.text and message.text.startswith("/"): + if message.text.startswith("/"): return # start command ignore elif message.text.startswith( @@ -218,9 +222,6 @@ def latest_handle_messages(message: Message, bot: TeleBot): # answer_it command ignore elif message.text.startswith("answer_it"): return - # if not text, ignore - elif not message.text: - return else: if chat_user_dict.get(chat_user_id): message.text += chat_message_dict[chat_id].text @@ -235,13 +236,40 @@ def answer_it_handler(message: Message, bot: TeleBot) -> None: """answer_it: /answer_it""" # answer the last message in the chat group who = "answer_it" - # get the last message in the chat chat_id = message.chat.id - latest_message = chat_message_dict.get(chat_id) - m = original_m = latest_message.text.strip() + full_answer = "" + local_image_path = "" + m = "" + original_m = "" + + # get the last message in the chat + if message.reply_to_message is not None: + latest_message = message.reply_to_message + else: + latest_message = chat_message_dict.get(chat_id) + + if latest_message.photo is not None: + max_size_photo = max(latest_message.photo, key=lambda p: p.file_size) + image_file = bot.get_file(max_size_photo.file_id).file_path + downloaded_file = bot.download_file(image_file) + local_image_path = "answer_it_temp.jpg" + with open(local_image_path, "wb") as temp_file: + temp_file.write(downloaded_file) + + m = original_m = remove_prompt_prefix(latest_message.caption.strip()) + ph_image_url = ph.upload_image(local_image_path) + full_answer += f"\n![Image]({ph_image_url})\n" + else: + m = original_m = remove_prompt_prefix(latest_message.text.strip()) + + if not m: + bot.reply_to(message, "The message was retrieved, but the prompt is empty.") + return + m = enrich_text_with_urls(m) - full_answer = f"Question:\n{m}\n" if len(m) < 300 else "" + full_answer += f"Question:\n{m}\n" if len(m) < 300 else "" + if len(m) > MESSAGE_MAX_LENGTH: a = ( "The message is too long, please shorten it." @@ -255,27 +283,35 @@ def answer_it_handler(message: Message, bot: TeleBot) -> None: #### Answers Thread #### executor = ThreadPoolExecutor(max_workers=Stream_Thread) if GEMINI_USE_THREAD and GOOGLE_GEMINI_KEY: - gemini_future = executor.submit(gemini_answer, latest_message, bot, m) + gemini_future = executor.submit( + gemini_answer, latest_message, bot, m, local_image_path + ) if CHATGPT_USE and CHATGPT_API_KEY: - chatgpt_future = executor.submit(chatgpt_answer, latest_message, bot, m) - if COHERE_USE and COHERE_API_KEY: + chatgpt_future = executor.submit( + chatgpt_answer, latest_message, bot, m, local_image_path + ) + if COHERE_USE and COHERE_API_KEY and not local_image_path: cohere_future = executor.submit(cohere_answer, latest_message, bot, m) - if QWEN_USE and QWEN_API_KEY: + if QWEN_USE and QWEN_API_KEY and not local_image_path: qwen_future = executor.submit(qwen_answer, latest_message, bot, m) if CLADUE_USE and ANTHROPIC_API_KEY: - claude_future = executor.submit(claude_answer, latest_message, bot, m) - if LLAMA_USE and LLAMA_API_KEY: + claude_future = executor.submit( + claude_answer, latest_message, bot, m, local_image_path + ) + if LLAMA_USE and LLAMA_API_KEY and not local_image_path: llama_future = executor.submit(llama_answer, latest_message, bot, m) #### Complete Message Thread #### executor2 = ThreadPoolExecutor(max_workers=Complete_Thread) - if CHATGPT_COMPLETE and CHATGPT_API_KEY: - complete_chatgpt_future = executor2.submit(complete_chatgpt, m) - if CLADUE_COMPLETE and ANTHROPIC_API_KEY: - complete_claude_future = executor2.submit(complete_claude, m) - if LLAMA_COMPLETE and LLAMA_API_KEY: + if not CHATGPT_USE and CHATGPT_COMPLETE and CHATGPT_API_KEY: + complete_chatgpt_future = executor2.submit( + complete_chatgpt, m, local_image_path + ) + if not CLADUE_USE and CLADUE_COMPLETE and ANTHROPIC_API_KEY: + complete_claude_future = executor2.submit(complete_claude, m, local_image_path) + if not LLAMA_USE and LLAMA_COMPLETE and LLAMA_API_KEY and not local_image_path: complete_llama_future = executor2.submit(complete_llama, m) - if COHERE_COMPLETE and COHERE_API_KEY: + if not COHERE_USE and COHERE_COMPLETE and COHERE_API_KEY and not local_image_path: complete_cohere_future = executor2.submit(complete_cohere, m) #### Gemini Answer Individual #### @@ -284,7 +320,12 @@ def answer_it_handler(message: Message, bot: TeleBot) -> None: g_s = "" g_reply_id = bot_reply_first(latest_message, g_who, bot) try: - g_r = convo.send_message(m, stream=True) + if local_image_path: + gemini_image_file = genai.upload_file(path=local_image_path) + g_r = convo.send_message([m, gemini_image_file], stream=True) + else: + g_r = convo.send_message(m, stream=True) + g_start = time.time() g_overall_start = time.time() for e in g_r: @@ -328,11 +369,11 @@ def answer_it_handler(message: Message, bot: TeleBot) -> None: anaswer_chatgpt, chatgpt_chat_id = chatgpt_future.result() full_chat_id_list.append(chatgpt_chat_id) full_answer += anaswer_chatgpt - if COHERE_USE and COHERE_API_KEY: + if COHERE_USE and COHERE_API_KEY and not local_image_path: answer_cohere, cohere_chat_id = cohere_future.result() full_chat_id_list.append(cohere_chat_id) full_answer += answer_cohere - if QWEN_USE and QWEN_API_KEY: + if QWEN_USE and QWEN_API_KEY and not local_image_path: answer_qwen, qwen_chat_id = qwen_future.result() full_chat_id_list.append(qwen_chat_id) full_answer += answer_qwen @@ -340,19 +381,19 @@ def answer_it_handler(message: Message, bot: TeleBot) -> None: answer_claude, claude_chat_id = claude_future.result() full_chat_id_list.append(claude_chat_id) full_answer += answer_claude - if LLAMA_USE and LLAMA_API_KEY: + if LLAMA_USE and LLAMA_API_KEY and not local_image_path: answer_llama, llama_chat_id = llama_future.result() full_chat_id_list.append(llama_chat_id) full_answer += answer_llama #### Complete Messages #### - if CHATGPT_COMPLETE and CHATGPT_API_KEY: + if not CHATGPT_USE and CHATGPT_COMPLETE and CHATGPT_API_KEY: full_answer += complete_chatgpt_future.result() - if CLADUE_COMPLETE and ANTHROPIC_API_KEY: + if not CLADUE_USE and CLADUE_COMPLETE and ANTHROPIC_API_KEY: full_answer += complete_claude_future.result() - if COHERE_COMPLETE and COHERE_API_KEY: + if not COHERE_USE and COHERE_COMPLETE and COHERE_API_KEY and not local_image_path: full_answer += complete_cohere_future.result() - if LLAMA_COMPLETE and LLAMA_API_KEY: + if not LLAMA_USE and LLAMA_COMPLETE and LLAMA_API_KEY and not local_image_path: full_answer += complete_llama_future.result() print(full_chat_id_list) @@ -391,14 +432,18 @@ def llm_background_ph_update(path: str, full_answer: str, m: str) -> str: return full_answer -def gemini_answer(latest_message: Message, bot: TeleBot, m): +def gemini_answer(latest_message: Message, bot: TeleBot, m, local_image_path): """gemini answer""" who = "Gemini Pro" # show something, make it more responsible reply_id = bot_reply_first(latest_message, who, bot) try: - r = convo.send_message(m, stream=True) + if local_image_path: + gemini_image_file = genai.upload_file(path=local_image_path) + r = convo.send_message([m, gemini_image_file], stream=True) + else: + r = convo.send_message(m, stream=True) s = "" start = time.time() overall_start = time.time() @@ -432,12 +477,25 @@ def gemini_answer(latest_message: Message, bot: TeleBot, m): return llm_answer(who, s), reply_id.message_id -def chatgpt_answer(latest_message: Message, bot: TeleBot, m): +def chatgpt_answer(latest_message: Message, bot: TeleBot, m, local_image_path): """chatgpt answer""" who = "ChatGPT Pro" reply_id = bot_reply_first(latest_message, who, bot) player_message = [{"role": "user", "content": m}] + if local_image_path: + player_message = [ + { + "role": "user", + "content": [ + {"type": "text", "text": m}, + { + "type": "image_url", + "image_url": {"url": image_to_data_uri(local_image_path)}, + }, + ], + } + ] try: r = client.chat.completions.create( @@ -472,14 +530,29 @@ def chatgpt_answer(latest_message: Message, bot: TeleBot, m): return llm_answer(who, s), reply_id.message_id -def claude_answer(latest_message: Message, bot: TeleBot, m): +def claude_answer(latest_message: Message, bot: TeleBot, m, local_image_path): """claude answer""" who = "Claude Pro" reply_id = bot_reply_first(latest_message, who, bot) + player_message = [{"role": "user", "content": m}] + if local_image_path: + player_message = [ + { + "role": "user", + "content": [ + {"type": "text", "text": m}, + { + "type": "image_url", + "image_url": {"url": image_to_data_uri(local_image_path)}, + }, + ], + } + ] + try: r = claude_client.chat.completions.create( - messages=[{"role": "user", "content": m}], + messages=player_message, max_tokens=4096, model=ANTHROPIC_MODEL, stream=True, @@ -710,12 +783,26 @@ def llm_summary(bot, full_answer, ph_s, reply_id) -> str: return s -def complete_chatgpt(m: str) -> str: +def complete_chatgpt(m: str, local_image_path: str) -> str: """we run chatgpt get the full answer""" who = "ChatGPT Pro" + player_message = [{"role": "user", "content": m}] + if local_image_path: + player_message = [ + { + "role": "user", + "content": [ + {"type": "text", "text": m}, + { + "type": "image_url", + "image_url": {"url": image_to_data_uri(local_image_path)}, + }, + ], + } + ] try: r = client.chat.completions.create( - messages=[{"role": "user", "content": m}], + messages=player_message, max_tokens=4096, model=CHATGPT_PRO_MODEL, ) @@ -727,12 +814,27 @@ def complete_chatgpt(m: str) -> str: return content -def complete_claude(m: str) -> str: +def complete_claude(m: str, local_image_path: str) -> str: """we run claude get the full answer""" who = "Claude Pro" + + player_message = [{"role": "user", "content": m}] + if local_image_path: + player_message = [ + { + "role": "user", + "content": [ + {"type": "text", "text": m}, + { + "type": "image_url", + "image_url": {"url": image_to_data_uri(local_image_path)}, + }, + ], + } + ] try: r = claude_client.chat.completions.create( - messages=[{"role": "user", "content": m}], + messages=player_message, max_tokens=4096, model=ANTHROPIC_MODEL, )