WeChatMsg/exporter/exporter_json.py

import json
import random
import os

from wxManager import Me, MessageType
from exporter.exporter import ExporterBase, remove_privacy_info, get_new_filename


class JsonStrategy:
    SPLIT_BY_TIME = 0  # 距离第一条消息的时间范围
    SPLIT_BY_INTERVALS = 1  # 相邻消息的时间间隔
    SLIDING_WINDOW = 2  # 滑动窗口法分割


class AssistantUser:
    SELF = 0  # 自己是ai助手
    CONTACT = 1  # 好友是ai助手


class JsonConfig:
    prompt: str = ''
    shuffle: bool = True  # 是否随机打乱数据
    train_ratio: int = 80  # 训练集占比（百分比）
    model: str = 'Alpaca'  # 可选：GLM4，ChatGLM3
    model_keys = {
        'GLM4': 'messages',
        'ChatGLM3': 'conversations'
    }
    strategy: int = JsonStrategy.SPLIT_BY_INTERVALS  # json导出策略
    intervals: int = 120  # 相邻两条消息的最大间隔时间
    span: int = 300  # 第一条消息跟最后一条消息的间隔时间
    window_size: int = 10  # 窗口大小
    step: int = 3  # 步长
    assistant = AssistantUser.SELF

    def get_model_keys(self):
        return self.model_keys.get(self.model, 'messages')


def modify(output, history):
    return output


def merge_content(conversions_list) -> list:
    """
    合并一组对话中连续发送的句子
    @param conversions_list:
    @return:
    """
    merged_data = []
    current_role = None
    current_content = ""
    str_time = ''
    for item in conversions_list:
        if 'str_time' in item:
            str_time = item['str_time']
        else:
            str_time = ''
        if current_role is None:
            current_role = item["role"]
            current_content = item["content"]
        elif current_role == item["role"]:
            current_content += "，" + item["content"]
        else:
            # merged_data.append({"role": current_role, "content": current_content, 'str_time': str_time})
            if len(current_content) < 3 and current_role == 'assistant':
                current_content = modify(current_content, merged_data)
            merged_data.append({"role": current_role, "content": current_content})
            current_role = item["role"]
            current_content = item["content"]
            str_time = item.get('str_time')

    # 处理最后一组
    if current_role is not None:
        # merged_data.append({"role": current_role, "content": current_content,'str_time': str_time})
        merged_data.append({"role": current_role, "content": current_content})
    return merged_data


def is_first_msg(conversions):
    if not conversions:
        return True
    else:
        return len(conversions) == 1 and conversions[0]['role'] == 'system'


def conversion_to_history(conversations):
    res = []
    has_system_prompt = conversations[0].get('role') == 'system'
    s_index, e_index = (1, len(conversations) - 3) if has_system_prompt else (0, len(conversations) - 2)
    for i in range(s_index, e_index, 2):
        res.append(
            [
                conversations[i].get('content'), conversations[i + 1].get('content')
            ]
        )
    return res


