Merge pull request #45 from F4ria/answer-reply

Enhance answer_it_handler to process image and replied messages

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

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)
+    full_answer = ""
-    m = original_m = latest_message.text.strip()
+    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)
+        chatgpt_future = executor.submit(
-    if COHERE_USE and COHERE_API_KEY:
+            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)
+        claude_future = executor.submit(
-    if LLAMA_USE and LLAMA_API_KEY:
+            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:
+    if not CHATGPT_USE and CHATGPT_COMPLETE and CHATGPT_API_KEY:
-        complete_chatgpt_future = executor2.submit(complete_chatgpt, m)
+        complete_chatgpt_future = executor2.submit(
-    if CLADUE_COMPLETE and ANTHROPIC_API_KEY:
+            complete_chatgpt, m, local_image_path
-        complete_claude_future = executor2.submit(complete_claude, m)
+        )
-    if LLAMA_COMPLETE and LLAMA_API_KEY:
+    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,
         )