Update app.py

app.py

@@ -2,14 +2,14 @@ import gradio as gr
 import requests
 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
 
-async def sort_key(key, rate_limit):
+async def sort_key(key, rate_limit, claude_opus):
     _key = key.strip()
     
     if _key.startswith("sk-or-v1-"):    
         return get_key_openrouter_info(_key)
         
     if _key.startswith("sk-ant-"):
-        return await get_key_ant_info(_key, rate_limit)
+        return await get_key_ant_info(_key, rate_limit, claude_opus)
         
     if _key.startswith("sk-"):    
         return get_key_oai_info(_key)
@@ -44,7 +44,7 @@ def get_key_oai_info(key):
     key_avai = check_key_availability(session, key)
     info_dict = {#"account_name": "",
                  "key_type": "OpenAI",
-                 "key_availability": True if key_avai else False,
+                 "key_availability": True if key_avai[0] else False,
                  "gpt4_availability": "",
                  "gpt4_32k_availability": "",
                  "default_org": "",                 
@@ -55,22 +55,27 @@ def get_key_oai_info(key):
                  "tokens_per_minute": "",
                  #"tokens_per_minute_left": "",
                  "quota": ""}
-    if key_avai:
-        info = get_subscription(key, session, key_avai)
+    if key_avai[0]:
+        if key_avai[0] == 200:            
+            info = get_subscription(key, session, key_avai[1])
+            info_dict["default_org"] = info["default_org"]        
+            info_dict["org_description"] = info["org_description"]
+            info_dict["organization"] = info["organization"]   
+        else:
+            empty_org = { 'id': '' }
+            info = get_subscription(key, session, [empty_org])
+            info_dict["organization"] = key_avai[1]
         info_dict["gpt4_availability"] = info["has_gpt4"]
         info_dict["gpt4_32k_availability"] = info["has_gpt4_32k"]
-        info_dict["default_org"] = info["default_org"]        
-        info_dict["org_description"] = info["org_description"]
-        info_dict["organization"] = info["organization"]      
         info_dict["models"] = info["models"]
         info_dict["requests_per_minute"] = info["rpm"]
         info_dict["tokens_per_minute"] = info["tpm"]
         info_dict["quota"] = info["quota"]
     return info_dict
 
-async def get_key_ant_info(key, rate_limit):
+async def get_key_ant_info(key, rate_limit, claude_opus):
     # Return a dictionary containing key information
-    key_avai = check_key_ant_availability(key)
+    key_avai = check_key_ant_availability(key, claude_opus)
     info_dict = {#"account_name": "",
                  "key_type": "Anthropic Claude",
                  "key_availability": key_avai[0],
@@ -236,6 +241,7 @@ with gr.Blocks() as demo:
         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)")
+            claude_opus = gr.Checkbox(label="Use claude-3-opus for filter_response check? (default: claude-3-haiku)")
             with gr.Row():
                 clear_button = gr.Button("Clear")
                 submit_button = gr.Button("Submit", variant="primary")
@@ -243,6 +249,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, rate_limit], outputs=[info], api_name="sort_key")
-demo.launch()
-1
+    submit_button.click(fn=sort_key, inputs=[key, rate_limit, claude_opus], outputs=[info], api_name="sort_key")
+demo.launch()