Update app.py

app.py

@@ -1,15 +1,15 @@
 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, check_key_or_availability, check_key_or_limits
+from api_usage import get_subscription, check_key_availability, check_key_ant_availability, check_ant_rate_limit, 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):
+async def sort_key(key, rate_limit):
     _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)
+        return await get_key_ant_info(_key, rate_limit)
         
     if _key.startswith("sk-"):    
         return get_key_oai_info(_key)
@@ -68,25 +68,29 @@ def get_key_oai_info(key):
         info_dict["quota"] = info["quota"]
     return info_dict
 
-def get_key_ant_info(key):
+async def get_key_ant_info(key, rate_limit):
     # Return a dictionary containing key information
     key_avai = check_key_ant_availability(key)
     info_dict = {#"account_name": "",
                  "key_type": "Anthropic Claude",
                  "key_availability": key_avai[0],
-                 "status": key_avai[1],
-                 "filter_response": key_avai[2],
-                 "requests_per_minute": key_avai[3] + ("" if key_avai[3] == "" else f" ({key_avai[4]} left)"),
-                 "tokens_per_minute": key_avai[5] + ("" if key_avai[5] == "" else f" ({key_avai[6]} left)"),
-                 "tier": key_avai[7]}
-    return info_dict
-
-def get_key_gemini_info(key):
-    key_avai = check_key_gemini_availability(key)
-    info_dict = {#"account_name": "",
-                 "key_type": "Google Gemini",
-                 "key_availability": key_avai[0],
-                 "models": key_avai[1]}
+                 "status": "",
+                 "filter_response": "",
+                 "requests_per_minute": "",
+                 "tokens_per_minute": "",
+                 "tier": "",
+                 "concurrent_rate_limit": ""}
+    
+    info_dict["status"] = key_avai[1]
+    info_dict["filter_response"] = key_avai[2]
+    info_dict["requests_per_minute"] = key_avai[3] + ("" if key_avai[3] == "" else f" ({key_avai[4]} left)")
+    info_dict["tokens_per_minute"] = key_avai[5] + ("" if key_avai[5] == "" else f" ({key_avai[6]} left)")
+    info_dict["tier"] = key_avai[7]
+        
+    if rate_limit:
+        rate = await check_ant_rate_limit(key)
+        info_dict["concurrent_rate_limit"] = rate
+        
     return info_dict
 
 def get_key_gemini_info(key):
@@ -225,6 +229,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             key = gr.Textbox(lines=1, max_lines=1, label="API Key")
+            rate_limit = gr.Checkbox(label="Check concurrent rate limit (API Claude) (experimental)")
             with gr.Row():
                 clear_button = gr.Button("Clear")
                 submit_button = gr.Button("Submit", variant="primary")
@@ -232,5 +237,5 @@ with gr.Blocks() as demo:
             info = gr.JSON(label="API Key Information")
 
     clear_button.click(fn=clear_inputs, inputs=[key], outputs=[key])
-    submit_button.click(fn=sort_key, inputs=[key], outputs=[info], api_name="sort_key")
+    submit_button.click(fn=sort_key, inputs=[key, rate_limit], outputs=[info], api_name="sort_key")
 demo.launch()