Update app.py

app.py

@@ -1,19 +1,30 @@
 import gradio as gr
 import anthropic
-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
+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
 
 def sort_key(key):
     _key = key.strip()
     if _key.startswith("sk-ant-"):
-        return get_key_ant_info(_key)
+        return get_key_ant_info(_key)        
+    elif _key.startswith("sk-"):    
+        return get_key_oai_info(_key)        
     elif _key.startswith("AIzaSy"):
-        return get_key_gemini_info(_key)
+        return get_key_gemini_info(_key)        
+    elif len(_key.split(':')) == 2 and _key.split(':')[1].islower() and len(_key.split(':')[1]) == 32:
+        endpoint = _key.split(':')[0]
+        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(_key)
-    elif _key.startswith("AKIA"):
-        return get_key_aws_info(_key)
+        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_aws_info(_key)        
+    elif len(_key) == 32:
+        return get_key_mistral_info(_key)        
     else:
-        return get_key_oai_info(_key)
+        return not_supported(_key)
+
 
 def get_key_oai_info(key):
     # Return a dictionary containing key information
@@ -63,34 +74,58 @@ def get_key_gemini_info(key):
                  "models": key_avai[1]}
     return info_dict
 
-def get_key_azure_info(key):
-    key_avai = check_key_azure_availability(key)
+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]}
+    return info_dict
+
+def get_key_azure_info(endpoint, api_key):
+    key_avai = check_key_azure_availability(endpoint, api_key)
+    info_dict = {#"account_name": "",
+             "key_type": "Microsoft Azure OpenAI",
+             "key_availability": key_avai[0],
+             "gpt35_availability": "",
+             "gpt4_availability": "",
+             "gpt4_turbo_availability": "",
+             "gpt4_32k_availability": "",             
+             "moderation_response": "",
+             "deployments": "",
+             "models": ""}
     if key_avai[0]:        
-        azure_deploy = get_azure_deploy(key)
-        status = get_azure_status(key, azure_deploy)
-        info_dict = {#"account_name": "",
-                     "key_type": "Microsoft Azure OpenAI",
-                     "key_availability": key_avai[0],
-                     "gpt4_32k_availability": status[1],
-                     "gpt4_turbo_availability": status[2],
-                     "gpt4_availability": status[3],
-                     "gpt35_availability": status[4],
-                     "models": key_avai[1],
-                     "deployments": azure_deploy,
-                     "moderation_response": status[0]}
-    else:
-        info_dict = {#"account_name": "",
-                     "key_type": "Microsoft Azure OpenAI",
-                     "key_availability": False,
-                     "models": None,
-                     "deployments": None,
-                     "moderation_response": ""}
+        azure_deploy = get_azure_deploy(endpoint, api_key)
+        status = get_azure_status(endpoint, api_key, azure_deploy)
+        info_dict["gpt35_availability"] = status[1],
+        info_dict["gpt4_availability"] = status[2],
+        info_dict["gpt4_turbo_availability"] = status[3],
+        info_dict["gpt4_32k_availability"] = status[4],
+        info_dict["moderation_response"] = status[0],
+        info_dict["models"] = key_avai[1],
+        info_dict["deployments"] = azure_deploy
+    return info_dict
+
+def get_key_mistral_info(key):
+    key_avai = check_key_mistral_availability(key)
+    info_dict = {#"account_name": "",
+                 "key_type": "Mistral AI",
+                 "key_availability": key_avai,
+                 "has_quota": ""}
+    if key_avai:
+        info_dict['has_quota'] = check_mistral_quota(key)
     return info_dict
 
 def get_key_aws_info(key):
     info_dict = {#"account_name": "",
-                     "key_type": "Amazon AWS Claude",
-                     "status": " Soon™ "}
+                 "key_type": "Amazon AWS Claude",
+                 "status": " Soon™ "}
+    return info_dict
+  
+def not_supported(key):
+    info_dict = {#"account_name": "",
+                     "key_type": "Not supported",
+                     "status": ""}
     return info_dict
     
 def clear_inputs(text):
@@ -98,7 +133,9 @@ def clear_inputs(text):
 
 with gr.Blocks() as demo:
     gr.Markdown('''
-    # OpenAI/Anthropic/Google Gemini/Microsoft Azure API Key Status Checker
+    # API Key Status Checker
+    
+    ### Supported: OpenAI/Anthropic/Google Gemini/Microsoft Azure/Mistral AI
     
     *(Based on shaocongma, CncAnon1 and Drago key checkers)*
     
@@ -107,7 +144,7 @@ with gr.Blocks() as demo:
 
     with gr.Row():
         with gr.Column():
-            key = gr.Textbox(lines=1, max_lines=1, label="OpenAI/Anthropic/Gemini/Azure API Key")
+            key = gr.Textbox(lines=1, max_lines=1, label="API Key")
             with gr.Row():
                 clear_button = gr.Button("Clear")
                 submit_button = gr.Button("Submit", variant="primary")