Commit 9709f8c

benny-dou <60535774+benny-dou@users.noreply.github.com>
2026-01-19 01:05:03
feat(message): add support for `message_thread_id` in prefix/suffix modification and replacement
1 parent eb84ebc
src/ai/texts/contexts.py
@@ -18,7 +18,7 @@ from ai.utils import BOT_TIPS, clean_context
 from asr.utils import GEMINI_AUDIO_EXT, downsampe_audio
 from config import AI, DOWNLOAD_DIR
 from database.r2 import head_cf_r2
-from messages.parser import parse_msg
+from messages.parser import get_thread_id, parse_msg
 from utils import convert_md, read_text
 
 if TYPE_CHECKING:
@@ -108,7 +108,8 @@ async def get_openai_response_contexts(client: Client, message: Message, api_key
             previous_response_id: str
         """
         key_hash = hashlib.sha256(api_key.encode()).hexdigest()
-        resp = await head_cf_r2(f"TTL/{cache_day}d/OpenAI/{model_id}/{key_hash}/{msg.chat.id}/{msg.id}")
+        tid = get_thread_id(msg)
+        resp = await head_cf_r2(f"TTL/{cache_day}d/OpenAI/{model_id}/{key_hash}/{msg.chat.id}/{msg.id}{'/' + str(tid) if tid else ''}")
         return glom(resp, "Metadata.response_id", default="") or ""
 
     previous_response_id = ""
src/ai/texts/openai_response.py
@@ -15,6 +15,7 @@ from ai.texts.contexts import get_openai_response_contexts
 from ai.utils import BOT_TIPS, EMOJI_REASONING_BEGIN, EMOJI_TEXT_BOT, beautify_llm_response, literal_eval
 from config import AI, PROXY, TEXT_LENGTH
 from database.r2 import set_cf_r2
+from messages.parser import get_thread_id
 from messages.progress import modify_progress
 from messages.utils import blockquote, count_without_entities, delete_message, smart_split
 from utils import number_to_emoji, strings_list
@@ -98,8 +99,9 @@ async def openai_responses_api(
                 day = round(cache_response_ttl // 86400)
                 for sent_msg in sent_messages:  # save the reponse to R2
                     key_hash = hashlib.sha256(api_key.encode()).hexdigest()
+                    tid = get_thread_id(sent_msg)
                     await set_cf_r2(
-                        f"TTL/{day}d/OpenAI/{model_id}/{key_hash}/{sent_msg.chat.id}/{sent_msg.id}",
+                        f"TTL/{day}d/OpenAI/{model_id}/{key_hash}/{sent_msg.chat.id}/{sent_msg.id}{'/' + str(tid) if tid else ''}",
                         data=resp["full_response"],
                         metadata={"response_id": resp["response_id"]},
                         silent=silent,
src/ai/chat_summary.py
@@ -180,7 +180,6 @@ async def ai_chat_summary(
     if kwargs.get("show_progress") and "progress" not in kwargs:
         res = await send2tg(client, message, texts=f"📝正在获取历史消息...\n⏩开始时间: {begin_time:%m-%d %H:%M:%S}\n⏯️结束时间: {end_time:%m-%d %H:%M:%S}", **kwargs)
         kwargs["progress"] = res[0]
-    __import__("ipdb").set_trace(context=15, cond=True)
     history_list = await get_history_info_list(client, info["cid"], offset_id, num_history, begin_time, end_time, filter_users)
     # parse the history contexts
     parsed = await parse_history_list(history_list)
src/messages/modify.py
@@ -8,6 +8,7 @@ from loguru import logger
 from pyrogram.types import Message
 from pyrogram.types.messages_and_media.message import Str
 
+from messages.parser import get_thread_id
 from utils import slim_cid
 
 
@@ -28,22 +29,28 @@ def add_prefix_suffix(message: Message) -> Message:
     """This function add prefix and suffix to the message via environment variable.
 
     Environment Variables:
-        MOD_MSG_ADD_PREFIX_C{CID}_U{UID} (str): Add prefix to message of {CID} and {UID}
-        MOD_MSG_ADD_SUFFIX_C{CID}_U{UID} (str): Add suffix to message of {CID} and {UID}
         MOD_MSG_ADD_PREFIX_C{CID} (str): Add prefix to message of {CID}
         MOD_MSG_ADD_SUFFIX_C{CID} (str): Add suffix to message of {CID}
+        MOD_MSG_ADD_PREFIX_C{CID}_T{TID} (str): Add prefix to message of {CID}-{TID}
+        MOD_MSG_ADD_SUFFIX_C{CID}_T{TID} (str): Add suffix to message of {CID}-{TID}
         MOD_MSG_ADD_PREFIX_U{UID} (str): Add prefix to message of {UID}
         MOD_MSG_ADD_SUFFIX_U{UID} (str): Add suffix to message of {UID}
+        MOD_MSG_ADD_PREFIX_C{CID}_U{UID} (str): Add prefix to message of {CID} and {UID}
+        MOD_MSG_ADD_SUFFIX_C{CID}_U{UID} (str): Add suffix to message of {CID} and {UID}
+        MOD_MSG_ADD_PREFIX_C{CID}_T{TID}_U{UID} (str): Add prefix to message of {CID}-{TID} and {UID}
+        MOD_MSG_ADD_SUFFIX_C{CID}_T{TID}_U{UID} (str): Add suffix to message of {CID}-{TID} and {UID}
 
     If there are multiple envvars for the same message, they will be applied in the order of
-        Chat Level -> User Level -> Chat_User Level
+        Chat Level -> Chat_Thread Level -> User Level -> Chat_User Level -> Chat_Thread_User Level
 
     Example:
-        Suppose the message is Message(chat_id=111, from_user=User(id=222), content="hello")
+        Suppose the message is Message(chat_id=111,from_user=User(id=222), content="hello", message_thread_id=333)
 
         # Single Environment Variable:
         MOD_MSG_ADD_PREFIX_C111 = "foo"  -> "foo hello" for chat_id=111
         MOD_MSG_ADD_SUFFIX_C111 = "bar"  -> "hello bar" for chat_id=111
+        MOD_MSG_ADD_PREFIX_C111_T333 = "foo"  -> "foo hello" for chat_id=111, message_thread_id=333
+        MOD_MSG_ADD_SUFFIX_C111_T333 = "bar"  -> "hello bar" for chat_id=111, message_thread_id=333
         MOD_MSG_ADD_PREFIX_U222 = "foo"  -> "foo hello" for user_id=222
         MOD_MSG_ADD_SUFFIX_U222 = "bar"  -> "hello bar" for user_id=222
         MOD_MSG_ADD_PREFIX_C111_U222 = "foo"  -> "foo hello" for chat_id=111 and user_id=222
@@ -65,6 +72,7 @@ def add_prefix_suffix(message: Message) -> Message:
     """
     uid = glom(message, "from_user.id", default=0) or 0
     cid = glom(message, "chat.id", default=0) or 0
+    tid = get_thread_id(message)
     cid = slim_cid(cid)
     texts = str(message.content).strip()
     # Chat level
@@ -74,6 +82,13 @@ def add_prefix_suffix(message: Message) -> Message:
     if suffix := os.getenv(f"MOD_MSG_ADD_SUFFIX_C{cid}"):
         texts = f"{texts} {suffix}"
 
+    # Chat_Thread level
+    suffix = os.getenv(f"MOD_MSG_ADD_SUFFIX_C{cid}_T{tid}")
+    if prefix := os.getenv(f"MOD_MSG_ADD_PREFIX_C{cid}_T{tid}"):
+        texts = f"{prefix} {texts}"
+    if suffix := os.getenv(f"MOD_MSG_ADD_SUFFIX_C{cid}_T{tid}"):
+        texts = f"{texts} {suffix}"
+
     # User level
     if prefix := os.getenv(f"MOD_MSG_ADD_PREFIX_U{uid}"):
         texts = f"{prefix} {texts}"
@@ -86,6 +101,12 @@ def add_prefix_suffix(message: Message) -> Message:
     if suffix := os.getenv(f"MOD_MSG_ADD_SUFFIX_C{cid}_U{uid}"):
         texts = f"{texts} {suffix}"
 
+    # Chat_Thread_User level
+    if prefix := os.getenv(f"MOD_MSG_ADD_PREFIX_C{cid}_T{tid}_U{uid}"):
+        texts = f"{prefix} {texts}"
+    if suffix := os.getenv(f"MOD_MSG_ADD_SUFFIX_C{cid}_T{tid}_U{uid}"):
+        texts = f"{texts} {suffix}"
+
     if texts.strip() != str(message.content).strip():  # If the message is modified
         if message.text:
             message.text = Str(texts)
@@ -100,8 +121,10 @@ def replace_prefix(message: Message) -> Message:
     Environment Variables:
         MOD_MSG_REPLACE_PREFIX_GLOBAL_{prefix} (str): Replace prefix to all messages (prefix is case-sensitive)
         MOD_MSG_REPLACE_PREFIX_C{CID}_{prefix} (str): Replace prefix to message of {CID} (prefix is case-sensitive)
+        MOD_MSG_REPLACE_PREFIX_CT{CID}_{TID}_{prefix} (str): Replace prefix to message of {CID}-{TID} (prefix is case-sensitive)
         MOD_MSG_REPLACE_PREFIX_U{UID}_{prefix} (str): Replace prefix to message of {UID} (prefix is case-sensitive)
         MOD_MSG_REPLACE_PREFIX_CU{CID}_{UID}_{prefix} (str): Replace prefix to message of {CID} and {UID} (prefix is case-sensitive)
+        MOD_MSG_REPLACE_PREFIX_CTU{CID}_{TID}_{UID}_{prefix} (str): Replace prefix to message of {CID}-{TID} and {UID} (prefix is case-sensitive)
 
     If there are multiple envvars for the same message, they will be applied in the order of
         Global Level -> Chat Level -> User Level -> Chat_User Level
@@ -126,6 +149,8 @@ def replace_prefix(message: Message) -> Message:
     """
     uid = glom(message, "from_user.id", default=0) or 0
     cid = glom(message, "chat.id", default=0) or 0
+    tid = get_thread_id(message)
+    cid = slim_cid(cid)
     texts = str(message.content).strip()
 
     # Global level
@@ -150,6 +175,17 @@ def replace_prefix(message: Message) -> Message:
             texts = texts.replace(prefix, env_var, count=1)
             logger.warning(f"Chat ({cid}) level prefix `{prefix} = {env_var}` -> {texts}")
 
+    # Chat_Thread level
+    chat_thread_envs = [x for x in os.environ if x.upper().startswith(f"MOD_MSG_REPLACE_PREFIX_CT{cid}_{tid}_")]
+    for env in sorted(chat_thread_envs, key=lambda x: len(x), reverse=True):
+        env_var = os.environ[env]
+        prefix = env[len(f"MOD_MSG_REPLACE_PREFIX_CT{cid}_{tid}_") :]
+        prefix = escape_strings(prefix)
+        env_var = escape_strings(env_var)
+        if texts.startswith(prefix):
+            texts = texts.replace(prefix, env_var, count=1)
+            logger.warning(f"Chat ({cid}), Thread ({tid}) level prefix `{prefix} = {env_var}` -> {texts}")
+
     # User level
     user_envs = [x for x in os.environ if x.upper().startswith(f"MOD_MSG_REPLACE_PREFIX_U{uid}_")]
     for env in sorted(user_envs, key=lambda x: len(x), reverse=True):
@@ -172,6 +208,17 @@ def replace_prefix(message: Message) -> Message:
             texts = texts.replace(prefix, env_var, count=1)
             logger.warning(f"Chat ({cid}), User ({uid}) level prefix `{prefix} = {env_var}` -> {texts}")
 
+    # Chat_Thread_User level
+    chat_thread_user_envs = [x for x in os.environ if x.upper().startswith(f"MOD_MSG_REPLACE_PREFIX_CTU{cid}_{tid}_{uid}_")]
+    for env in sorted(chat_thread_user_envs, key=lambda x: len(x), reverse=True):
+        env_var = os.environ[env]
+        prefix = env[len(f"MOD_MSG_REPLACE_PREFIX_CTU{cid}_{tid}_{uid}_") :]
+        prefix = escape_strings(prefix)
+        env_var = escape_strings(env_var)
+        if texts.startswith(prefix):
+            texts = texts.replace(prefix, env_var, count=1)
+            logger.warning(f"Chat ({cid}), Thread ({tid}), User ({uid}) level prefix `{prefix} = {env_var}` -> {texts}")
+
     if texts != str(message.content).strip():  # If the message is modified
         if message.text:
             message.text = Str(texts)
src/messages/parser.py
@@ -13,7 +13,7 @@ from pyrogram.types import Chat, Message
 
 from config import TZ, cache
 from others.emoji import CTYPE_EMOJI, MTYPE_EMOJI
-from utils import nowdt, slim_cid
+from utils import nowdt
 
 
 def parse_msg(message: Message, *, silent: bool = False, verbose: bool = False, use_cache: bool = True) -> dict:
@@ -31,13 +31,14 @@ def parse_msg(message: Message, *, silent: bool = False, verbose: bool = False,
     chandle = glom(message, "chat.username", default="") or ""
     uid = glom(message, "from_user.id", default=1) or 1  # uid must > 0
     cid = glom(message, "chat.id", default=0) or 0
+    tid = get_thread_id(message)
     mid = glom(message, "id", default=0) or 0
     media_group_id = glom(message, "media_group_id", default=0) or 0
     is_bot = glom(message, "from_user.is_bot", default=False)
     text = message.content
     dt = message.date.astimezone(ZoneInfo(TZ)) if isinstance(message.date, datetime) else nowdt(TZ)
     time = f"{dt:%Y-%m-%d %H:%M:%S}"
-    message_url = f"https://t.me/{chandle}/{mid}" if chandle else f"https://t.me/c/{slim_cid(cid)}/{mid}"
+    message_url = glom(message, "link", default="") or ""
     # parse user attributes
     first_name = glom(message, "from_user.first_name", default="") or ""
     last_name = glom(message, "from_user.last_name", default="") or ""
@@ -103,6 +104,7 @@ def parse_msg(message: Message, *, silent: bool = False, verbose: bool = False,
         "chandle": str(chandle),
         "uid": int(uid),
         "cid": int(cid),
+        "tid": int(tid),
         "mid": int(mid),
         "media_group_id": int(media_group_id),
         "is_bot": bool(is_bot),
@@ -162,3 +164,14 @@ def parse_chat(chat: Chat, *, use_cache: bool = True) -> dict:
     if use_cache:
         cache.set(f"parse_chat-{chat.id}", info, ttl=120)  # cache the same msg for 2 minutes
     return info
+
+
+def get_thread_id(message: Message) -> int:
+    """Get the thread ID of a message.
+
+    If the message is not a reply, return the message ID.
+    """
+    tid = glom(message, "message_thread_id", default=0) or 0
+    if not tid and glom(message, "chat.is_forum", default=False):
+        tid = 1  # this message is sent to `General` topic thread
+    return tid