refactor: 精简request.post调用、更改余额部分默认文字

geektown · Mar 29, 2023 · 5817e2c · 5817e2c
1 parent 4aae38f
commit 5817e2c
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 58 deletions.
diff --git a/ChuanhuChatbot.py b/ChuanhuChatbot.py
@@ -108,7 +108,7 @@
                         visible=not HIDE_MY_KEY,
                         label="API-Key",
                     )
-                    usageTxt = gr.Markdown(get_usage(my_api_key), elem_id="usage_display")
+                    usageTxt = gr.Markdown("**发送消息** 或 **提交key** 以显示余额", elem_id="usage_display")
                     model_select_dropdown = gr.Dropdown(
                         label="选择模型", choices=MODELS, multiselect=False, value=MODELS[0]
                     )

diff --git a/modules/chat_func.py b/modules/chat_func.py
@@ -58,39 +58,21 @@ def get_response(
     else:
         timeout = timeout_all
 
-    # 获取环境变量中的代理设置
-    http_proxy = os.environ.get("HTTP_PROXY") or os.environ.get("http_proxy")
-    https_proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("https_proxy")
-
-    # 如果存在代理设置，使用它们
-    proxies = {}
-    if http_proxy:
-        logging.info(f"使用 HTTP 代理: {http_proxy}")
-        proxies["http"] = http_proxy
-    if https_proxy:
-        logging.info(f"使用 HTTPS 代理: {https_proxy}")
-        proxies["https"] = https_proxy
+    proxies = get_proxies()
 
     # 如果有自定义的api-url，使用自定义url发送请求，否则使用默认设置发送请求
     if shared.state.api_url != API_URL:
         logging.info(f"使用自定义API URL: {shared.state.api_url}")
-    if proxies:
-        response = requests.post(
-            shared.state.api_url,
-            headers=headers,
-            json=payload,
-            stream=True,
-            timeout=timeout,
-            proxies=proxies,
-        )
-    else:
-        response = requests.post(
-            shared.state.api_url,
-            headers=headers,
-            json=payload,
-            stream=True,
-            timeout=timeout,
-        )
+
+    response = requests.post(
+        shared.state.api_url,
+        headers=headers,
+        json=payload,
+        stream=True,
+        timeout=timeout,
+        proxies=proxies,
+    )
+
     return response
 
 

diff --git a/modules/openai_func.py b/modules/openai_func.py
@@ -2,6 +2,7 @@
 import logging
 from modules.presets import timeout_all, BALANCE_API_URL,standard_error_msg,connection_timeout_prompt,error_retrieve_prompt,read_timeout_prompt
 from modules import shared
+from modules.utils import get_proxies
 import os
 
 
@@ -13,39 +14,19 @@ def get_usage_response(openai_api_key):
 
     timeout = timeout_all
 
-    # 获取环境变量中的代理设置
-    http_proxy = os.environ.get("HTTP_PROXY") or os.environ.get("http_proxy")
-    https_proxy = os.environ.get(
-        "HTTPS_PROXY") or os.environ.get("https_proxy")
-
-    # 如果存在代理设置，使用它们
-    proxies = {}
-    if http_proxy:
-        logging.info(f"使用 HTTP 代理: {http_proxy}")
-        proxies["http"] = http_proxy
-    if https_proxy:
-        logging.info(f"使用 HTTPS 代理: {https_proxy}")
-        proxies["https"] = https_proxy
-
-    # 如果有代理，使用代理发送请求，否则使用默认设置发送请求
+    proxies = get_proxies()
     """
     暂不支持修改
     if shared.state.balance_api_url != BALANCE_API_URL:
         logging.info(f"使用自定义BALANCE API URL: {shared.state.balance_api_url}")
     """
-    if proxies:
-        response = requests.get(
-            BALANCE_API_URL,
-            headers=headers,
-            timeout=timeout,
-            proxies=proxies,
-        )
-    else:
-        response = requests.get(
-            BALANCE_API_URL,
-            headers=headers,
-            timeout=timeout,
-        )
+    response = requests.get(
+        BALANCE_API_URL,
+        headers=headers,
+        timeout=timeout,
+        proxies=proxies,
+    )
+
     return response
 
 def get_usage(openai_api_key):
@@ -68,3 +49,6 @@ def get_usage(openai_api_key):
     except requests.exceptions.ReadTimeout:
         status_text = standard_error_msg + read_timeout_prompt + error_retrieve_prompt
         return status_text
+    except Exception as e:
+        logging.error(f"获取API使用情况失败:"+str(e))
+        return standard_error_msg + error_retrieve_prompt
diff --git a/modules/utils.py b/modules/utils.py
@@ -437,3 +437,23 @@ def transfer_input(inputs):
         gr.Button.update(visible=False),
         gr.Button.update(visible=True),
     )
+
+
+def get_proxies():
+    # 获取环境变量中的代理设置
+    http_proxy = os.environ.get("HTTP_PROXY") or os.environ.get("http_proxy")
+    https_proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("https_proxy")
+
+    # 如果存在代理设置，使用它们
+    proxies = {}
+    if http_proxy:
+        logging.info(f"使用 HTTP 代理: {http_proxy}")
+        proxies["http"] = http_proxy
+    if https_proxy:
+        logging.info(f"使用 HTTPS 代理: {https_proxy}")
+        proxies["https"] = https_proxy
+
+    if proxies == {}:
+        proxies = None
+
+    return proxies