Update api_usage.py

api_usage.py

@@ -4,13 +4,12 @@ from datetime import datetime
 from dateutil.relativedelta import relativedelta
 import openai
 
-queryUrl = 'https://api.openai.com/v1/chat/completions'
+baseUrl = 'https://api.openai.com/v1'
 GPT_TYPES = ["gpt-3.5-turbo", "gpt-4", "gpt-4-32k"]
 rate_limit_per_model = {
-    "gpt-3.5-turbo-trial": 2000,
-    "gpt-3.5-turbo-pay": 3500,
+    "gpt-3.5-turbo": 2000, # new pay turbo will have 2000 RPM for the first 48 hours then become 3500
     "gpt-4": 200,
-    "gpt-4-32k": 1000  # No actual clue, rare enough
+    "gpt-4-32k": 1000
 }
 body_gpt = {
     "gpt-3.5-turbo": {"model": "gpt-3.5-turbo", "max_tokens": 1, "messages": [{'role':'user', 'content': ''}]},
@@ -18,95 +17,83 @@ body_gpt = {
     "gpt-4-32k": {"model": "gpt-4-32k", "max_tokens": 1, "messages": [{'role':'user', 'content': ''}]}
 }
 
-
 def get_headers(key):
     headers = {'Authorization': f'Bearer {key}'}
     return headers
 
 def get_subscription(key):    
-    headers = get_headers(key)    
-    #results = r.json()
-    
-    key_highest_model = ""
-    if check_gpt4_32k_availability():
-        key_highest_model = GPT_TYPES[2]
-    elif check_gpt4_availability():
-        key_highest_model = GPT_TYPES[1]
-    else:
-        key_highest_model = GPT_TYPES[0]    
-
+    headers = get_headers(key)
+      
     if check_key_availability():
-        rpm = ""
+        rpm = "0"        
+        tpm = "0"
         org = ""
         quota = ""
-        r = requests.post(queryUrl, headers=headers, json=body_gpt[key_highest_model])
+        has_gpt4_32k = False
+        has_gpt4 = False
+        available_models = [model["root"] for model in openai.Model.list()["data"]]
+        key_highest_model = ""
+        
+        if check_gpt4_32k_availability(available_models):
+            key_highest_model = GPT_TYPES[2]
+            has_gpt4_32k = True
+            has_gpt4 = True
+        elif check_gpt4_availability(available_models):
+            key_highest_model = GPT_TYPES[1]
+            has_gpt4 = True
+        else:
+            key_highest_model = GPT_TYPES[0]  
+            
+        r = requests.post(f"{baseUrl}/chat/completions", headers=headers, json=body_gpt[key_highest_model])
         result = r.json()
-        if "id" in result:
-            rpm = r.headers['x-ratelimit-limit-requests']
-            org = r.headers['openai-organization']
-            quota = check_key_type(key_highest_model, int(rpm))         
+        
+        if "id" in result:            
+            rpm = r.headers.get("x-ratelimit-limit-requests", "0")
+            tpm = r.headers.get("x-ratelimit-limit-tokens", "0")
+            org = r.headers.get('openai-organization', "")
+            quota = check_key_type(key_highest_model, int(rpm))      
         else:
-            e = result["error"]["code"]
+            e = result.get("error", {}).get("code", "")
             quota = f"Error: {e}"
+            org = get_org_name(key)
             
-        return {"organization": org,
-                "rpm": rpm,
+        return {"has_gpt4_32k": has_gpt4_32k,
+                "has_gpt4": has_gpt4,
+                "organization": org,
+                "rpm": f"{rpm} ({key_highest_model})",
+                "tpm": f"{tpm}",
                 "quota": quota}           
     else:
-        return {"organization": "",
+        return {"has_gpt4_32k": has_gpt4_32k,
+                "has_gpt4": has_gpt4,
+                "organization": "",
                 "rpm": "",
+                "tpm": "",
                 "quota": ""}
-                #"has_payment_method": False,
-                #"hard_limit_usd": "",
-                #"plan": ""}
 
-#def get_usage(key):
-#    if check_key_availability():
-#        start_date = datetime.now().strftime('%Y-%m-01')
-#        end_date = (datetime.now() + relativedelta(months=1)).strftime('%Y-%m-01')
-#        queryUrl = f'https://api.openai.com/dashboard/billing/usage?start_date={start_date}&end_date={end_date}'
-#        headers = get_headers(key)
-#        r = requests.get(queryUrl, headers=headers)
-#        return r.json()
-#    else:
-#        return ""
+def get_org_name(key):
+    headers=get_headers(key)
+    r = requests.post(f"{baseUrl}/images/generations", headers=headers)
+    return r.headers['openai-organization']
 
 def check_key_type(model, rpm):
-    if model == GPT_TYPES[0]:
-        if rpm > rate_limit_per_model['gpt-3.5-turbo-pay']:
-            return "yes | pay, possibly big key"
-        elif rpm > rate_limit_per_model['gpt-3.5-turbo-trial'] and rpm <= rate_limit_per_model['gpt-3.5-turbo-pay']:
-            return "yes | pay"
-        else:
-            return "yes | trial"
+    if rpm >= rate_limit_per_model[model]:
+        return "yes | pay"
     else:
-        if rpm < rate_limit_per_model[model]:
-            return "yes | trial"
-        elif rpm == rate_limit_per_model[model]:
-            return "yes | pay"
-        else:
-            return "yes | pay, possibly big key"
+        return "yes | trial"
     
-def check_gpt4_availability():
-    if check_key_availability():
-        available_models = [model["root"] for model in openai.Model.list()["data"]]
-        if 'gpt-4' in available_models:
-            return True
-        else:
-            return False
+def check_gpt4_availability(available_models):
+    if 'gpt-4' in available_models:
+        return True
     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
+def check_gpt4_32k_availability(available_models):
+    if 'gpt-4-32k' in available_models:
+        return True
     else:
         return False
-
+        
 def check_key_availability():
     try:
         openai.Model.list()
@@ -116,8 +103,6 @@ def check_key_availability():
 
 if __name__ == "__main__":
     key = os.getenv("OPENAI_API_KEY")
-#    results = get_usage(key)
-#    print(results)
 
     results = get_subscription(key)
     for k, v in results.items():