Update app.py

app.py

@@ -4,15 +4,15 @@ import requests
 import re
 from typing import List, Dict, Tuple, Any # 添加类型提示
 
-# 假设api_usage模块根据用户环境存在
-# 导入需要的函数（实际使用时需要此模块）
+# api_usage 模块假定存在于适合用户的环境中.
+# 导入所需的函数（实际使用时需要相应模块）
 from api_usage import (
     get_subscription,
     check_key_availability,
     get_orgs_me,
     check_key_ant_availability,
     check_ant_rate_limit,
-    check_key_gemini_availability, # 假设此函数返回(bool, str)元组
+    check_key_gemini_availability,
     check_key_azure_availability,
     get_azure_status,
     get_azure_deploy,
@@ -32,15 +32,16 @@ from api_usage import (
 )
 
 # ─────────────────────────────────────────
-# 特定密钥的帮助函数（保持原样 - 包括Gemini）
+# Key-specific helper functions (OpenAI 函数已修改)
 # ─────────────────────────────────────────
 def get_key_oai_info(key: str) -> Dict[str, Any]:
     session = requests.Session()
-    status, org_data = check_key_availability(session, key)
+    raw_status_code, org_data_or_error = check_key_availability(session, key)
 
     info_dict = {
         "key_type": "OpenAI",
-        "key_availability": True if status else False,
+        "key_availability": False,
+        "has_sufficient_quota": True, # 默认设为True，如有问题则改为False
         "gpt4_availability": "",
         "gpt4_32k_availability": "",
         "default_org": "",
@@ -51,17 +52,68 @@ def get_key_oai_info(key: str) -> Dict[str, Any]:
         "tokens_per_minute": "",
         "quota": "",
         "all_models": "",
+        "status_message": ""
     }
 
-    if not status:
+    org_data_for_subscription = None
+
+    if raw_status_code == 200:
+        info_dict["key_availability"] = True
+        org_data_for_subscription = org_data_or_error
+    elif raw_status_code == 401: # Unauthorized
+        info_dict["status_message"] = "Unauthorized: Invalid API key."
+        info_dict["has_sufficient_quota"] = False
+        return info_dict
+    elif raw_status_code == 403: # Forbidden
+        status_me, orgs_me_data = get_orgs_me(session, key)
+        if status_me == 200:
+            info_dict["key_availability"] = True
+            org_data_for_subscription = orgs_me_data
+        else:
+            info_dict["status_message"] = f"Forbidden, and get_orgs_me failed (status: {status_me}). Key might be inactive or lack permissions."
+            info_dict["has_sufficient_quota"] = False
+            return info_dict
+    elif raw_status_code == 429: # Too Many Requests (Rate Limit or Quota)
+        info_dict["key_availability"] = True # 密钥本身可能是有效的
+        info_dict["has_sufficient_quota"] = False # 被视为配额不足
+        info_dict["status_message"] = "Rate limit or quota likely exceeded (initial check)."
+        if isinstance(org_data_or_error, dict) and "error" in org_data_or_error:
+            error_details = org_data_or_error["error"]
+            current_quota_message = error_details.get("message", "Quota details unavailable from initial check")
+            info_dict["quota"] = current_quota_message
+            if "insufficient_quota" not in current_quota_message.lower(): # 如果详细消息中没有insufficient_quota，则在状态消息中添加
+                 info_dict["status_message"] += f" Error: {current_quota_message}"
+    else: # 기타 에러
+        info_dict["status_message"] = f"Key check failed (status: {raw_status_code})."
+        if isinstance(org_data_or_error, dict) and "error" in org_data_or_error:
+            info_dict["status_message"] += f" Error: {org_data_or_error['error'].get('message', str(org_data_or_error))}"
+        elif isinstance(org_data_or_error, str):
+             info_dict["status_message"] += f" Details: {org_data_or_error}"
+        info_dict["has_sufficient_quota"] = False
         return info_dict
 
-    if status == 403:
-        status_me, orgs_me = get_orgs_me(session, key)
+    if not info_dict["key_availability"]:
+        # 并非由于429错误的其他原因 key_availability가 False면 반환
+        if raw_status_code != 429:
+             return info_dict
+
+    if not org_data_for_subscription and info_dict["key_availability"]:
+        status_me, orgs_me_data = get_orgs_me(session, key)
         if status_me == 200:
-            org_data = orgs_me
+            org_data_for_subscription = orgs_me_data
+        else:
+            info_dict["status_message"] = (info_dict["status_message"] + " Could not identify organization for subscription.").strip()
+            info_dict["key_availability"] = False
+            info_dict["has_sufficient_quota"] = False
+            return info_dict
+
+    if not org_data_for_subscription and info_dict["key_availability"]:
+        info_dict["status_message"] = (info_dict["status_message"] + " Organization data for subscription is missing.").strip()
+        info_dict["key_availability"] = False
+        info_dict["has_sufficient_quota"] = False
+        return info_dict
 
-    subscription_info = get_subscription(key, session, org_data)
+    subscription_info = get_subscription(key, session, org_data_for_subscription)
     if subscription_info:
         info_dict.update(
             {
@@ -73,10 +125,51 @@ def get_key_oai_info(key: str) -> Dict[str, Any]:
                 "models": subscription_info.get("models", ""),
                 "requests_per_minute": subscription_info.get("rpm", ""),
                 "tokens_per_minute": subscription_info.get("tpm", ""),
-                "quota": subscription_info.get("quota", ""),
+                "quota": subscription_info.get("quota") if subscription_info.get("quota") is not None else info_dict.get("quota", ""),
                 "all_models": subscription_info.get("all_models", ""),
             }
         )
+
+        error_info = subscription_info.get("error")
+        if error_info and isinstance(error_info, dict):
+            err_type = error_info.get("type", "").lower()
+            err_code = error_info.get("code", "").lower()
+            err_msg = error_info.get("message", "").lower()
+
+            if "insufficient_quota" in err_type or \
+               "insufficient_quota" in err_code or \
+               "you exceeded your current quota" in err_msg or \
+               "payment required" in err_msg or \
+               ("billing" in err_msg and "issue" in err_msg):
+                info_dict["has_sufficient_quota"] = False
+                new_quota_message = f"Insufficient: {error_info.get('message', err_type)}"
+                info_dict["quota"] = new_quota_message
+                current_status_msg = info_dict["status_message"]
+                additional_msg = f" Quota/Billing issue from error object: {error_info.get('message', err_type)}"
+                if additional_msg.strip() not in current_status_msg :
+                    info_dict["status_message"] = (current_status_msg + additional_msg).strip()
+
+        if "account_deactivated" in str(subscription_info).lower() or \
+           "payment_failed" in str(subscription_info).lower():
+            info_dict["has_sufficient_quota"] = False
+            current_status_msg = info_dict["status_message"]
+            additional_msg = " Account issue (e.g., deactivated, payment failed) from subscription info."
+            if additional_msg.strip() not in current_status_msg:
+                 info_dict["status_message"] = (current_status_msg + additional_msg).strip()
+    else:
+        if info_dict["key_availability"]:
+            info_dict["status_message"] = (info_dict["status_message"] + " Failed to retrieve full subscription details.").strip()
+            info_dict["has_sufficient_quota"] = False
+
+    if info_dict.get("has_sufficient_quota") is True:
+        current_quota_field_value = info_dict.get("quota")
+        if current_quota_field_value:
+            if "insufficient_quota" in str(current_quota_field_value).lower():
+                info_dict["has_sufficient_quota"] = False
+                insufficient_msg = "Insufficient quota indicated in 'quota' field."
+                if insufficient_msg not in info_dict["status_message"]:
+                    info_dict["status_message"] = (info_dict["status_message"] + f" {insufficient_msg}").strip()
+
     return info_dict
 
 
@@ -111,19 +204,22 @@ async def get_key_ant_info(key: str, rate_limit: bool, claude_model: str) -> Dic
 
     return info_dict
 
-
-# 保持用户原始的get_key_gemini_info函数
+########## 更改的部分 ##########
+# 1. get_key_gemini_info 修改函数以包含模型列表
 def get_key_gemini_info(key: str) -> Dict[str, Any]:
-    """获取Gemini密钥信息（用户原始版本）。"""
-    key_avai = check_key_gemini_availability(key) # 假设返回(bool, str)元组
+    """Gemini 获取关键信息并包含模型列表."""
+    # check_key_gemini_availability가 (bool, str, list_of_models) 假设返回元组
+    key_avai = check_key_gemini_availability(key)
     info_dict = {
-        "key": key, # 原版中有key字段
+        "key": key,
+        "key_type": "Google Gemini",
         "key_availability": key_avai[0],
         "status": key_avai[1],
+        # 通过检查返回值长度安全地获取模型列表
+        "models": key_avai[2] if len(key_avai) > 2 else []
     }
     return info_dict
 
-
 def get_key_azure_info(endpoint: str, api_key: str) -> Dict[str, Any]:
     key_avai = check_key_azure_availability(endpoint, api_key)
     info_dict = {
@@ -367,172 +463,156 @@ def not_supported(key: str) -> Dict[str, Any]:
 
 
 # ─────────────────────────────────────────
-#  单个密钥异步处理（已更新）
+#  单一键异步处理（已更新）
 # ─────────────────────────────────────────
 async def process_single_key(key: str, rate_limit: bool, claude_model: str) -> Dict[str, Any]:
-    """分析给定的单个密钥并返回信息字典。包含Gemini密钥的工作状态标志。"""
     _key = key.strip()
 
     if not _key:
         return {"key": "", "key_type": "Empty", "key_availability": False}
 
-    # OpenRouter
     if re.match(re.compile(r"sk-or-v1-[a-z0-9]{64}"), _key):
         result = get_key_openrouter_info(_key)
-        return {"key": _key, **result} # key_type在get_key_openrouter_info中返回
-
-    # Anthropic Claude
+        return {"key": _key, **result}
     if re.match(re.compile(r"sk-ant-api03-[a-zA-Z0-9\-_]{93}AA"), _key) or \
        (_key.startswith("sk-ant-") and len(_key) == 93) or \
        (len(_key) == 89 and re.match(re.compile(r"sk-[a-zA-Z0-9]{86}"), _key)):
         result = await get_key_ant_info(_key, rate_limit, claude_model)
-        return {"key": _key, **result} # 包含key_type
-
-    # Stability
+        return {"key": _key, **result}
     if re.match(re.compile(r"sk-[a-zA-Z0-9]{48}"), _key) and len(_key) == 51 and "T3BlbkFJ" not in _key:
         result = get_key_stability_info(_key)
-        return {"key": _key, **result} # 包含key_type
-
-    # Deepseek
+        return {"key": _key, **result}
     if re.match(re.compile(r"sk-[a-f0-9]{32}"), _key):
          result = get_key_deepseek_info(_key)
-         return {"key": _key, **result} # 包含key_type
-
-    # OpenAI (应该位于其他sk-模式之后)
+         return {"key": _key, **result}
     if _key.startswith("sk-"):
         result = get_key_oai_info(_key)
-        return {"key": _key, **result} # 包含key_type
+        return {"key": _key, **result}
 
-    # Google Gemini处理（已修改）
+    ########## 更改的部分 ##########
+    # 2. process_single_key 在函数中添加特定模型存在性检查逻辑
     if _key.startswith("AIzaSy"):
-        gemini_info = get_key_gemini_info(_key) # 调用原始函数
-        # 添加工作状态标志（检查status）
+        gemini_info = get_key_gemini_info(_key)
         is_working = gemini_info.get("key_availability") and gemini_info.get("status") == "Working"
-        # 构建结果字典：原始信息 + key_type + 工作状态标志
-        # gemini_info已经包含'key'
+
+        # 检查特定模型("gemini-2.5-pro-preview-06-05")是否在可用模型列表中
+        available_models = gemini_info.get("models", [])
+        has_target_model = "gemini-2.5-pro-preview-06-05" in available_models
+
         result = {
-            "key_type": "Google Gemini", # 明确添加类型信息
             **gemini_info,
-            "is_gemini_working": is_working # 工作状态标志
+            "is_gemini_working": is_working,
+            "has_specific_gemini_model": has_target_model # 添加用于保存结果的标志
         }
         return result
 
-    # NovelAI
     if _key.startswith("pst-"):
         result = get_key_nai_info(_key)
-        return {"key": _key, **result} # 包含key_type
-
-    # Replicate
+        return {"key": _key, **result}
     if (_key.startswith("r8_") and len(_key) == 40) or (_key.islower() and len(_key) == 40):
         result = get_key_replicate_info(_key)
-        return {"key": _key, **result} # 包含key_type
-
-    # xAI
+        return {"key": _key, **result}
     if _key.startswith("xai-"):
         result = get_key_xai_info(_key)
-        return {"key": _key, **result} # 包含key_type
-
-    # Azure endpoint: "name:key"
+        return {"key": _key, **result}
     if len(_key.split(":")) == 2:
         name, potential_key = _key.split(":", 1)
         if re.fullmatch(r'[a-fA-F0-9]{32}', potential_key) and "openai.azure.com" not in name:
              endpoint = f"https://{name}.openai.azure.com/"
              api_key = potential_key
              result = get_key_azure_info(endpoint, api_key)
-             return {"key": _key, **result} # 包含key_type
-
-    # Azure endpoint: "https://xxx.openai.azure.com;key"
+             return {"key": _key, **result}
     if ";" in _key and "openai.azure.com" in _key.split(";")[0]:
         endpoint, api_key = _key.split(";", 1)
         result = get_key_azure_info(endpoint, api_key)
-        return {"key": _key, **result} # 包含key_type
-
-    # AWS
+        return {"key": _key, **result}
     if _key.startswith("AKIA") and len(_key.split(":")[0]) == 20 and _key.split(":")[0].isalnum() and _key.split(":")[0].isupper() and len(_key.split(':')) == 2:
         result = await get_key_aws_info(_key)
-        return {"key": _key, **result} # 包含key_type
-
-    # ElevenLabs
+        return {"key": _key, **result}
     if re.fullmatch(r"[a-f0-9]{32}", _key) or re.fullmatch(r"sk_[a-f0-9]{48}", _key):
         result = get_key_elevenlabs_info(_key)
-        return {"key": _key, **result} # 包含key_type
-
-    # Mistral
+        return {"key": _key, **result}
     if re.fullmatch(r"[a-zA-Z0-9]{32}", _key) and not _key.startswith('sk-'):
         result = get_key_mistral_info(_key)
-        return {"key": _key, **result} # 包含key_type
-
-    # Groq
+        return {"key": _key, **result}
     if re.match(re.compile(r"gsk_[a-zA-Z0-9]{20}WGdyb3FY[a-zA-Z0-9]{24}"), _key):
         result = get_key_groq_info(_key)
-        return {"key": _key, **result} # 包含key_type
-
-    # GCP - refresh token
+        return {"key": _key, **result}
     if re.match(re.compile(r"[\w\-]+:[\w\-@\.]+:.+:.+"), _key):
         parts = _key.split(':')
         if len(parts) >= 4:
             result = await get_key_gcp_info(_key, 0)
-            return {"key": _key, **result} # 包含key_type
-
-    # GCP - service account
+            return {"key": _key, **result}
     if re.match(re.compile(r"[\w\-]+:[\w\-@\.]+:.+\\n"), _key):
          parts = _key.split(':')
          if len(parts) >= 3:
             result = await get_key_gcp_info(_key, 1)
-            return {"key": _key, **result} # 包含key_type
-
-    # Not supported
-    result = not_supported(_key) # 包含key_type, key_availability, status
-    # 添加密钥值并返回
+            return {"key": _key, **result}
+    result = not_supported(_key)
     return {"key": _key, **result}
 
 
 # ─────────────────────────────────────────
-#  多密钥异步处理函数（已更新）
+#  多个key异步处理函数（已更新）
 # ─────────────────────────────────────────
-async def sort_keys(text: str, rate_limit: bool, claude_model: str) -> Tuple[List[Dict[str, Any]], str]:
-    """分析文本框中输入的多个密钥（按行），并返回完整结果和有效Gemini密钥列表。"""
+async def sort_keys(text: str, rate_limit: bool, claude_model: str) -> Tuple[List[Dict[str, Any]], str, str, str]:
     keys = [k.strip() for k in text.splitlines() if k.strip()]
     if not keys:
-        return [], ""
+        return [], "", "", ""
 
     tasks = [process_single_key(k, rate_limit, claude_model) for k in keys]
     results = await asyncio.gather(*tasks)
 
-    # 仅筛选有效的Gemini密钥（使用is_gemini_working标志）
     working_gemini_keys = []
+    working_oai_keys = []
+    working_anthropic_keys = []
+
     for result in results:
-        # 检查is_gemini_working键是否存在且为True
-        if result.get("is_gemini_working"): # 使用get()避免键不存在时出错
-            # 原始get_key_gemini_info包含'key'字段，可直接使用
-            working_gemini_keys.append(result["key"])
+        key_value = result.get("key")
+        if not key_value:
+            continue
 
-    # 返回JSON列表结果和Gemini密钥列表字符串
-    return results, "\n".join(working_gemini_keys)
+        key_type = result.get("key_type")
+
+        ########## 更改的部分 ##########
+        # 3. sort_keys 加强函数中Gemini键的过滤条件
+        # 'is_gemini_working'和'has_specific_gemini_model' 检查所有标志
+        if result.get("is_gemini_working") and result.get("has_specific_gemini_model"):
+            working_gemini_keys.append(key_value)
+        elif key_type == "OpenAI" and \
+             result.get("key_availability") is True and \
+             result.get("has_sufficient_quota") is True:
+            working_oai_keys.append(key_value)
+        elif key_type == "Anthropic Claude" and result.get("key_availability") is True:
+            working_anthropic_keys.append(key_value)
+
+    return results, "\n".join(working_oai_keys), "\n".join(working_anthropic_keys), "\n".join(working_gemini_keys)
 
 
 # ─────────────────────────────────────────
-#  UI工具函数
+#  UI util （已更新）
 # ─────────────────────────────────────────
-def clear_inputs(): # 移除输入参数（不需要key_box值）
-    return "", "", "" # 重置所有输入/输出
+def clear_inputs():
+    return "", "", "", "", ""
 
 
 # ─────────────────────────────────────────
-#  Gradio界面（已更新）
+#  Gradio UI （已更新）
 # ─────────────────────────────────────────
 with gr.Blocks() as demo:
     gr.Markdown(
         """
-# 多API密钥状态检查器（Gemini增强版）
-*（基于shaocongma、CncAnon1、su、Drago、kingbased密钥检查器）*
-检查各种API密钥的状态和详细信息，包括OpenAI、Anthropic、Gemini、Azure、Mistral、Replicate、AWS Claude、OpenRouter、Vertex AI（GCP Anthropic）、Groq、NovelAI、ElevenLabs、xAI、Stability AI和DeepSeek。
-**密钥格式：**
-*   **AWS:** `AWS_ACCESS_KEY_ID:AWS_SECRET_ACCESS_KEY`（root可能不准确）
-*   **Azure:** `RESOURCE_NAME:API_KEY` **或** `https://RESOURCE_NAME.openai.azure.com;API_KEY`
-*   **GCP服务账号:** `PROJECT_ID:CLIENT_EMAIL:PRIVATE_KEY`（确保在密钥中包含`\\n`作为换行符）
-*   **GCP刷新令牌:** `PROJECT_ID:CLIENT_ID:CLIENT_SECRET:REFRESH_TOKEN`
-*   **其他密��:** 供应商提供的标准格式。
+# Multi-API Key Status Checker (OAI, Anthropic, Gemini Enhanced)
+*(Based on shaocongma, CncAnon1, su, Drago, kingbased key checkers)*
+Check the status and details of various API keys including OpenAI, Anthropic, Gemini, Azure, Mistral, Replicate, AWS Claude, OpenRouter, Vertex AI (GCP Anthropic), Groq, NovelAI, ElevenLabs, xAI, Stability AI, and DeepSeek.
+This version highlights working OpenAI (with sufficient quota), Anthropic, and Gemini keys in separate text boxes.
+
+**Key Formats:**
+*   **AWS:** `AWS_ACCESS_KEY_ID:AWS_SECRET_ACCESS_KEY` (root might not be accurate)
+*   **Azure:** `RESOURCE_NAME:API_KEY` **or** `https://RESOURCE_NAME.openai.azure.com;API_KEY`
+*   **GCP Service Account:** `PROJECT_ID:CLIENT_EMAIL:PRIVATE_KEY` (ensure `\\n` is included for newlines in the key)
+*   **GCP Refresh Token:** `PROJECT_ID:CLIENT_ID:CLIENT_SECRET:REFRESH_TOKEN`
+*   **Other keys:** Standard format provided by the vendor.
 """
     )
 
@@ -548,47 +628,60 @@ with gr.Blocks() as demo:
             key_box = gr.Textbox(
                 lines=5,
                 max_lines=20,
-                label="API密钥（每行一个）",
-                placeholder="在此处输入API密钥，每行一个。",
+                label="API Key(s) - One per line",
+                placeholder="Enter one API key per line here.",
             )
             with gr.Row():
                  claude_model = gr.Dropdown(
                     claude_options,
                     value="claude-3-haiku-20240307",
-                    label="Claude模型（用于过滤/并发检查）",
+                    label="Claude Model (for filter/concurrent check)",
                     scale=3
                  )
-                 rate_limit = gr.Checkbox(label="检查Claude并发限制（实验性）", scale=1)
+                 rate_limit = gr.Checkbox(label="Check Claude concurrent limit (exp.)", scale=1)
 
             with gr.Row():
-                clear_button = gr.Button("清除所有")
-                submit_button = gr.Button("检查密钥", variant="primary")
+                clear_button = gr.Button("Clear All")
+                submit_button = gr.Button("Check Keys", variant="primary")
 
         with gr.Column(scale=2):
-            info = gr.JSON(label="API密钥信息（所有结果）", open=True)
-            # 添加Gemini专用输出框
+            info = gr.JSON(label=" API Key Information (All Results)", open=True)
+            oai_keys_output = gr.Textbox(
+                label="Working OpenAI Keys (Sufficient Quota)",
+                info="Lists OpenAI keys confirmed as working and having sufficient quota.",
+                lines=3,
+                max_lines=10,
+                interactive=False,
+            )
+            anthropic_keys_output = gr.Textbox(
+                label="Working Anthropic Keys",
+                info="Lists Anthropic keys confirmed as working (key_availability is True).",
+                lines=3,
+                max_lines=10,
+                interactive=False,
+            )
+
+            ########## 变更的部分 ##########
+            # 4. Gemini 更新结果框的标签和说明，以明确新的筛选标准
             gemini_keys_output = gr.Textbox(
-                label="有效的Gemini密钥",
-                info="列出确认为'有效'的Gemini密钥。",
+                label="Working Gemini Keys (with gemini-2.5-pro-preview-06-05)",
+                info="Lists Gemini keys that are 'Working' AND have 'gemini-2.5-pro-preview-06-05' in their available models.",
                 lines=3,
                 max_lines=10,
-                interactive=False, # 用户不可编辑
+                interactive=False,
             )
 
-    # 点击清除按钮时重置输入框和输出框
     clear_button.click(
-        fn=clear_inputs, # 使用修改后的函数
+        fn=clear_inputs,
         inputs=None,
-        outputs=[key_box, info, gemini_keys_output]
+        outputs=[key_box, info, oai_keys_output, anthropic_keys_output, gemini_keys_output]
     )
 
-    # 点击提交按钮时调用sort_keys函数并分配结果
     submit_button.click(
         fn=sort_keys,
         inputs=[key_box, rate_limit, claude_model],
-        outputs=[info, gemini_keys_output], # 将结果映射到两个输出组件
+        outputs=[info, oai_keys_output, anthropic_keys_output, gemini_keys_output],
         api_name="sort_keys",
     )
 
-# demo.launch(share=True) # 需要时激活共享链接
-demo.launch()
+demo.launch()