SD_LarkBot/utils/config_io.py

import json
import os
from utils.logger import logger


def save_config_to_json(config_data, message_id):
    # 确保目录存在，如果不存在则创建
    output_directory = f"data/input_images/{message_id}"
    os.makedirs(output_directory, exist_ok=True)

    # 生成JSON文件路径
    json_path = os.path.join(output_directory, "config.json")

    # 将字典转换为JSON并保存到文件中
    with open(json_path, 'w') as f:
        json.dump(config_data, f)

def load_config_from_json(message_id):
    # 生成JSON文件路径
    json_path = os.path.join("data/input_images", message_id, "config.json")

    # 检查文件是否存在，如果不存在则返回默认字典
    if not os.path.exists(json_path):
        return {"user_message": "", "action": "", "lora": "", "animal": "", "if_download": False, "img_missions": {}}

    # 从文件中读取JSON数据并转换为字典
    with open(json_path, 'r') as f:
        config_data = json.load(f)

    return config_data

def add_to_user_reply(event_id, user_message_id, reply_message_id):
    output_directory = "data"
    os.makedirs(output_directory, exist_ok=True)
    json_path = os.path.join(output_directory, "user_reply.json")

    # 检查文件是否存在，如果不存在则创建
    if os.path.exists(json_path):
        # 从文件中读取JSON数据并转换为字典
        with open(json_path, 'r') as f:
            config_data = json.load(f)
        if event_id in config_data:
            return
        else:
            config_data[event_id] = [user_message_id, reply_message_id]
            with open(json_path, 'w') as f:
                json.dump(config_data, f)
            return
    else:
        config_data={
            event_id : [user_message_id, reply_message_id]
        }
        with open(json_path, 'w') as f:
            json.dump(config_data, f)
        return

def if_has_processed(event_id):
    output_directory = "data"
    os.makedirs(output_directory, exist_ok=True)
    json_path = os.path.join(output_directory, "user_reply.json")

    # 检查文件是否存在，如果不存在则创建
    if os.path.exists(json_path):
        # 从文件中读取JSON数据并转换为字典
        with open(json_path, 'r') as f:
            config_data = json.load(f)
        if event_id in config_data:
            return True
        else:
            return False
    else:
        return False
    
# 从文件中读取lora选项
def load_lora_options():
    json_path = "lora_options.json"
    with open(json_path, 'r', encoding='utf-8') as f:
        lora_options = json.load(f)
    return lora_options

# 从文件中读取animal选项
def load_animal_options():
    json_path = "animal_options.json"
    with open(json_path, 'r', encoding='utf-8') as f:
        animal_options = json.load(f)
    return animal_options


# 检测任务是否完成
def if_mission_done(reply_message_id):
    config_data = load_config_from_json(reply_message_id)
    for key, value in config_data['img_missions'].items():
        if value != 'done':
            return False
    return True

# 检测任务是否存在失败
def if_mission_fail(reply_message_id):
    config_data = load_config_from_json(reply_message_id)
    for key, value in config_data['img_missions'].items():
        if value == 'fail':
            return True
    return False