Update api_usage.py

api_usage.py

@@ -4,34 +4,70 @@ from datetime import datetime
 from dateutil.relativedelta import relativedelta
 import openai
 
+GPT_TYPES = ["gpt-3.5-turbo", "gpt-4", "gpt-4-32k"]
+
+rate_limit_per_model = {
+    "gpt-3.5-turbo": 3500,
+    "gpt-4": 200,
+    "gpt-4-32k": 10  # No actual clue, rare enough
+}
 
 def get_headers(key):
     headers = {'Authorization': f'Bearer {key}'}
     return headers
 
 def get_subscription(key):
-    queryUrl = 'https://api.openai.com/dashboard/billing/subscription'
-    headers = get_headers(key)
-    r = requests.get(queryUrl, headers=headers)
-    results = r.json()
+    queryUrl = 'https://api.openai.com/v1/chat/completions'
+    headers = get_headers(key)    
+    #results = r.json()
+    body_turbo = {"model": "gpt-3.5-turbo", "max_tokens": 1, "messages": [{'role':'user', 'content': ''}]}
+    body_gpt4 = {"model": "gpt-4", "max_tokens": 1, "messages": [{'role':'user', 'content': ''}]}
+
     if check_key_availability():
-        has_payment_method = results["has_payment_method"]
+        rpm = ""
+        org = ""
+        quota = ""
+        try:
+            r = requests.post(queryUrl, headers=headers, json=body_gpt4 if check_gpt4_availability() else body_turbo)
+            rpm = r.headers['x-ratelimit-limit-requests']
+            org = r.headers['openai-organization']
+            quota = ""
+        except Exception as e:
+            error_message = str(e)
+            if "You exceeded your current quota" in error_message:
+                rpm = "";
+                org = "";
+                quota = "No quota left."
+            elif "Your account is not active" in error_message:
+                rpm = "";
+                org = "";
+                quota = "Error: Your account is not active, please check your billing details on our website."
+            else:
+                rpm = "";
+                org = "";
+                quota = f"Unexpected Error at check_key: {error_message}"
+        
+        #has_payment_method = results["has_payment_method"]
         # hard_limit = results["hard_limit"]
-        soft_limit_usd = results["soft_limit_usd"]
-        hard_limit_usd = results["hard_limit_usd"]
-        plan_title = results["plan"]["title"]
-        plan_id = results["plan"]["id"]
-        account_name = results["account_name"]
-        return {"account_name": account_name,
-                "has_payment_method": has_payment_method,
-                "soft_limit_usd": soft_limit_usd,
-                "hard_limit_usd": hard_limit_usd,
-                "plan": plan_title + ", " + plan_id}
+        #soft_limit_usd = results["soft_limit_usd"]
+        #hard_limit_usd = results["hard_limit_usd"]
+        #plan_title = results["plan"]["title"]
+        #plan_id = results["plan"]["id"]
+        #account_name = results["account_name"]
+        return {"organization": org,
+                "rpm": rpm,
+                "quota": quota}
+                #"has_payment_method": has_payment_method,
+                #"soft_limit_usd": soft_limit_usd,
+                #"hard_limit_usd": hard_limit_usd,
+                #"plan": plan_title + ", " + plan_id}               
     else:
-        return {"account_name": "",
-                "has_payment_method": False,
-                "hard_limit_usd": "",
-                "plan": ""}
+        return {"organization": "",
+                "rpm": "",
+                "quota": "Invalid key or key has been revoked!"}
+                #"has_payment_method": False,
+                #"hard_limit_usd": "",
+                #"plan": ""}
 
 #def get_usage(key):
 #    if check_key_availability():
@@ -54,6 +90,16 @@ def check_gpt4_availability():
     else:
         return False
 
+def check_gpt4_32k_availability():
+    if check_key_availability():
+        available_models = [model["root"] for model in openai.Model.list()["data"]]
+        if 'gpt-4-32k' in available_models:
+            return True
+        else:
+            return False
+    else:
+        return False
+
 def check_key_availability():
     try:
         openai.Model.list()