Update api_usage.py

api_usage.py

@@ -13,7 +13,8 @@ import asyncio, aiohttp
 import aiohttp
 
 BASE_URL = 'https://api.openai.com/v1'
-GPT_TYPES = ["gpt-3.5-turbo", "gpt-4", "gpt-4-32k", "gpt-4-32k-0314", "gpt-4o", "gpt-4-turbo"]
+GPT_TYPES = ["gpt-3.5-turbo", "gpt-4", "gpt-4-32k", "gpt-4-32k-0314", "gpt-4o", "gpt-4-turbo", "chatgpt-4o-latest" ]
+GPT_TYPES_ORDER = ["gpt-4-32k", "gpt-4-32k-0314", "chatgpt-4o-latest", "gpt-4o", "gpt-4-turbo", "gpt-4", "gpt-3.5-turbo"]
 
 TOKEN_LIMIT_PER_TIER_TURBO = {
     "free": 40000,
@@ -42,199 +43,12 @@ RPM_LIMIT_PER_BUILD_TIER_ANT = {
 } # https://docs.anthropic.com/claude/reference/rate-limits
 
 
-def get_headers(key, org_id:str = None):
+def get_headers(key, org_id=None):
     headers = {'Authorization': f'Bearer {key}'}
     if org_id:
         headers["OpenAI-Organization"] = org_id
     return headers
 
-def get_subscription(key, session, org_list):
-    has_gpt4 = False
-    has_gpt4_32k = False
-    has_gpt4_32k_0314 = False
-    default_org = ""
-    org_description = []
-    org = []
-    rpm = []
-    tpm = []    
-    quota = []
-    list_models = []
-    list_models_avai = set()
-    
-    for org_in in org_list:
-        if len(org_list) < 2: # mismatch_organization
-            headers = get_headers(key)
-            available_models = get_models(session, key)
-        else:
-            headers = get_headers(key, org_in['id'])
-            available_models = get_models(session, key, org_in['id'])
-        if org_in['id']:
-            if org_in['is_default']:
-                default_org = org_in['name']            
-            org_description.append(f"{org_in['description']} (Created: {datetime.utcfromtimestamp(org_in['created'])} UTC" + (", personal)" if org_in['personal'] else ")"))
-        if 'No perm' in available_models:
-            available_models.extend(GPT_TYPES)
-        has_gpt4_32k = True if GPT_TYPES[2] in available_models else False
-        has_gpt4_32k_0314 = True if GPT_TYPES[3] in available_models else False
-        has_gpt4 = True if GPT_TYPES[1] in available_models else False
-        
-        if has_gpt4_32k_0314 or has_gpt4_32k:
-            if org_in['id']:
-                org.append(f"{org_in['id']} ({org_in['name']}, {org_in['title']}, {org_in['role']})")
-            if has_gpt4_32k:
-                list_models_avai.update(GPT_TYPES)
-                if 'No perm' in available_models:
-                    status_formated = format_status(GPT_TYPES, session, headers)
-                else:
-                    status_formated = format_status([GPT_TYPES[2], GPT_TYPES[4], GPT_TYPES[5], GPT_TYPES[1], GPT_TYPES[0]], session, headers)
-                rpm.append(status_formated[0])
-                tpm.append(status_formated[1])
-                quota.append(status_formated[2])
-                if 'No perm' in available_models:
-                    lst_string = ""
-                    length = len(status_formated[3])
-                    count = 1
-                    for k, v in status_formated[3].items():
-                        if v:
-                            if count < length:
-                                lst_string += f'{k}, '
-                                continue
-                            else:
-                                lst_string += f' {k} '
-                        if v == False:
-                            list_models_avai.remove(k)
-                            if k == GPT_TYPES[2]:
-                                has_gpt4_32k = False
-                            elif k == GPT_TYPES[1]:
-                                has_gpt4 = False
-                            elif k == GPT_TYPES[0]:
-                                has_35 = False
-                            elif k == GPT_TYPES[4]:
-                                has_4o = False
-                        count += 1
-                    lst_string += '(No get model permission)'
-                    #list_models.append(f"gpt-4-32k, gpt-4o, gpt-4-turbo, gpt-4, gpt-3.5-turbo (No get model permission)")
-                    list_models.append(lst_string)
-                else:
-                    list_models.append(f"gpt-4-32k, gpt-4o, gpt-4-turbo, gpt-4, gpt-3.5-turbo ({len(available_models)} total)")
-            else:
-                list_models_avai.update([GPT_TYPES[3], GPT_TYPES[1], GPT_TYPES[0]])
-                status_formated = format_status([GPT_TYPES[3], GPT_TYPES[4], GPT_TYPES[5], GPT_TYPES[1], GPT_TYPES[0]], session, headers)
-                rpm.append(status_formated[0])
-                tpm.append(status_formated[1])
-                quota.append(status_formated[2])
-                list_models.append(f"gpt-4-32k-0314, gpt-4o, gpt-4-turbo, gpt-4, gpt-3.5-turbo ({len(available_models)} total)")                
-            
-        elif has_gpt4:
-            if org_in['id']:
-                org.append(f"{org_in['id']} ({org_in['name']}, {org_in['title']}, {org_in['role']})")
-            list_models_avai.update([GPT_TYPES[1], GPT_TYPES[0]])
-            status_formated = format_status([GPT_TYPES[4], GPT_TYPES[5], GPT_TYPES[1], GPT_TYPES[0]], session, headers)
-            rpm.append(status_formated[0])
-            tpm.append(status_formated[1])
-            quota.append(status_formated[2])
-            list_models.append(f"gpt-4o, gpt-4-turbo, gpt-4, gpt-3.5-turbo ({len(available_models)} total)")
-            
-        else:
-            if org_in['id']:
-                org.append(f"{org_in['id']} ({org_in['name']}, {org_in['title']}, {org_in['role']})")
-            list_models_avai.update([GPT_TYPES[0]])
-            status_formated = format_status([GPT_TYPES[0]], session, headers)
-            rpm.append(status_formated[0])
-            tpm.append(status_formated[1])
-            quota.append(status_formated[2])
-            list_models.append(f"gpt-3.5-turbo ({len(available_models)} total)")
-    
-    return {"has_gpt4_32k": True if GPT_TYPES[2] in list_models_avai else False,
-            "has_gpt4": True if GPT_TYPES[1] in list_models_avai else False,
-            "default_org": default_org,
-            "organization": [o for o in org],
-            "org_description": org_description,
-            "models": list_models,
-            "rpm": rpm,
-            "tpm": tpm,
-            "quota": quota}
-          
-def send_oai_completions(oai_stuff):
-    session = oai_stuff[0]
-    headers = oai_stuff[1]
-    model = oai_stuff[2]
-    model_status = False
-    try:
-        req_body = {"model": model, "max_tokens": 1}
-        rpm_string = ""
-        tpm_string = ""
-        quota_string = ""
-        r = session.post(f"{BASE_URL}/chat/completions", headers=headers, json=req_body, timeout=10)
-        result = r.json()
-        if "error" in result:
-            e = result.get("error", {}).get("code", "")
-            if e == None or e == 'missing_required_parameter':
-                rpm_num = int(r.headers.get("x-ratelimit-limit-requests", 0))
-                model_status = True
-                tpm_num = int(r.headers.get('x-ratelimit-limit-tokens', 0))
-                tpm_left = int(r.headers.get('x-ratelimit-remaining-tokens', 0))
-                _rpm = '{:,}'.format(rpm_num).replace(',', ' ')
-                _tpm = '{:,}'.format(tpm_num).replace(',', ' ')
-                _tpm_left = '{:,}'.format(tpm_left).replace(',', ' ')
-                rpm_string = f"{_rpm} ({model})"
-                #tpm_string = f"{_tpm} ({_tpm_left} left, {model})"
-                tpm_string = f"{_tpm} ({model})"
-                dictCount = 0
-                dictLength = len(TOKEN_LIMIT_PER_TIER_GPT4)
-                            
-                # Check if gpt-4 has custom tpm (600k for example), if not, proceed with 3turbo's tpm
-                if model == GPT_TYPES[1]:
-                    for k, v in TOKEN_LIMIT_PER_TIER_GPT4.items():
-                        if tpm_num == v:
-                            break
-                        else:
-                            dictCount+=1                      
-                            if dictCount == dictLength:
-                                quota_string = "yes | custom-tier"
-                elif model == GPT_TYPES[0] and quota_string == "":
-                    quota_string = check_key_tier(rpm_num, tpm_num, TOKEN_LIMIT_PER_TIER_TURBO, headers)
-            else:
-                rpm_string = f"0 ({model})"
-                tpm_string = f"0 ({model})"
-                quota_string = e
-        return rpm_string, tpm_string, quota_string, model, model_status
-    except Exception as e:
-        #print(e)
-        return "", "", "", model, model_status
-        
-def format_status(list_models_avai, session, headers):
-    rpm = []
-    tpm = []
-    model_status = {}
-    quota = ""
-    args = [(session, headers, model) for model in list_models_avai]
-    with concurrent.futures.ThreadPoolExecutor() as executer:
-        for result in executer.map(send_oai_completions, args):
-            rpm.append(result[0])
-            tpm.append(result[1])
-            model_status[result[3]] = result[4]
-            if result[2]:
-                if quota == 'yes | custom-tier':
-                    continue
-                else:
-                    quota = result[2]
-    rpm_str = ""
-    tpm_str = ""
-    for i in range(len(rpm)):                
-        rpm_str += rpm[i] + (" | " if i < len(rpm)-1 else "")
-        tpm_str += tpm[i] + (" | " if i < len(rpm)-1 else "")
-    return rpm_str, tpm_str, quota, model_status
-
-def check_key_tier(rpm, tpm, dict, headers):
-    dictItemsCount = len(dict)
-    dictCount = 0
-    for k, v in dict.items():        
-        if tpm == v:
-            return f"yes | {k}"
-        dictCount+=1
-        if (dictCount == dictItemsCount):
-            return "yes | custom-tier"
 
 def get_orgs(session, key):
     headers=get_headers(key)
@@ -242,33 +56,205 @@ def get_orgs(session, key):
         rq = session.get(f"{BASE_URL}/organizations", headers=headers, timeout=10)
         return 200, rq.json()['data']
     except:
-        if rq.status_code == 403:
+        if rq.status_code == 403 or (rq.status_code == 401 and 'insufficient permissions' in f'{rq.text}'):
             return 403, rq.json()['error']['message']
         else:
             return False, False
-            
-def get_models(session, key, org: str = None):
-    if org != None:
-        headers = get_headers(key, org)
-    else:
-        headers = get_headers(key)
-        
+
+def get_orgs_me(session, key):
+    headers=get_headers(key)
     try:
-        rq = session.get(f"{BASE_URL}/models", headers=headers, timeout=10)
-        avai_models = rq.json()
-        if rq.status_code == 403:
-            return ['No perm']
-        list_models = [model["id"] for model in avai_models["data"]] #[model["id"] for model in avai_models["data"] if model["id"] in GPT_TYPES]
+        rq = session.get(f"{BASE_URL}/me", headers=headers, timeout=10)
+        if rq.status_code == 200:
+            return 200, rq.json()['orgs']['data']
+        else:
+            return rq.status_code, ""
     except:
-        list_models = []
-    return list_models
-        
+        return False, rq.json()['error']['message']
+
+def get_models(session, key, org_id=None):
+    headers = get_headers(key, org_id)
+    try:
+        response = session.get(f"{BASE_URL}/models", headers=headers, timeout=10)
+        response.raise_for_status()
+        all_models = response.json().get("data", [])
+        available_models = [model["id"] for model in all_models]
+        return available_models, all_models
+    except requests.exceptions.RequestException:
+        if response.status_code == 403:
+            return ['No perm'], []
+        return [], []
+
+
 def check_key_availability(session, key):
+    return get_orgs(session, key)
+
+def get_subscription(key, session, org_data):
+    default_org = ""
+    org_description = []
+    organizations = []
+    rpm_list = []
+    tpm_list = []
+    quota_list = []
+    models_list = []
+    int_models = []
+    all_models = []
+    has_gpt4 = False
+    has_gpt4_32k = False
+
+    if isinstance(org_data, str):
+        organizations.append(org_data)
+        org_list = [{"id": ""}]
+    else:
+        org_list = org_data
+    #print(org_list)
+    for org in org_list:
+        org_id = org.get('id')
+        headers = get_headers(key, org_id)
+        all_available_models_info = get_models(session, key, org_id)
+        all_available_models = all_available_models_info[0]
+        available_models = [model for model in all_available_models if model in GPT_TYPES]
+        if 'No perm' in available_models:
+            available_models = GPT_TYPES
+
+        if org.get('is_default'):
+            default_org = org.get('name', '')
+
+        if org_id:
+            organizations.append(f"{org_id} ({org.get('name')}, {org.get('title')}, {org.get('role')})")
+            org_description.append(f"{org.get('description')} (Created: {datetime.utcfromtimestamp(org.get('created'))} UTC"
+                                    + (", personal)" if org.get('personal') else ")"))
+
+        status = format_status(available_models, session, headers)
+        rpm_list.append(" | ".join(status["rpm"]))
+        tpm_list.append(" | ".join(status["tpm"]))
+
+        quota = determine_quota_for_org(status)
+        quota_list.append(quota)
+
+        models_list.append(", ".join(status["models"]) + f" ({len(all_available_models)} total)")
+        all_models.append(all_available_models)
+
+        has_gpt4 = has_gpt4 or GPT_TYPES[1] in available_models
+        has_gpt4_32k = has_gpt4_32k or GPT_TYPES[2] in available_models
+
+    return {
+        "has_gpt4": has_gpt4,
+        "has_gpt4_32k": has_gpt4_32k,
+        "default_org": default_org,
+        "organization": organizations,
+        "org_description": org_description,
+        "models": models_list,
+        "rpm": rpm_list,
+        "tpm": tpm_list,
+        "quota": quota_list,
+        "all_models": all_models
+    }
+
+
+def determine_quota_for_org(status):
+    for model, quota in zip(status["models"], status["quota"]):
+        if model == "gpt-4":
+            if quota == "custom-tier":
+                return quota
+
+    for model, quota in zip(status["models"], status["quota"]):
+        if model == "gpt-3.5-turbo":
+            return quota
+
+    for quota in status["quota"]:
+        if quota not in ["unknown", "custom-tier"]:
+            return quota
+
+    return "unknown"
+
+def format_status(models, session, headers):
+    rpm = []
+    tpm = []
+    quota = []
+    model_status = {}
+
+    args = [(session, headers, model) for model in models]
+
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        results = executor.map(lambda x: send_oai_completions(*x), args)
+
+    model_results = {result["model"]: result for result in results}
+
+    sorted_models = [model for model in GPT_TYPES_ORDER if model in model_results]
+
+    for model in sorted_models:
+        result = model_results[model]
+        rpm.append(result["rpm"])
+        tpm.append(result["tpm"])
+        quota.append(result["quota"])
+        model_status[model] = result["status"]
+
+    return {
+        "rpm": rpm,
+        "tpm": tpm,
+        "quota": quota,
+        "models": sorted_models,
+        "model_status": model_status
+    }
+
+def send_oai_completions(session, headers, model):
     try:
-        orgs = get_orgs(session, key)
-        return orgs
-    except Exception as e:
-        return False, False
+        response = session.post(
+            f"{BASE_URL}/chat/completions",
+            headers=headers,
+            json={"model": model, "max_tokens": 1},
+            timeout=10
+        )
+
+        result = response.json()
+        rpm = int(response.headers.get("x-ratelimit-limit-requests", 0))
+        tpm = int(response.headers.get("x-ratelimit-limit-tokens", 0))
+        quota_string = ""
+
+        if "error" in result:
+            error_code = result.get("error", {}).get("code", "")
+            if error_code in [None, "missing_required_parameter"]:
+                quota_string = check_tier(tpm, TOKEN_LIMIT_PER_TIER_GPT4 if model == "gpt-4" else TOKEN_LIMIT_PER_TIER_TURBO)
+                return {
+                    "rpm": f"{rpm:,} ({model})",
+                    "tpm": f"{tpm:,} ({model})",
+                    "quota": quota_string,
+                    "model": model,
+                    "status": True
+                }
+            else:
+                return {
+                    "rpm": f"0 ({model})",
+                    "tpm": f"0 ({model})",
+                    "quota": error_code,
+                    "model": model,
+                    "status": False
+                }
+
+        quota_string = check_tier(tpm, TOKEN_LIMIT_PER_TIER_GPT4 if model == "gpt-4" else TOKEN_LIMIT_PER_TIER_TURBO)
+        return {
+            "rpm": f"{rpm:,} ({model})",
+            "tpm": f"{tpm:,} ({model})",
+            "quota": quota_string,
+            "model": model,
+            "status": True
+        }
+
+    except requests.exceptions.RequestException:
+        return {
+            "rpm": "0",
+            "tpm": "0",
+            "quota": "request_failed",
+            "model": model,
+            "status": False
+        }
+
+def check_tier(tpm, token_limits):
+    for tier, limit in token_limits.items():
+        if tpm == limit:
+            return tier
+    return "custom-tier"
 
 async def fetch_ant(async_session, json_data):
     url = 'https://api.anthropic.com/v1/messages'
@@ -347,8 +333,17 @@ async def check_key_ant_availability(key, claude_model):
     tpm = ""
     tpm_left = ""
     tier = ""
-    
+    tpm_input = ""
+    tpm_input_left = ""
+    tpm_output = ""
+    tpm_output_left = ""
+    models = ""
     async with aiohttp.ClientSession(headers=headers) as async_session:
+        async with async_session.get(url='https://api.anthropic.com/v1/models') as m_response:
+            if m_response.status == 200:
+                models_info = await m_response.json()
+                models = [model['id'] for model in models_info['data']]
+        
         async with async_session.post(url=url, json=json_data) as response:
             result = await response.json()
             if response.status == 200:
@@ -356,15 +351,19 @@ async def check_key_ant_availability(key, claude_model):
                 rpm_left = response.headers.get('anthropic-ratelimit-requests-remaining', '')
                 tpm = response.headers.get('anthropic-ratelimit-tokens-limit', '')
                 tpm_left = response.headers.get('anthropic-ratelimit-tokens-remaining', '')
+                tpm_input = response.headers.get('anthropic-ratelimit-input-tokens-limit', '')
+                tpm_input_left = response.headers.get('anthropic-ratelimit-input-tokens-remaining', '')
+                tpm_output = response.headers.get('anthropic-ratelimit-output-tokens-limit', '')
+                tpm_output_left = response.headers.get('anthropic-ratelimit-output-tokens-remaining', '')
                 tier = check_ant_tier(rpm)
                 msg = result.get('content', [''])[0].get('text', '')
-                return True, "Working", msg, rpm, rpm_left, tpm, tpm_left, tier
+                return True, "Working", msg, rpm, rpm_left, tpm, tpm_left, tier, tpm_input, tpm_input_left, tpm_output, tpm_output_left, models
             else:
                 #err_type = result.get('error', '').get('type', '')
                 err_msg = result.get('error', '').get('message', '')
                 if response.status == 401:
-                    return False, f'Error: {response.status}', err_msg, rpm, rpm_left, tpm, tpm_left, tier
-                return True, f'Error: {response.status}', err_msg, rpm, rpm_left, tpm, tpm_left, tier
+                    return False, f'Error: {response.status}', err_msg, rpm, rpm_left, tpm, tpm_left, tier, tpm_input, tpm_input_left, tpm_output, tpm_output_left, models
+                return True, f'Error: {response.status}', err_msg, rpm, rpm_left, tpm, tpm_left, tier, tpm_input, tpm_input_left, tpm_output, tpm_output_left, models
 
 def check_key_gemini_availability(key):
     avai = False
@@ -642,7 +641,7 @@ def is_model_working(form_info, model_info):
             return model_info['agreementAvailability']['errorMessage']
         return "No"
     except:
-        return "No"
+        return "Maybe"
         
 async def get_model_status(session, key, secret, region, model_name, form_info):
     model_info = await bedrock_model_available(session, key, secret, region, f"anthropic.{model_name}")
@@ -686,7 +685,7 @@ async def check_bedrock_claude_status(session, key, secret):
         if region and model_name:
             if msg == "Maybe":
                 invoke_info = await send_signed_request_bedrock(session, payload, f"anthropic.{model_name}", key, secret, region)
-                if 'messages.0' in invoke_info.get('message') or 'many requests' in invoke_info.get('message'):
+                if 'messages.0' in invoke_info.get('message') or 'many requests' in invoke_info.get('message') or 'equal to 1' in invoke_info.get('message'):
                     models[model_name].append(f'{region}')
             else:
                 models[model_name].append(region)
@@ -883,6 +882,76 @@ def check_elevenlabs_status(key):
     else:
         return False, user_info[1], ""
 
+def check_xai_status(key):
+    url = 'https://api.x.ai/v1/'
+    headers = {"authorization": f"Bearer {key}"}
+    response = requests.get(url+"api-key", headers=headers)
+    if response.status_code == 200:
+        response_models = requests.get(url+"models", headers=headers)
+        models = response_models.json()
+        if 'data' in models:
+            return True, response.json(), models['data']
+        return True, response.json(), models
+    else:
+        return False, response.json(), ""
+        
+def check_stripe_status(key):
+    response = requests.get('https://api.stripe.com/v1/balance', auth=(key, ''))
+    
+    if response.status_code == 200:
+        return True, response.json()
+    else:
+        return False, response.json()
+
+def check_stability_info(url, headers):
+    response = requests.get(url=url, headers=headers)
+    if response.status_code == 200:
+        return True, response.json()
+    else:
+        return False, response.json()
+
+def check_stability_status(key):
+    headers = {"Authorization": f"Bearer {key}"}
+    
+    status, account = check_stability_info('https://api.stability.ai/v1/user/account', headers)
+    
+    if 'Incorrect API key' in f'{account}':
+        return False, account, "", ""
+    
+    _, models = check_stability_info('https://api.stability.ai/v1/engines/list', headers)
+    _, credit = check_stability_info('https://api.stability.ai/v1/user/balance', headers)
+    
+    return True, account, models, credit
+
+def check_deepseek_balance(key):
+    url = 'https://api.deepseek.com/user/balance'
+    headers = {"Authorization": f"Bearer {key}"}
+    response = requests.get(url, headers=headers)
+    
+    if response.status_code == 200:
+        return response.json()
+    else:
+        return ""
+
+def check_deepseek_models(key):
+    url = 'https://api.deepseek.com/models'
+    headers = {"Authorization": f"Bearer {key}"}
+    response = requests.get(url, headers=headers)
+    
+    if response.status_code == 200:
+        return True, response.json()['data']
+    else:
+        return False, response.json()
+
+def check_deepseek_status(key):
+    status, models = check_deepseek_models(key)
+    if not status and 'no such user' in f'{models}':
+        return False, models, ""
+    
+    balance_info = check_deepseek_balance(key)
+    
+    return True, models, balance_info
+
 if __name__ == "__main__":
     key = os.getenv("OPENAI_API_KEY")
     key_ant = os.getenv("ANTHROPIC_API_KEY")