MessageUtils更新

pytests/utils_test/message_utils_test.py

@@ -203,7 +203,7 @@ def load_message_via_file(monkeypatch):
 def dummy_number_to_short_id(original_id: int, salt: str, length: int = 6) -> str:
     return "X" * length  # 返回固定的字符串，长度由参数决定，模拟生成短ID的行为
 
-def dummy_is_bot_self(user_id: str) -> bool:
+def dummy_is_bot_self(user_id: str, platform) -> bool:
     return user_id == "bot_self"
 
 def load_utils_via_file(monkeypatch):
@@ -212,6 +212,8 @@ def load_utils_via_file(monkeypatch):
     # Mock math_utils 模块，供 from .math_utils import number_to_short_id 使用
     math_utils_mod = ModuleType("src.common.utils.math_utils")
     math_utils_mod.number_to_short_id = dummy_number_to_short_id
+    math_utils_mod.TimestampMode = type("TimestampMode", (), {"NORMAL": "%Y-%m-%d %H:%M:%S", "NORMAL_NO_YMD": "%H:%M:%S", "RELATIVE": "relative"})
+    math_utils_mod.translate_timestamp_to_human_readable = lambda timestamp, mode: "2024-01-01 12:00:00"  # 返回固定的时间字符串
     monkeypatch.setitem(sys.modules, "src.common.utils.math_utils", math_utils_mod)
 
     # 确保包层级模块存在于 sys.modules 中，使相对导入能正确解析
@@ -252,7 +254,7 @@ async def test_build_readable_message_basic(monkeypatch):
     user_info = UserInfo(user_id="u1", user_nickname="Alice")
     msg.message_info = MessageInfo(user_info=user_info)
     msg.raw_message = MessageSequence([TextComponent("Hello world")])
-    text, mapping = await MessageUtils.build_readable_message([msg], anonymize=False, show_lineno=True)
+    text, mapping, _ = await MessageUtils.build_readable_message([msg], anonymize=False, show_lineno=True)
     assert "[1] Alice说：Hello world" in text
     assert mapping == {}
 
@@ -270,7 +272,7 @@ async def test_build_readable_message_anonymize(monkeypatch):
     user_info = UserInfo(user_id="u42", user_nickname="Bob")
     msg.message_info = MessageInfo(user_info=user_info)
     msg.raw_message = MessageSequence([TextComponent("Secret text")])
-    text, mapping = await MessageUtils.build_readable_message([msg], anonymize=True, show_lineno=False)
+    text, mapping, _ = await MessageUtils.build_readable_message([msg], anonymize=True, show_lineno=False)
     # 根据实现，original_name 为 user_nickname，因此文本中应包含原始名称
     assert "XXXXXX说：" in text
     assert "u42" in mapping
@@ -291,7 +293,7 @@ async def test_build_readable_message_replace_bot(monkeypatch):
     user_info = UserInfo(user_id="bot_self", user_nickname="SomeBot")
     msg.message_info = MessageInfo(user_info=user_info)
     msg.raw_message = MessageSequence([TextComponent("ping")])
-    text, mapping = await MessageUtils.build_readable_message([msg], replace_bot_name=True, target_bot_name="MAIBot")
+    text, mapping, _ = await MessageUtils.build_readable_message([msg], replace_bot_name=True, target_bot_name="MAIBot")
     assert "MAIBot说：ping" in text
 
 
@@ -309,7 +311,7 @@ async def test_build_readable_message_image_extraction(monkeypatch):
     msg.session_id = "s_img"
     msg.raw_message = MessageSequence([img])
     msg.message_info = MessageInfo(UserInfo(user_id="ui_img", user_nickname="ImgUser"))
-    text, mapping = await MessageUtils.build_readable_message([msg], extract_pictures=True)
+    text, mapping, _ = await MessageUtils.build_readable_message([msg], extract_pictures=True)
     # 应包含图片描述占位
     assert "图片1" in text
     # mapping 不为空（匿名化未开启则为空）
