newbot/config.py

# encoding:utf-8
import argparse
import json
import logging
import os
import pickle

from common.log import logger

# 将所有可用的配置项写在字典里, 请使用小写字母
# 此处的配置值无实际意义，程序不会读取此处的配置，仅用于提示格式，请将配置加入到config.json中
available_setting = {
    "debug": False,
    "api_base": "http://fastgpt.ascrm.cn/api/v1",
    "api_key": "fastgpt-sKABkv3PTHxlFZYPn9Mo35HHsZSdzdFNBH4XeWIRn5CwdkG7aXqEDmXwDwK",
    "token": "bowen-test",
    "open_ai_model": "gpt-4o",
    "open_ai_temperature": 0.7,
    "open_ai_max_tokens": 1024,
    "open_ai_top_p": 1,
    "open_ai_frequency_penalty": 0,
    "open_ai_presence_penalty": 0,
    "open_ai_stop": ["<|im_end|>"],
    "open_ai_stream": True,
    "contacts_white_list": [],
    "appdata_dir": ".",

    "agent_id": 0,
    "role": "",
    "background": "",
    "examples": "",
    "dataset_id": "",
    "collection_id": "",
    "custom_agent_base": "",
    "custom_agent_key": "",
    "openai_base": "",
    "openai_key": "",
    "dataset_base": "",
    "dataset_key": "",
}


class Config(dict):
    def __init__(self, d=None):
        super().__init__()
        if d is None:
            d = {}
        for k, v in d.items():
            self[k] = v
        # user_datas: 用户数据，key为用户名，value为用户数据，也是dict
        self.user_datas = {}

    def __getitem__(self, key):
        if key not in available_setting:
            raise Exception("key {} not in available_setting".format(key))
        return super().__getitem__(key)

    def __setitem__(self, key, value):
        if key not in available_setting:
            raise Exception("key {} not in available_setting".format(key))
        return super().__setitem__(key, value)

    def get(self, key, default=None):
        try:
            return self[key]
        except KeyError as e:
            return default
        except Exception as e:
            raise e

    # Make sure to return a dictionary to ensure atomic
    def get_user_data(self, user) -> dict:
        if self.user_datas.get(user) is None:
            self.user_datas[user] = {}
        return self.user_datas[user]

    def load_user_datas(self):
        try:
            with open(os.path.join(get_appdata_dir(), "user_datas.pkl"), "rb") as f:
                self.user_datas = pickle.load(f)
                logger.info("[Config] User datas loaded.")
        except FileNotFoundError as e:
            logger.info("[Config] User datas file not found, ignore.")
        except Exception as e:
            logger.info("[Config] User datas error: {}".format(e))
            self.user_datas = {}

    def save_user_datas(self):
        try:
            with open(os.path.join(get_appdata_dir(), "user_datas.pkl"), "wb") as f:
                pickle.dump(self.user_datas, f)
                logger.info("[Config] User datas saved.")
        except Exception as e:
            logger.info("[Config] User datas error: {}".format(e))


config = Config()


# parser = argparse.ArgumentParser(description='消息中间处理程序')
# parser.add_argument('-c', '--config', help='设置配置文件，默认值是 ./config.json')


def load_config():
    global config, parser

    # args = parser.parse_args()

    # parser.print_help()

    # if args.config:
    #     config_path = args.config
    # else:
    config_path = "./config.json"
    if not os.path.exists(config_path):
        logger.info("配置文件不存在，将使用config-template.json模板")
        config_path = "./config-template.json"

    config_str = read_file(config_path)
    logger.debug("[INIT] config str: {}".format(config_str))

    # 将json字符串反序列化为dict类型
    config = Config(json.loads(config_str))

    # override config with environment variables.
    # Some online deployment platforms (e.g. Railway) deploy project from github directly. So you shouldn't put your secrets like api key in a config file, instead use environment variables to override the default config.
    for name, value in os.environ.items():
        name = name.lower()
        if name in available_setting:
            logger.info("[INIT] override config by environ args: {}={}".format(name, value))
            try:
                config[name] = eval(value)
            except:
                if value == "false":
                    config[name] = False
                elif value == "true":
                    config[name] = True
                else:
                    config[name] = value

    if config.get("debug", False):
        logger.setLevel(logging.DEBUG)
        logger.debug("[INIT] set log level to DEBUG")

    logger.info("[INIT] load config: {}".format(config))

    # config.load_user_datas()
    #
    # logger.info("[INIT] load user datas: {}".format(config.get_user_data("api_base")))


def set_config(name, value):
    global config
    name = name.lower()
    if name in available_setting:
        logger.info("[INIT] override config by environ args: {}={}".format(name, value))
        try:
            config[name] = eval(value)
        except:
            if value == "false":
                config[name] = False
            elif value == "true":
                config[name] = True
            else:
                config[name] = value


def get_root():
    return os.path.dirname(os.path.abspath(__file__))


def read_file(path):
    with open(path, mode="r", encoding="utf-8") as f:
        return f.read()


def conf():
    return config


def get_appdata_dir():
    data_path = os.path.join(get_root(), conf().get("appdata_dir", ""))
    if not os.path.exists(data_path):
        logger.info("[INIT] data path not exists, create it: {}".format(data_path))
        os.makedirs(data_path)
    return data_path


def subscribe_msg():
    trigger_prefix = conf().get("single_chat_prefix", [""])[0]
    msg = conf().get("subscribe_msg", "")
    return msg.format(trigger_prefix=trigger_prefix)


# global plugin config
plugin_config = {}


def write_plugin_config(pconf: dict):
    """
    写入插件全局配置
    :param pconf: 全量插件配置
    """
    global plugin_config
    for k in pconf:
        plugin_config[k.lower()] = pconf[k]


def pconf(plugin_name: str) -> dict:
    """
    根据插件名称获取配置
    :param plugin_name: 插件名称
    :return: 该插件的配置项
    """
    return plugin_config.get(plugin_name.lower())


# 全局配置，用于存放全局生效的状态
global_config = {
    "admin_users": []
}