Update app.py

app.py

@@ -9,7 +9,8 @@ from flask import Flask, request, jsonify
 logging.basicConfig(level=logging.INFO,
                     format='%(asctime)s - %(levelname)s - %(message)s')
 
-API_ENDPOINT = "https://api.siliconflow.cn/v1/user/info"  # 修改为获取用户信息的接口
+API_ENDPOINT = "https://api.siliconflow.cn/v1/user/info"
+TEST_MODEL_ENDPOINT = "https://api.siliconflow.cn/v1/chat/completions"
 
 app = Flask(__name__)
 
@@ -25,7 +26,7 @@ def get_credit_summary(api_key):
         response = requests.get(API_ENDPOINT, headers=headers)
         response.raise_for_status()
         data = response.json().get("data", {})
-        total_balance = data.get("totalBalance", 0)  # 获取总余额
+        total_balance = data.get("totalBalance", 0)
         return {"total_balance": total_balance}
     except requests.exceptions.RequestException as e:
         logging.error(f"获取额度信息失败，API Key：{api_key}，错误信息：{e}")
@@ -34,19 +35,73 @@ def get_credit_summary(api_key):
         logging.error(f"解析额度信息失败，API Key：{api_key}，错误信息：{e}")
         return None
 
+def test_model_availability(api_key, model_name):
+    """
+    测试指定的模型是否可用。
+    """
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    try:
+        response = requests.post(TEST_MODEL_ENDPOINT,
+                                 headers=headers,
+                                 json={
+                                     "model": model_name,
+                                     "messages": [{"role": "user", "content": "hi"}],
+                                     "max_tokens": 10,
+                                     "stream": False
+                                 },
+                                 timeout=10)
+        response.raise_for_status()
+        # 检查是否是429错误
+        if response.status_code == 429:
+            return True
+        
+        response.json()  # 尝试解析 JSON 响应
+        return True
+    except requests.exceptions.RequestException as e:
+        logging.error(f"测试模型 {model_name} 可用性失败，API Key：{api_key}，错误信息：{e}")
+        return False
+    except ValueError:
+        logging.error(f"测试模型 {model_name} 可用性失败，API Key：{api_key}，响应不是有效的 JSON 格式")
+        return False
+
 def load_keys():
     """
-    从环境变量中加载 keys，并记录到日志中，同时获取并记录每个 key 的额度信息。
+    从环境变量中加载 keys，并根据额度和模型可用性进行分类，然后记录到日志中。
     """
     keys_str = os.environ.get("KEYS")
+    test_model = os.environ.get("TEST_MODEL", "Pro/google/gemma-2-9b-it")
+
+    invalid_keys = []
+    free_keys = []
+    unverified_keys = []
+    valid_keys = []
+
     if keys_str:
         keys = [key.strip() for key in keys_str.split(',')]
         logging.info(f"加载的 keys：{keys}")
 
         for key in keys:
             credit_summary = get_credit_summary(key)
-            if credit_summary:
-                logging.info(f"API Key：{key}，额度信息：{credit_summary}")
+            if credit_summary is None:
+                invalid_keys.append(key)
+            else:
+                total_balance = credit_summary.get("total_balance", 0)
+                if total_balance <= 0:
+                    free_keys.append(key)
+                else:
+                    if test_model_availability(key, test_model):
+                        valid_keys.append(key)
+                    else:
+                        unverified_keys.append(key)
+
+        logging.info(f"无效 KEY：{invalid_keys}")
+        logging.info(f"免费 KEY：{free_keys}")
+        logging.info(f"未实名 KEY：{unverified_keys}")
+        logging.info(f"有效 KEY：{valid_keys}")
+
     else:
         logging.warning("环境变量 KEYS 未设置。")
 
@@ -69,45 +124,22 @@ def check_tokens():
     处理前端发送的 Token 检测请求。
     """
     tokens = request.json.get('tokens', [])
+    test_model = os.environ.get("TEST_MODEL", "Pro/google/gemma-2-9b-it")
+
     results = []
     for token in tokens:
-        is_valid = False
-        balance = 0
-        error_message = ""
-
-        # 先尝试进行一个简单的请求来验证 Token 的有效性
-        try:
-            headers = {
-                "Authorization": f"Bearer {token}",
-                "Content-Type": "application/json"
-            }
-            test_response = requests.post("https://api.siliconflow.cn/v1/chat/completions",
-                                          headers=headers,
-                                          json={
-                                              "model": "Qwen/Qwen2.5-72B-Instruct",
-                                              "messages": [{"role": "user", "content": "hi"}],
-                                              "max_tokens": 100,
-                                              "stream": False
-                                          },
-                                          timeout=5)
-
-            if test_response.status_code == 200:
-                is_valid = True
-                credit_summary = get_credit_summary(token)
-                if credit_summary:
-                    balance = credit_summary.get("total_balance", 0)
+        credit_summary = get_credit_summary(token)
+        if credit_summary is None:
+            results.append({"token": token, "type": "无效 KEY", "balance": 0, "message": "无法获取额度信息"})
+        else:
+            total_balance = credit_summary.get("total_balance", 0)
+            if total_balance <= 0:
+                results.append({"token": token, "type": "免费 KEY", "balance": total_balance, "message": "额度不足"})
             else:
-                error_message = test_response.json().get("message", "未知错误")
-
-        except requests.exceptions.RequestException as e:
-            error_message = str(e)
-
-        results.append({
-            "token": token,
-            "isValid": is_valid,
-            "balance": balance,
-            "message": error_message
-        })
+                if test_model_availability(token, test_model):
+                    results.append({"token": token, "type": "有效 KEY", "balance": total_balance, "message": "可以使用指定模型"})
+                else:
+                    results.append({"token": token, "type": "未实名 KEY", "balance": total_balance, "message": "无法使用指定模型"})
 
     return jsonify(results)