Commit 9709f8c
Changed files (5)
src
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