Alexzf01/AAAChatGPT

1from collections import defaultdict2from contextlib import contextmanager3import os4import logging5import sys6import json7 8from . import shared9 10 11__all__ = [12    "my_api_key",13    "authflag",14    "auth_list",15    "dockerflag",16    "retrieve_proxy",17    "log_level",18    "advance_docs",19    "update_doc_config",20    "multi_api_key",21]22 23# 添加一个统一的config文件，避免文件过多造成的疑惑（优先级最低）24# 同时，也可以为后续支持自定义功能提供config的帮助25if os.path.exists("config.json"):26    with open("config.json", "r", encoding='utf-8') as f:27        config = json.load(f)28else:29    config = {}30 31## 处理docker if we are running in Docker32dockerflag = config.get("dockerflag", False)33if os.environ.get("dockerrun") == "yes":34    dockerflag = True35 36## 处理 api-key 以及 允许的用户列表37my_api_key = config.get("openai_api_key", "") # 在这里输入你的 API 密钥38my_api_key = os.environ.get("my_api_key", my_api_key)39 40## 多账户机制41multi_api_key = config.get("multi_api_key", False) # 是否开启多账户机制42if multi_api_key:43    api_key_list = config.get("api_key_list", [])44    if len(api_key_list) == 0:45        logging.error("多账号模式已开启，但api_key_list为空，请检查config.json")46        sys.exit(1)47    shared.state.set_api_key_queue(api_key_list)48 49auth_list = config.get("users", []) # 实际上是使用者的列表50authflag = len(auth_list) > 0  # 是否开启认证的状态值，改为判断auth_list长度51 52# 处理自定义的api_host，优先读环境变量的配置，如果存在则自动装配53api_host = os.environ.get("api_host", config.get("api_host", ""))54if api_host:55    shared.state.set_api_host(api_host)56 57if dockerflag:58    if my_api_key == "empty":59        logging.error("Please give a api key!")60        sys.exit(1)61    # auth62    username = os.environ.get("USERNAME")63    password = os.environ.get("PASSWORD")64    if not (isinstance(username, type(None)) or isinstance(password, type(None))):65        auth_list.append((os.environ.get("USERNAME"), os.environ.get("PASSWORD")))66        authflag = True67else:68    if (69        not my_api_key70        and os.path.exists("api_key.txt")71        and os.path.getsize("api_key.txt")72    ):73        with open("api_key.txt", "r") as f:74            my_api_key = f.read().strip()75    if os.path.exists("auth.json"):76        authflag = True77        with open("auth.json", "r", encoding='utf-8') as f:78            auth = json.load(f)79            for _ in auth:80                if auth[_]["username"] and auth[_]["password"]:81                    auth_list.append((auth[_]["username"], auth[_]["password"]))82                else:83                    logging.error("请检查auth.json文件中的用户名和密码！")84                    sys.exit(1)85 86@contextmanager87def retrieve_openai_api(api_key = None):88    old_api_key = os.environ.get("OPENAI_API_KEY", "")89    if api_key is None:90        os.environ["OPENAI_API_KEY"] = my_api_key91        yield my_api_key92    else:93        os.environ["OPENAI_API_KEY"] = api_key94        yield api_key95    os.environ["OPENAI_API_KEY"] = old_api_key96 97## 处理log98log_level = config.get("log_level", "INFO")99logging.basicConfig(100    level=log_level,101    format="%(asctime)s [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s",102)103 104## 处理代理：105http_proxy = config.get("http_proxy", "")106https_proxy = config.get("https_proxy", "")107http_proxy = os.environ.get("HTTP_PROXY", http_proxy)108https_proxy = os.environ.get("HTTPS_PROXY", https_proxy)109 110# 重置系统变量，在不需要设置的时候不设置环境变量，以免引起全局代理报错111os.environ["HTTP_PROXY"] = ""112os.environ["HTTPS_PROXY"] = ""113 114@contextmanager115def retrieve_proxy(proxy=None):116    """117    1, 如果proxy = NONE，设置环境变量，并返回最新设置的代理118    2，如果proxy ！= NONE，更新当前的代理配置，但是不更新环境变量119    """120    global http_proxy, https_proxy121    if proxy is not None:122        http_proxy = proxy123        https_proxy = proxy124        yield http_proxy, https_proxy125    else:126        old_var = os.environ["HTTP_PROXY"], os.environ["HTTPS_PROXY"]127        os.environ["HTTP_PROXY"] = http_proxy128        os.environ["HTTPS_PROXY"] = https_proxy129        yield http_proxy, https_proxy # return new proxy130 131        # return old proxy132        os.environ["HTTP_PROXY"], os.environ["HTTPS_PROXY"] = old_var133 134 135## 处理advance docs136advance_docs = defaultdict(lambda: defaultdict(dict))137advance_docs.update(config.get("advance_docs", {}))138def update_doc_config(two_column_pdf):139    global advance_docs140    if two_column_pdf:141        advance_docs["pdf"]["two_column"] = True142    else:143        advance_docs["pdf"]["two_column"] = False144 145    logging.info(f"更新后的文件参数为：{advance_docs}")