Alexzf01/AAAChatGPT

1from modules.presets import COMPLETION_URL, BALANCE_API_URL, USAGE_API_URL, API_HOST2import os3import queue4 5class State:6    interrupted = False7    multi_api_key = False8    completion_url = COMPLETION_URL9    balance_api_url = BALANCE_API_URL10    usage_api_url = USAGE_API_URL11 12    def interrupt(self):13        self.interrupted = True14 15    def recover(self):16        self.interrupted = False17 18    def set_api_host(self, api_host):19        self.completion_url = f"https://{api_host}/v1/chat/completions"20        self.balance_api_url = f"https://{api_host}/dashboard/billing/credit_grants"21        self.usage_api_url = f"https://{api_host}/dashboard/billing/usage"22        os.environ["OPENAI_API_BASE"] = f"https://{api_host}/v1"23 24    def reset_api_host(self):25        self.completion_url = COMPLETION_URL26        self.balance_api_url = BALANCE_API_URL27        self.usage_api_url = USAGE_API_URL28        os.environ["OPENAI_API_BASE"] = f"https://{API_HOST}/v1"29        return API_HOST30 31    def reset_all(self):32        self.interrupted = False33        self.completion_url = COMPLETION_URL34    35    def set_api_key_queue(self, api_key_list):36        self.multi_api_key = True37        self.api_key_queue = queue.Queue()38        for api_key in api_key_list:39            self.api_key_queue.put(api_key)40 41    def switching_api_key(self, func):42        if not hasattr(self, "api_key_queue"):43            return func44        45        def wrapped(*args, **kwargs):46            api_key = self.api_key_queue.get()47            args = list(args)[1:]48            ret = func(api_key, *args, **kwargs)49            self.api_key_queue.put(api_key)50            return ret51 52        return wrapped53        54 55state = State()56