Add API dollar usage and fix style 增加当月API美金使用量 (#456)

* add dollar usage and fix style

* conditional rednering of usage

ChuanhuChatbot.py

@@ -9,7 +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
+from modules.openai_func import get_usage, get_dollar_usage_for_current_month
 
 logging.basicConfig(
     level=logging.DEBUG,
@@ -108,7 +108,7 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
                         visible=not HIDE_MY_KEY,
                         label="API-Key",
                     )
-                    usageTxt = gr.Markdown("**发送消息** 或 **提交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]
                     )
@@ -272,14 +272,20 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
         fn=get_usage, inputs=[user_api_key], outputs=[usageTxt], show_progress=False
     )
 
+    get_dollar_usage_args = dict(
+        fn=get_dollar_usage_for_current_month, inputs=[user_api_key], outputs=[usageTxt], show_progress=False
+    )
+
     # Chatbot
     cancelBtn.click(cancel_outputing, [], [])
 
     user_input.submit(**transfer_input_args).then(**chatgpt_predict_args).then(**end_outputing_args)
     user_input.submit(**get_usage_args)
+    user_input.submit(**get_dollar_usage_args)
 
     submitBtn.click(**transfer_input_args).then(**chatgpt_predict_args).then(**end_outputing_args)
     submitBtn.click(**get_usage_args)
+    submitBtn.click(**get_dollar_usage_args)
 
     emptyBtn.click(
         reset_state,
@@ -306,6 +312,7 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
         show_progress=True,
     ).then(**end_outputing_args)
     retryBtn.click(**get_usage_args)
+    retryBtn.click(**get_dollar_usage_args)
 
     delFirstBtn.click(
         delete_first_conversation,
@@ -338,9 +345,13 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
         show_progress=True,
     )
     reduceTokenBtn.click(**get_usage_args)
+    reduceTokenBtn.click(**get_dollar_usage_args)
 
     # ChatGPT
-    keyTxt.change(submit_key, keyTxt, [user_api_key, status_display]).then(**get_usage_args)
+    if get_usage(user_api_key) == "**您的免费额度已用完**":
+        keyTxt.change(submit_key, keyTxt, [user_api_key, status_display]).then(**get_dollar_usage_args)
+    else:
+        keyTxt.change(submit_key, keyTxt, [user_api_key, status_display]).then(**get_usage_args)
 
     # Template
     templateRefreshBtn.click(get_template_names, None, [templateFileSelectDropdown])

assets/custom.css

@@ -37,8 +37,9 @@ footer {
 
 /* usage_display */
 #usage_display {
-    height: 1em;
+    height: 2em;
 }
+
 #usage_display p{
     padding: 0 1em;
     font-size: .85em;

modules/openai_func.py

@@ -1,10 +1,19 @@
 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.presets import (
+    timeout_all,
+    USAGE_API_URL,
+    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
-
+from datetime import datetime, timedelta
 
 def get_usage_response(openai_api_key):
     headers = {
@@ -42,7 +51,52 @@ def get_usage(openai_api_key):
             logging.error(f"API使用情况解析失败:"+str(e))
             balance = 0
             total_used=0
-        return f"**API使用情况**（已用/余额）\u3000{total_used}$ / {balance}$"
+            return f"**API使用情况解析失败**"
+        if balance == 0:
+            return f"**您的免费额度已用完**"
+        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
+    except Exception as e:
+        logging.error(f"获取API使用情况失败:"+str(e))
+        return standard_error_msg + error_retrieve_prompt
+
+def get_usage_for_current_month_response(openai_api_key):
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {openai_api_key}",
+    }
+
+    timeout = timeout_all
+    today = datetime.today()
+    first_day_of_month = today.replace(day=1).date()
+    last_day_of_month = get_last_day_of_month(today).date()
+    
+    proxies = get_proxies()
+    response = requests.get(
+        f"{USAGE_API_URL}?start_date={first_day_of_month}&end_date={last_day_of_month}",
+        headers=headers,
+        timeout=timeout,
+        proxies=proxies,
+    )
+        
+    return response
+
+def get_dollar_usage_for_current_month(openai_api_key):
+    try:
+        response=get_usage_for_current_month_response(openai_api_key=openai_api_key)
+        usage = 0
+        upper_limit = "120" # hardcoded as 120 dollars for now
+        try:
+            usage = round(response.json().get("total_usage")/100, 2) if response.json().get(
+                "total_usage") else 0
+        except Exception as e:
+            logging.error(f"API使用情况解析失败:"+str(e))
+        return f"**本月API使用情况**（已用/限额）\u3000${usage} / ${upper_limit}"
     except requests.exceptions.ConnectTimeout:
         status_text = standard_error_msg + connection_timeout_prompt + error_retrieve_prompt
         return status_text
@@ -52,3 +106,9 @@ def get_usage(openai_api_key):
     except Exception as e:
         logging.error(f"获取API使用情况失败:"+str(e))
         return standard_error_msg + error_retrieve_prompt
+    
+def get_last_day_of_month(any_day):
+    # The day 28 exists in every month. 4 days later, it's always next month
+    next_month = any_day.replace(day=28) + timedelta(days=4)
+    # subtracting the number of the current day brings us back one month
+    return next_month - timedelta(days=next_month.day)

modules/presets.py

@@ -5,6 +5,7 @@ import gradio as gr
 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"
+USAGE_API_URL="https://api.openai.com/dashboard/billing/usage"
 HISTORY_DIR = "history"
 TEMPLATES_DIR = "templates"