class JsonExporter(ExporterBase):
    def __init__(
            self,
            database,
            contact,
            output_dir,
            type_,  # 导出文件类型
            message_types: set[MessageType] = None,  # 导出的消息类型
            time_range=None,  # 导出的日期范围
            group_members: set[str] = None,  # 群聊中只导出这些人的聊天记录
            progress_callback=None,  # 进度回调函数，func(progress:float)
            finish_callback=None,  # 导出完成回调函数
            json_config: JsonConfig = None
    ):
        super().__init__(database, contact, output_dir, type_, message_types, time_range, group_members,
                         progress_callback, finish_callback)  # 调用父类的构造函数
        if json_config:
            self.json_config: JsonConfig = json_config
        else:
            self.json_config = JsonConfig()

    def is_user(self, is_send):
        """
        判断一条消息是否是user角色发送的
        @param is_send:
        @return:
        """
        return is_send ^ (self.json_config.assistant == AssistantUser.SELF)

    def system_prompt(self):
        system = {
            "role": "system",
            "content": self.json_config.prompt.replace(
                '{{name}}', Me().name
            ).replace(
                '{{remark}}', self.contact.remark
            )
        }
        return system

    def message_to_conversion(self, group):
        conversions = [self.system_prompt()] if self.json_config.prompt else []
        # 确保最后一条消息是assistant发出的
        while len(group) and self.is_user(group[-1].is_sender):
            group.pop()
        for message in group:
            is_send = message.is_sender
            text = remove_privacy_info(message.content)
            # 确保第一条消息必须是user发出的
            if is_first_msg(conversions) and not self.is_user(is_send):
                continue
            if self.is_user(is_send):
                json_msg = {
                    "role": "user",
                    "content": text
                }
            else:
                json_msg = {
                    "role": "assistant",
                    "content": text
                }
            json_msg['str_time'] = message.str_time
            conversions.append(json_msg)
        if len(conversions) == 1:
            return []
        return merge_content(conversions)

    def split_by_time(self, length=300):
        """
        通过第一条消息和最后一条消息的时间间隔分割数据集
        @param length:
        @return:
        """
        messages = self.database.get_messages_by_type(self.contact.wxid, type_=MessageType.Text,
                                                      time_range=self.time_range)
        start_time = 0
        res = []
        i = 0
        while i < len(messages):
            message = messages[i]
            timestamp = message.timestamp
            is_send = message.is_sender
            group = []
            while i < len(messages) and timestamp - start_time < length:
                group.append(message)
                i += 1
                if i >= len(messages):
                    break
                message = messages[i]
                timestamp = message.timestamp
                is_send = message.is_sender
            while not self.is_user(is_send):
                group.append(message)
                i += 1
                if i >= len(messages):
                    break
                message = messages[i]
                timestamp = message.timestamp
                is_send = message.is_sender
            start_time = timestamp
            if len(group) > 4:
                res.append(group)
        return res

    def split_by_intervals(self, max_diff_seconds=300):
        """
        通过相邻两条消息的时间间隔分割数据集
        @param max_diff_seconds:
        @return:
        """
        messages = self.database.get_messages_by_type(self.contact.wxid, type_=MessageType.Text,
                                                      time_range=self.time_range)
        res = []
        i = 0
        current_group = []
        while i < len(messages):
            message = messages[i]
            timestamp = message.timestamp
            is_send = message.is_sender
            while not self.is_user(is_send) and i + 1 < len(messages):
                i += 1
                message = messages[i]
                is_send = message.is_sender
            current_group = [messages[i]]
            i += 1
            while i < len(messages) and messages[i].timestamp - current_group[-1].timestamp <= max_diff_seconds:
                current_group.append(messages[i])
                i += 1
            while i < len(messages) and not self.is_user(messages[i].is_sender):
                current_group.append(messages[i])
                i += 1
            if len(current_group) > 4:
                res.append(current_group)
        return res

    def split_by_window(self, window_size=10, step=3):
        """
        滑动窗口切分数据集
        @param window_size:
        @param step:
        @return:
        """
        messages = self.database.get_messages_by_type(self.contact.wxid, type_=MessageType.Text,
                                                      time_range=self.time_range)
        res = []
        i = 0
        while i < len(messages):
            message = messages[i]
            timestamp = message.timestamp
            is_send = message.is_sender
            current_group = []
            j = i
            while not self.is_user(is_send) and j + 1 < len(messages) and j - i < window_size:
                j += 1
                message = messages[j]
                is_send = message.is_sender
            current_group = [messages[j]]
            j += 1
            while j < len(messages) and j - i < window_size:
                current_group.append(messages[j])
                j += 1
            res.append(current_group)
            i += step
        return res

    def export(self):
        print(f"【开始导出 json {self.contact.remark}】")
        origin_path = self.origin_path
        filename = os.path.join(origin_path, f"{self.contact.remark}.json")
        filename = get_new_filename(filename)
        messages_groups = []
        match self.json_config.strategy:
            case JsonStrategy.SPLIT_BY_INTERVALS:
                messages_groups = self.split_by_intervals(self.json_config.intervals)
            case JsonStrategy.SPLIT_BY_TIME:
                messages_groups = self.split_by_time(self.json_config.span)
            case JsonStrategy.SLIDING_WINDOW:
                messages_groups = self.split_by_window(self.json_config.window_size, self.json_config.step)
        dataset = []
        self.update_progress_callback(0.5)
        for group in messages_groups:
            conversations = self.message_to_conversion(group)
            if conversations:
                if self.json_config.model == 'Alpaca':
                    has_system_prompt = conversations[0].get('role') == 'system'
                    dataset.append(
                        {
                            'system': conversations[0].get('content') if has_system_prompt else '',
                            'instruction': conversations[-2].get('content'),
                            'input': '',
                            'output': conversations[-1].get('content'),
                            'history': conversion_to_history(conversations),
                        }
                    )
                else:
                    dataset.append({
                        self.json_config.get_model_keys(): conversations
                    })
        if self.json_config.shuffle:
            # 打乱列表顺序
            random.shuffle(dataset)
        with open(filename, "w", encoding="utf-8") as f:
            json.dump(dataset, f, ensure_ascii=False, indent=4)
        print(f"【完成导出 json {self.contact.remark}】")
        self.update_progress_callback(1)
        self.finish_callback(self.exporter_id)