feat: 新增显示当前api_key的余额 (#372)

* feat: 支持显示余额

* fix: 修复余额解析错误

* perf: 优化展示效果；优化获取使用情况的方法

* 优化显示样式

* style: 统一变量名；文字靠左对齐

* Update ChuanhuChatbot.py

* perf: 更新key时自动更新API使用情况

* chore: 微调usage_display样式

---------

Co-authored-by: mzlegion <mzhao.ouo@gmail.com>
Co-authored-by: Chuan Hu <51039745+GaiZhenbiao@users.noreply.github.com>
Co-authored-by: Keldos <i@keldos.me>

ChuanhuChatbot.py

@@ -9,6 +9,7 @@ from modules.utils import *
 from modules.presets import *
 from modules.overwrites import *
 from modules.chat_func import *
+from modules.openai_func import get_usage
 
 logging.basicConfig(
     level=logging.DEBUG,
@@ -101,6 +102,8 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
                         visible=not HIDE_MY_KEY,
                         label="API-Key",
                     )
+                    usageTxt = gr.Markdown(get_usage(my_api_key), elem_id="usage_display")
+                    usageUpdateBtn = gr.Button("🔄 更新API使用情况")
                     model_select_dropdown = gr.Dropdown(
                         label="选择模型", choices=MODELS, multiselect=False, value=MODELS[0]
                     )
@@ -260,11 +263,13 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
         fn=transfer_input, inputs=[user_input], outputs=[user_question, user_input, submitBtn, cancelBtn], show_progress=True
     )
 
-    keyTxt.submit(submit_key, keyTxt, [user_api_key, status_display])
-    keyTxt.change(submit_key, keyTxt, [user_api_key, status_display])
+    keyTxt.submit(submit_key, keyTxt, [user_api_key, status_display, usageTxt])
+    keyTxt.change(submit_key, keyTxt, [user_api_key, status_display, usageTxt])
     # Chatbot
     cancelBtn.click(cancel_outputing, [], [])
 
+    usageUpdateBtn.click(get_usage, [user_api_key], [usageTxt], show_progress=True)
+
     user_input.submit(**transfer_input_args).then(**chatgpt_predict_args).then(**end_outputing_args)
 
     submitBtn.click(**transfer_input_args).then(**chatgpt_predict_args).then(**end_outputing_args)

assets/custom.css

@@ -19,6 +19,17 @@
 #chuanhu_chatbot, #status_display {
     transition: all 0.6s;
 }
+
+/* usage_display */
+#usage_display {
+    height: 1em;
+}
+#usage_display p{
+    padding: 0 1em;
+    font-size: .85em;
+    font-family: monospace;
+    color: var(--body-text-color-subdued);
+}
 /* list */
 ol:not(.options), ul:not(.options) {
     padding-inline-start: 2em !important;

modules/openai_func.py

@@ -0,0 +1,71 @@
+import requests
+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
+import os
+
+
+def get_usage_response(openai_api_key):
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {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
+
+    # 如果有代理，使用代理发送请求，否则使用默认设置发送请求
+    """
+    暂不支持修改
+    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,
+        )
+    return response
+
+def get_usage(openai_api_key):
+    try:
+        response=get_usage_response(openai_api_key=openai_api_key)
+        print(response.json())
+        try:
+            balance = response.json().get("total_available") if response.json().get(
+                "total_available") else 0
+            total_used = response.json().get("total_used") if response.json().get(
+                "total_used") else 0
+        except Exception as e:
+            logging.error(f"API使用情况解析失败:"+str(e))
+            balance = 0
+            total_used=0 
+        return f"**API使用情况**（已用/余额）\u3000{total_used}$ / {balance}$"
+    except requests.exceptions.ConnectTimeout:
+        status_text = standard_error_msg + connection_timeout_prompt + error_retrieve_prompt
+        return status_text
+    except requests.exceptions.ReadTimeout:
+        status_text = standard_error_msg + read_timeout_prompt + error_retrieve_prompt
+        return status_text
+    

modules/presets.py

@@ -4,6 +4,7 @@ import gradio as gr
 # ChatGPT 设置
 initial_prompt = "You are a helpful assistant."
 API_URL = "https://api.openai.com/v1/chat/completions"
+BALANCE_API_URL="https://api.openai.com/dashboard/billing/credit_grants"
 HISTORY_DIR = "history"
 TEMPLATES_DIR = "templates"
 

modules/utils.py

@@ -22,6 +22,7 @@ from pygments.formatters import HtmlFormatter
 
 from modules.presets import *
 import modules.shared as shared
+from modules.openai_func import get_usage
 
 logging.basicConfig(
     level=logging.INFO,
@@ -349,7 +350,7 @@ def submit_key(key):
     key = key.strip()
     msg = f"API密钥更改为了{hide_middle_chars(key)}"
     logging.info(msg)
-    return key, msg
+    return key, msg, get_usage(key)
 
 
 def sha1sum(filename):