File size: 2,374 Bytes
458ae42
55f3e5a
 
1f45019
3e8b0ee
 
 
55f3e5a
79eae90
55f3e5a
 
79eae90
458ae42
3e8b0ee
 
 
 
 
 
 
1f45019
 
 
 
 
 
79eae90
458ae42
79eae90
1f45019
 
458ae42
3e8b0ee
55f3e5a
79eae90
458ae42
55f3e5a
 
1f45019
55f3e5a
3e8b0ee
 
 
79eae90
1f45019
55f3e5a
 
 
 
 
 
 
 
 
7b5a1c0
55f3e5a
 
7b5a1c0
 
55f3e5a
 
 
 
1f45019
3e8b0ee
 
8971a40
 
 
fcf8580
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
68
69
70
from modules.presets import CHAT_COMPLETION_URL, BALANCE_API_URL, USAGE_API_URL, API_HOST, OPENAI_API_BASE, IMAGES_COMPLETION_URL
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
    images_completion_url = IMAGES_COMPLETION_URL

    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.images_completion_url = f"{api_host}/v1/images/generations"
        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 + "/v1"

    def reset_api_host(self):
        self.chat_completion_url = CHAT_COMPLETION_URL
        self.images_completion_url = IMAGES_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")