cdryzun/tg_bot_collections
1
0
Fork 0
mirror of https://github.com/cdryzun/tg_bot_collections.git synced 2025-04-28 16:17:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #45 from F4ria/answer-reply

Browse Source 
Enhance answer_it_handler to process image and replied messages
This commit is contained in:
yihong 2024-07-15 12:11:31 +08:00 committed by GitHub
commit 059db58c5c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 185 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

47
handlers/__init__.py
View File

@ -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",

176
handlers/useful.py
View File

@ -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,
)