Update app.py

app.py

@@ -1,31 +1,42 @@
 import gradio as gr
 import requests
-from api_usage import get_subscription, check_key_availability, check_key_ant_availability, check_key_gemini_availability, check_key_azure_availability, get_azure_status, get_azure_deploy, check_key_mistral_availability, check_mistral_quota, check_key_replicate_availability, check_key_aws_availability
+from api_usage import get_subscription, check_key_availability, check_key_ant_availability, check_key_gemini_availability, check_key_azure_availability, get_azure_status, get_azure_deploy, check_key_mistral_availability, check_mistral_quota, check_key_replicate_availability, check_key_aws_availability, check_key_or_availability, check_key_or_limits
 
 def sort_key(key):
     _key = key.strip()
+    
+    if _key.startswith("sk-or-v1-"):    
+        return get_key_openrouter_info(_key)
+        
     if _key.startswith("sk-ant-"):
-        return get_key_ant_info(_key)        
-    elif _key.startswith("sk-"):    
-        return get_key_oai_info(_key)        
-    elif _key.startswith("AIzaSy"):
+        return get_key_ant_info(_key)
+        
+    if _key.startswith("sk-"):    
+        return get_key_oai_info(_key)
+        
+    if _key.startswith("AIzaSy"):
         return get_key_gemini_info(_key)
-    elif (_key.startswith("r8_") and len(_key) == 40) or (_key.islower() and len(_key) == 40):
+        
+    if (_key.startswith("r8_") and len(_key) == 40) or (_key.islower() and len(_key) == 40):
         return get_key_replicate_info(_key)
-    elif len(_key.split(':')) == 2 and _key.split(':')[1].islower() and len(_key.split(':')[1]) == 32 and "openai.azure.com" not in _key.split(':')[1]:
+        
+    if len(_key.split(':')) == 2 and _key.split(':')[1].islower() and len(_key.split(':')[1]) == 32 and "openai.azure.com" not in _key.split(':')[1]:
         endpoint = f"{_key.split(':')[0]}.openai.azure.com"
         api_key = _key.split(':')[1]
-        return get_key_azure_info(endpoint, api_key)        
-    elif "openai.azure.com" in _key.split(';')[0]:
+        return get_key_azure_info(endpoint, api_key)
+        
+    if "openai.azure.com" in _key.split(';')[0]:
         endpoint = _key.split(';')[0]
         api_key = _key.split(';')[1]
-        return get_key_azure_info(endpoint, api_key)        
-    elif _key.startswith("AKIA") and len(_key.split(':')[0]) == 20 and _key.split(':')[0].isupper():
+        return get_key_azure_info(endpoint, api_key)
+        
+    if _key.startswith("AKIA") and len(_key.split(':')[0]) == 20 and _key.split(':')[0].isupper():
         return get_key_aws_info(_key)
-    elif len(_key) == 32:
-        return get_key_mistral_info(_key)    
-    else:
-        return not_supported(_key)
+        
+    if len(_key) == 32:
+        return get_key_mistral_info(_key)
+        
+    return not_supported(_key)
 
 def get_key_oai_info(key):
     # Return a dictionary containing key information
@@ -157,11 +168,39 @@ def get_key_aws_info(key):
         info_dict['admin'] = key_avai[3]
         info_dict['quarantine'] = key_avai[4] 
         info_dict['iam_full_access'] = key_avai[5]
-        info_dict['billing'] = key_avai[8]
         info_dict['claude_v2_enabled_region'] = key_avai[6]
         info_dict['claude_sonnet_enabled_region'] = key_avai[7]
         info_dict['cost_and_usage'] = key_avai[8]        
     return info_dict
+
+def get_key_openrouter_info(key):    
+    key_avai = check_key_or_availability(key)
+    info_dict = {#"account_name": "",
+                 "key_type": "OpenRouter",
+                 "key_availability": key_avai[0],
+                 "is_free_tier": "",
+                 "usage": "",
+                 "balance": "",
+                 "limit": "",
+                 "limit_remaining": "",
+                 "rate_limit_per_minite": "",
+                 "4_turbo_per_request_tokens_limit": "",
+                 "sonnet_per_request_tokens_limit": "",
+                 "opus_per_request_tokens_limit": ""}
+    if key_avai[0]:
+        models_info = check_key_or_limits(key)
+        info_dict['is_free_tier'] = key_avai[1]['is_free_tier']
+        info_dict['limit'] = key_avai[1]['limit']
+        info_dict['limit_remaining'] = key_avai[1]['limit_remaining']
+        info_dict['usage'] = f"${format(key_avai[1]['usage'], '.4f')}"
+        info_dict['balance'] = f"${format(models_info[0], '.4f')}"
+        info_dict['rate_limit_per_minite'] = key_avai[2]
+        info_dict['4_turbo_per_request_tokens_limit'] = models_info[1]['openai/gpt-4-turbo-preview']
+        info_dict['sonnet_per_request_tokens_limit'] = models_info[1]['anthropic/claude-3-sonnet:beta']
+        info_dict['opus_per_request_tokens_limit'] = models_info[1]['anthropic/claude-3-opus:beta']
+    else:
+        info_dict['usage'] = key_avai[1]
+    return info_dict
   
 def not_supported(key):
     info_dict = {#"account_name": "",
@@ -174,7 +213,7 @@ def clear_inputs(text):
 
 with gr.Blocks() as demo:
     gr.Markdown('''
-    # OpenAI/Anthropic/Gemini/Azure/Mistral/Replicate/AWS Claude API Key Status Checker
+    # OpenAI/Anthropic/Gemini/Azure/Mistral/Replicate/AWS Claude/OpenRouter API Key Status Checker
     
     *(Based on shaocongma, CncAnon1, Drago and kingbased key checkers)*