From 5bbe3f12d6729de2d6479823258e8dbcfc2326c4 Mon Sep 17 00:00:00 2001
From: Soulter <905617992@qq.com>
Date: Mon, 25 Sep 2023 13:25:38 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20OpenAI=E5=AE=98=E6=96=B9=E6=A8=A1?=
 =?UTF-8?q?=E5=9E=8B=E6=94=AF=E6=8C=81=E5=88=87=E6=8D=A2=E8=B4=A6=E5=8F=B7?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 model/command/command_openai_official.py   | 34 +++++++++++++++++++++-
 model/provider/provider_openai_official.py |  2 ++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/model/command/command_openai_official.py b/model/command/command_openai_official.py
index bc2fc5b6a..02c96888a 100644
--- a/model/command/command_openai_official.py
+++ b/model/command/command_openai_official.py
@@ -61,6 +61,8 @@ class CommandOpenAIOfficial(Command):
             return True, self.draw(message)
         elif self.command_start_with(message, "key"):
             return True, self.key(message)
+        elif self.command_start_with(message, "switch"):
+            return True, self.switch(message)
         
         if self.command_start_with(message, "/"):
             return True, (False, "未知指令", "unknown_command")
@@ -137,7 +139,7 @@ class CommandOpenAIOfficial(Command):
                 continue
             if 'sponsor' in key_stat[key]:
                 sponsor = key_stat[key]['sponsor']
-            chatgpt_cfg_str += f"  |-{index}: {key_stat[key]['used']}/{max} {sponsor}赞助{tag}\n"
+            chatgpt_cfg_str += f"  |-{index}: {key[-8:]} {key_stat[key]['used']}/{max} {sponsor}{tag}\n"
             index += 1
         return True, f"⭐使用情况({str(gg_count)}个已用):\n{chatgpt_cfg_str}", "status"
 
@@ -161,6 +163,36 @@ class CommandOpenAIOfficial(Command):
         else:
             return True, "该Key被验证为无效。也许是输入错误了，或者重试。", "key"
 
+    def switch(self, message: str):
+        '''
+        切换账号
+        '''
+        l = message.split(" ")
+        if len(l) == 1:
+            _, ret, _ = self.status()
+            curr_ = self.provider.get_curr_key()
+            if curr_ is None:
+                ret += "当前您未选择账号。输入/switch <账号序号>切换账号。"
+            else:
+                ret += f"当前您选择的账号为：{curr_[-8:]}。输入/switch <账号序号>切换账号。"
+            return True, ret, "switch"
+        elif len(l) == 2:
+            try:
+                key_stat = self.provider.get_key_stat()
+                index = int(l[1])
+                if index > len(key_stat) or index < 1:
+                    return True, "账号序号不合法。", "switch"
+                else:
+                    ret = self.provider.check_key(list(key_stat.keys())[index-1])
+                    if ret:
+                        return True, f"账号切换成功。", "switch"
+                    else:
+                        return True, f"账号切换失败，可能超额或超频。", "switch"
+            except BaseException as e:
+                return True, "未知错误: "+str(e), "switch"
+        else:
+            return True, "参数过多。", "switch"
+
     def unset(self, session_id: str):
         if self.provider is None:
             return False, "未启动OpenAI ChatGPT语言模型.", "unset"
diff --git a/model/provider/provider_openai_official.py b/model/provider/provider_openai_official.py
index 12e2d65b4..2f920769d 100644
--- a/model/provider/provider_openai_official.py
+++ b/model/provider/provider_openai_official.py
@@ -355,6 +355,8 @@ class ProviderOpenAIOfficial(Provider):
         return self.key_stat
     def get_key_list(self):
         return self.key_list
+    def get_curr_key(self):
+        return openai.api_key
     
     # 添加key
     def append_key(self, key, sponsor):