File size: 2,161 Bytes
e99bd71
5e4ca56
e4df10e
6f97306
e4df10e
2c5812c
 
e4df10e
e99bd71
5e4ca56
 
e99bd71
2c5812c
 
 
 
 
 
 
f650935
6f97306
f650935
 
 
 
e99bd71
 
f650935
 
6f97306
2c5812c
5e4ca56
e99bd71
5e4ca56
 
6f97306
5e4ca56
2c5812c
 
 
e99bd71
f650935
e4df10e
 
 
 
 
 
 
 
 
017dec1
e4df10e
 
017dec1
 
e4df10e
 
 
 
f650935
2c5812c
 
1a5edd8
 
 
9186425
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
from modules.presets import CHAT_COMPLETION_URL, BALANCE_API_URL, USAGE_API_URL, API_HOST, OPENAI_API_BASE
import os
import queue
import openai

class State:
    interrupted = False
    multi_api_key = False
    chat_completion_url = CHAT_COMPLETION_URL
    balance_api_url = BALANCE_API_URL
    usage_api_url = USAGE_API_URL
    openai_api_base = OPENAI_API_BASE

    def interrupt(self):
        self.interrupted = True

    def recover(self):
        self.interrupted = False

    def set_api_host(self, api_host: str):
        api_host = api_host.rstrip("/")
        if not api_host.startswith("http"):
            api_host = f"https://{api_host}"
        if api_host.endswith("/v1"):
            api_host = api_host[:-3]
        self.chat_completion_url = f"{api_host}/v1/chat/completions"
        self.openai_api_base = f"{api_host}/v1"
        self.balance_api_url = f"{api_host}/dashboard/billing/credit_grants"
        self.usage_api_url = f"{api_host}/dashboard/billing/usage"
        os.environ["OPENAI_API_BASE"] = api_host

    def reset_api_host(self):
        self.chat_completion_url = CHAT_COMPLETION_URL
        self.balance_api_url = BALANCE_API_URL
        self.usage_api_url = USAGE_API_URL
        os.environ["OPENAI_API_BASE"] = f"https://{API_HOST}"
        return API_HOST

    def reset_all(self):
        self.interrupted = False
        self.chat_completion_url = CHAT_COMPLETION_URL

    def set_api_key_queue(self, api_key_list):
        self.multi_api_key = True
        self.api_key_queue = queue.Queue()
        for api_key in api_key_list:
            self.api_key_queue.put(api_key)

    def switching_api_key(self, func):
        if not hasattr(self, "api_key_queue"):
            return func

        def wrapped(*args, **kwargs):
            api_key = self.api_key_queue.get()
            args[0].api_key = api_key
            ret = func(*args, **kwargs)
            self.api_key_queue.put(api_key)
            return ret

        return wrapped


state = State()

modules_path = os.path.dirname(os.path.realpath(__file__))
chuanhu_path = os.path.dirname(modules_path)
assets_path = os.path.join(chuanhu_path, "web_assets")