@@ -333,7 +335,7 @@ async def test_build_readable_message_anonymize_and_replace_bot_name_and_lineno(
     msg2.message_info = MessageInfo(UserInfo(user_id="bot_self", user_nickname="SomeBot"))
     msg1.raw_message = MessageSequence([TextComponent("Hi")])
     msg2.raw_message = MessageSequence([TextComponent("Hello")])
-    text, mapping = await MessageUtils.build_readable_message(
+    text, mapping, _ = await MessageUtils.build_readable_message(
         [msg1, msg2],
         anonymize=True,
         replace_bot_name=True,
@@ -361,7 +363,7 @@ async def test_build_readable_message_with_at(monkeypatch):
     msg.session_id = "s_at"
     msg.raw_message = MessageSequence([at_comp])
     msg.message_info = MessageInfo(UserInfo(user_id="u_main", user_nickname="MainUser"))
-    text, mapping = await MessageUtils.build_readable_message([msg], anonymize=True, replace_bot_name=True, target_bot_name="MAIBot")
+    text, mapping, _ = await MessageUtils.build_readable_message([msg], anonymize=True, replace_bot_name=True, target_bot_name="MAIBot")
     # 验证主消息和@组件中的用户信息都被处理
     assert "XXXXXX说：" in text  # 主消息用户被匿名化
     assert "XXXXXX说：@XXXXXX" in text  # @组件用户被匿名化

src/bw_learner/expression_learner.py

@@ -53,7 +53,7 @@ class ExpressionLearner:
         if not self._messages_cache:
             logger.debug("没有消息可供学习，跳过学习过程")
             return
-        readable_message, _ = await MessageUtils.build_readable_message(
+        readable_message, _, _ = await MessageUtils.build_readable_message(
             self._messages_cache,
             anonymize=True,
             show_lineno=True,

src/common/utils/math_utils.py

@@ -47,7 +47,7 @@ def number_to_short_id(original_id: int, salt: str, length: int = 6) -> str:
     return short_id
 
 
-def translate_timestamp_to_human_readable(timestamp: float, mode: TimestampMode) -> str:
+def translate_timestamp_to_human_readable(timestamp: float, mode: TimestampMode | str) -> str:
     """将时间戳按照指定模式转换为人类可读的格式
 
     Args:
@@ -56,6 +56,11 @@ def translate_timestamp_to_human_readable(timestamp: float, mode: TimestampMode)
     Returns:
         str: 转换后的时间字符串
     """
+    if isinstance(mode, str):
+        if mode.upper() in TimestampMode.__members__:
+            mode = TimestampMode[mode.upper()]
+        else:
+            raise ValueError(f"不支持的时间戳转换模式: {mode}")
     if mode in [TimestampMode.NORMAL, TimestampMode.NORMAL_NO_YMD]:
         return time.strftime(mode.value, time.localtime(timestamp))
     elif mode == TimestampMode.RELATIVE:

src/common/utils/utils_message.py

@@ -1,6 +1,5 @@
-from enum import Enum
 from maim_message import MessageBase, Seg
-from typing import List, Tuple, Optional, Dict, TYPE_CHECKING
+from typing import List, Tuple, Optional, Dict, TYPE_CHECKING, Callable
 
 import base64
 import hashlib
@@ -136,7 +135,7 @@ class MessageUtils:
 
     @staticmethod
     def store_message_to_db(message: "SessionMessage"):
-        """存储消息到数据库"""
+        """存储消息到数据库，此方法没有update机制"""
         from src.common.database.database import get_db_session
 
         with get_db_session() as session:
@@ -152,10 +151,12 @@ class MessageUtils:
         extract_pictures: bool = False,
         replace_bot_name: bool = False,
         target_bot_name: Optional[str] = None,
-        timestamp_mode: Optional[TimestampMode] = None,
+        timestamp_mode: Optional[TimestampMode | str] = None,
         show_message_id_prefix: bool = False,
+        read_mark_time: Optional[float] = None,
         truncate_message: bool = False,
-    ) -> Tuple[str, Dict[str, Tuple[str, str]]]:
+        truncate_func: Optional[Callable[[float], Tuple[Optional[int], str]]] = None,
+    ) -> Tuple[str, Dict[str, Tuple[str, str]], List[str]]:
         """
         将消息构建为LLM可读的文本格式
 
@@ -168,17 +169,21 @@ class MessageUtils:
             target_bot_name (Optional[str]): 如果replace_bot_name为True，指定要替换的机器人名称，比如可以把机器人名称替换为“你”
             timestamp_mode (Optional[TimestampMode]): 时间戳显示模式，默认为None表示不显示时间戳
             show_message_id_prefix (bool): 是否在每条消息前显示消息ID前缀
-            truncate_message (bool): 是否截断过长的消息文本，避免生成过长的输入给LLM
+            truncate_message (bool): 是否启用消息文本截断功能，截断过长的消息文本
+            truncate_func (Optional[Callable[[float], Tuple[Optional[int], str]]]) 截断函数，接受消息的百分位位置(0-1)，返回一个元组(文本长度限制(可为None表不切割), 替换内容)
         Returns:
-            return (Tuple[str, Dict[str, Tuple[str, str]]]): 构建后的消息文本，以及映射表（匿名ID, 原始名称）
+            return (Tuple[str, Dict[str, Tuple[str, str]], List[str]]): 构建后的消息文本，映射表 {用户ID: (匿名ID, 原始名称)}，消息编号列表
         """
         msg_list: List["SessionMessage"] = messages
         user_id_mapping: Dict[str, Tuple[str, str]] = {}  # user_id -> (匿名ID, 原始名称)
+        message_ids: List[str] = []  # 存储消息编号的列表
         copied: bool = False  # 标记是否已经复制过消息列表，避免不必要的复制开销
         img_map: Optional[Dict[str, Tuple[int, str]]] = None
         emoji_map: Optional[Dict[str, Tuple[int, str]]] = None
         if replace_bot_name and not target_bot_name:
             raise ValueError("当replace_bot_name为True时，必须指定target_bot_name参数")
+
+        # 匿名化和机器人名称处理
         if anonymize or replace_bot_name:
             user_id_mapping = {}  # 利用弱引用直接传入并得到修改结果
             anonymous_messages: List["SessionMessage"] = []
@@ -198,28 +203,99 @@ class MessageUtils:
             copied = True
 
         processed_plain_texts: List[str] = []
+
+        # 将图片提取到内容最前面
         if extract_pictures:
             img_map = {}  # binary_hash -> (图片ID, 描述信息)
             emoji_map = {}  # binary_hash -> (表情ID, 描述信息)
             msg_list = [
                 MessageUtils._extract_pictures_from_message(msg, img_map, emoji_map, copied) for msg in msg_list
             ]
+            processed_plain_texts.append("图片信息和表情信息：")
             processed_plain_texts.extend(f"[图片{img_id}: {desc}]" for img_id, desc in img_map.values())
             processed_plain_texts.append("")  # 图片和表情之间添加一个换行，避免连在一起
             processed_plain_texts.extend(f"[表情{emoji_id}: {desc}]" for emoji_id, desc in emoji_map.values())
-            processed_plain_texts.append("")  # 表情和消息文本之间添加两个换行，避免连在一起
+            processed_plain_texts.extend(("", "聊天记录信息："))
 
-        lineno_counter = 1
+        msg_count = len(msg_list)
-        for msg in msg_list:
+        read_mark_added_flag: bool = False  # 标记是否已经添加过已读标签，确保只添加一次
+        for i, msg in enumerate(msg_list):
             await msg.process()
             plain_text: str = msg.processed_plain_text  # type: ignore
-            usr_info = msg.message_info.user_info
+            if truncate_message:  # 消息截断逻辑
-            usr_name = usr_info.user_cardname or usr_info.user_nickname or "未知用户"
+                percentile = i / msg_count
-            header = f"[{lineno_counter}] {usr_name}说：" if show_lineno else f"{usr_name}说："
+                if not read_mark_time:  # 没有已读标签
-            lineno_counter += 1
+                    plain_text = MessageUtils._truncate_message(
+                        percentile,
+                        plain_text,
+                        truncate_func or MessageUtils._default_truncate_func,
+                    )
+                elif msg.timestamp.timestamp() < read_mark_time:
+                    plain_text = MessageUtils._truncate_message(
+                        percentile,
+                        plain_text,
+                        truncate_func or MessageUtils._default_truncate_func,
+                    )
+                elif not read_mark_added_flag:
+                    read_mark_added_flag = True
+                    processed_plain_texts.append("\n--- 以上消息是你已经看过，请关注以下未读的新消息---\n")
+            header, message_id = MessageUtils._build_line_header(
+                msg,
+                show_message_id_prefix,
+                i + 1,
+                show_lineno,
+                timestamp_mode,
+            )
+            if message_id is not None:
+                message_ids.append(message_id)
             processed_plain_texts.append("".join([header, plain_text]))
 
-        return "\n".join(processed_plain_texts), user_id_mapping
+        return "\n".join(processed_plain_texts), user_id_mapping, message_ids
+
+    @staticmethod
+    def _build_line_header(
+        message: "SessionMessage",
+        show_message_id_prefix: bool,
+        counter: int,
+        show_lineno: bool,
+        timestamp_mode: Optional[str | TimestampMode] = None,
+    ) -> Tuple[str, Optional[str]]:
+        usr_info = message.message_info.user_info
+        usr_name = usr_info.user_cardname or usr_info.user_nickname or "未知用户"
+        header_parts = [f"{usr_name}说："]
+        message_id = None
+        if show_message_id_prefix:
+            rand_id = f"{counter}{random.randint(10, 99)}"
+            message_id = f"m{rand_id}"
+            header_parts.insert(0, f"[消息ID: {message_id}]")
+        if timestamp_mode:
+            timestamp_str = translate_timestamp_to_human_readable(message.timestamp.timestamp(), mode=timestamp_mode)
+            header_parts.insert(0, f"[{timestamp_str}]")
+        if show_lineno:
+            header_parts.insert(0, f"[{counter}]")
+        return " ".join(header_parts), message_id
+
+    @staticmethod
+    def _truncate_message(
+        percentile: float, original_content: str, truncate_func: Callable[[float], Tuple[Optional[int], str]]
+    ):
+        limit, replacement = truncate_func(percentile)
+        if limit:
+            return f"{original_content[:limit]}{replacement}"
+        else:
+            return original_content
+
+    @staticmethod
+    def _default_truncate_func(percentile: float) -> Tuple[int, str]:
+        """默认的截断函数，根据消息在消息列表中的位置返回不同的截断长度和替换内容"""
+        if percentile < 0.3:
+            return 400, "......（内容太长了）"
+        elif percentile < 0.5:
+            return 200, "......（内容太长了）"
+        elif percentile < 0.8:
+            return 100, "......（有点记不清了）"
+        else:
+            return 50, "......（记不清了）"
 
     @staticmethod
     def _process_usr_info(
@@ -232,11 +308,13 @@ class MessageUtils:
     ):
         """处理消息中的用户信息，进行匿名化显示"""
         new_message = message.deepcopy()
+        platform = message.platform
         new_component_list = [
             MessageUtils._process_msg_component(
                 component,
                 anonymize_mapping,
                 salt,
+                platform,
                 anonymize,
                 replace_bot_name,
                 target_bot_name,
@@ -254,7 +332,7 @@ class MessageUtils:
             anonymous_name = anonymize_mapping[msg_usr_info.user_id][0]
             new_message.message_info.user_info.user_nickname = anonymous_name
             new_message.message_info.user_info.user_cardname = anonymous_name
-        if replace_bot_name and target_bot_name and is_bot_self(msg_usr_info.user_id):
+        if replace_bot_name and target_bot_name and is_bot_self(msg_usr_info.user_id, platform):
             new_message.message_info.user_info.user_nickname = target_bot_name
             new_message.message_info.user_info.user_cardname = target_bot_name
         return new_message
@@ -264,6 +342,7 @@ class MessageUtils:
         component: StandardMessageComponents,
         anonymize_mapping: Dict[str, Tuple[str, str]],
         salt: str,
+        platform: str,
         anonymize: bool,
         replace_bot_name: bool,
         target_bot_name: Optional[str] = None,
@@ -274,6 +353,7 @@ class MessageUtils:
                 component,
                 anonymize_mapping,
                 salt,
+                platform,
                 anonymize,
                 replace_bot_name,
                 target_bot_name,
@@ -283,6 +363,7 @@ class MessageUtils:
                 component,
                 anonymize_mapping,
                 salt,
+                platform,
                 anonymize,
                 replace_bot_name,
                 target_bot_name,
@@ -292,6 +373,7 @@ class MessageUtils:
                 component,
                 anonymize_mapping,
                 salt,
+                platform,
                 anonymize,
                 replace_bot_name,
                 target_bot_name,
@@ -303,6 +385,7 @@ class MessageUtils:
         component: AtComponent,
         anonymize_mapping: Dict[str, Tuple[str, str]],
         salt: str,
+        platform: str,
         anonymize: bool,
         replace_bot_name: bool,
         target_bot_name: Optional[str] = None,
@@ -319,7 +402,7 @@ class MessageUtils:
             anonymous_name = anonymize_mapping[user_id][0]
             component.target_user_nickname = anonymous_name
             component.target_user_cardname = anonymous_name
-        if replace_bot_name and target_bot_name and is_bot_self(user_id):
+        if replace_bot_name and target_bot_name and is_bot_self(user_id, platform):
             component.target_user_nickname = target_bot_name
             component.target_user_cardname = target_bot_name
         return component
@@ -329,6 +412,7 @@ class MessageUtils:
         component: ForwardNodeComponent,
         anonymize_mapping: Dict[str, Tuple[str, str]],
         salt: str,
+        platform: str,
         anonymize: bool,
         replace_bot_name: bool,
         target_bot_name: Optional[str] = None,
@@ -354,7 +438,7 @@ class MessageUtils:
                 anonymous_name = anonymize_mapping[user_id][0]
                 comp.user_nickname = anonymous_name
                 comp.user_cardname = anonymous_name
-            if replace_bot_name and target_bot_name and is_bot_self(user_id):
+            if replace_bot_name and target_bot_name and is_bot_self(user_id, platform):
                 comp.user_nickname = target_bot_name
                 comp.user_cardname = target_bot_name
             comp.content = [  # 递归处理转发消息中的组件
@@ -362,6 +446,7 @@ class MessageUtils:
                     c,
                     anonymize_mapping,
                     salt,
+                    platform,
                     anonymize,
                     replace_bot_name,
                     target_bot_name,
@@ -375,6 +460,7 @@ class MessageUtils:
         component: ReplyComponent,
         anonymize_mapping: Dict[str, Tuple[str, str]],
         salt: str,
+        platform: str,
         anonymize: bool,
         replace_bot_name: bool,
         target_bot_name: Optional[str] = None,
@@ -391,7 +477,7 @@ class MessageUtils:
                 anonymous_name = anonymize_mapping[user_id][0]
                 component.target_message_sender_nickname = anonymous_name
                 component.target_message_sender_cardname = anonymous_name
-            if replace_bot_name and target_bot_name and is_bot_self(user_id):
+            if replace_bot_name and target_bot_name and is_bot_self(user_id, platform):
                 component.target_message_sender_nickname = target_bot_name
                 component.target_message_sender_cardname = target_bot_name
         else:
@@ -447,10 +533,10 @@ class MessageUtils:
 
 
 # TODO: 这个函数的实现非常临时，后续需要替换为更完善的实现，比如直接从配置文件中读取机器人自己的ID，或者通过API获取机器人自己的信息等
-def is_bot_self(user_id: str) -> bool:
+def is_bot_self(user_id: str, platform: str) -> bool:
     """
     判断用户ID是否是机器人自己
 
     临时方法，后续会替换为更完善的实现
     """
-    return user_id == "bot_self"
+    return user_id == "bot_self" and platform == "test_platform"