Update app.py

app.py

@@ -24,6 +24,12 @@ free_keys_global = []
 unverified_keys_global = []
 valid_keys_global = []
 
+# 使用 ThreadPoolExecutor 管理并发
+executor = concurrent.futures.ThreadPoolExecutor(max_workers=20)
+
+# 为每个模型维护一个独立的 key 索引
+model_key_indices = {}
+
 def get_credit_summary(api_key):
     """
     使用 API 密钥获取额度信息。
@@ -48,7 +54,7 @@ def get_credit_summary(api_key):
         logging.error(f"total_balance 无法转换为浮点数，API Key：{api_key}，错误信息：{e}")
         return None
 
-FREE_MODEL_TEST_KEY = "sk-bmjbjzleaqfgtqfzmcnsbagxrlohriadnxqrzfocbizaxukw"  # 硬编码免费模型测试 KEY
+FREE_MODEL_TEST_KEY = "sk-bmjbjzleaqfgtqfzmcnsbagxrlohriadnxqrzfocbizaxukw"
 
 def test_model_availability(api_key, model_name):
     """
@@ -68,7 +74,6 @@ def test_model_availability(api_key, model_name):
                                      "stream": False
                                  },
                                  timeout=10)
-        # 正常返回或返回 429 都被认为是免费模型
         if response.status_code == 429 or response.status_code == 200 :
             return True
         else:
@@ -83,16 +88,11 @@ def refresh_models():
     """
     global all_models, free_models
 
-    # 使用 FREE_MODEL_TEST_KEY 获取所有模型列表
     all_models = get_all_models(FREE_MODEL_TEST_KEY)
-
     free_models = []
 
-    # 使用 ThreadPoolExecutor 并发执行测试
-    with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:  # 限制最大工作线程数为 10，可根据实际情况调整
-        # 创建一个模型到 future 的映射
+    with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
         future_to_model = {executor.submit(test_model_availability, FREE_MODEL_TEST_KEY, model): model for model in all_models}
-
         for future in concurrent.futures.as_completed(future_to_model):
             model = future_to_model[future]
             try:
@@ -108,39 +108,43 @@ def refresh_models():
 def load_keys():
     """
     从环境变量中加载 keys，并根据额度和模型可用性进行分类，然后记录到日志中。
+    使用线程池并发处理每个 key。
     """
     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 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:
+        with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
+            future_to_key = {executor.submit(process_key, key, test_model): key for key in keys}
+
+            invalid_keys = []
+            free_keys = []
+            unverified_keys = []
+            valid_keys = []
+
+            for future in concurrent.futures.as_completed(future_to_key):
+                key = future_to_key[future]
+                try:
+                    key_type = future.result()
+                    if key_type == "invalid":
+                        invalid_keys.append(key)
+                    elif key_type == "free":
+                        free_keys.append(key)
+                    elif key_type == "unverified":
                         unverified_keys.append(key)
+                    elif key_type == "valid":
+                        valid_keys.append(key)
+                except Exception as exc:
+                    logging.error(f"处理 KEY {key} 生成异常: {exc}")
 
         logging.info(f"无效 KEY：{invalid_keys}")
         logging.info(f"免费 KEY：{free_keys}")
         logging.info(f"未实名 KEY：{unverified_keys}")
         logging.info(f"有效 KEY：{valid_keys}")
-      
-        # 更新全局的 key 列表
+
         global invalid_keys_global, free_keys_global, unverified_keys_global, valid_keys_global
         invalid_keys_global = invalid_keys
         free_keys_global = free_keys
@@ -150,6 +154,23 @@ def load_keys():
     else:
         logging.warning("环境变量 KEYS 未设置。")
 
+def process_key(key, test_model):
+    """
+    处理单个 key，判断其类型。
+    """
+    credit_summary = get_credit_summary(key)
+    if credit_summary is None:
+        return "invalid"
+    else:
+        total_balance = credit_summary.get("total_balance", 0)
+        if total_balance <= 0:
+            return "free"
+        else:
+            if test_model_availability(key, test_model):
+                return "valid"
+            else:
+                return "unverified"
+
 def get_all_models(api_key):
     """
     获取所有模型列表。
@@ -162,7 +183,6 @@ def get_all_models(api_key):
         response = requests.get(MODELS_ENDPOINT, headers=headers, params={"sub_type": "chat"})
         response.raise_for_status()
         data = response.json()
-        # 确保 data 是字典且包含 'data' 键，'data' 对应的值是一个列表
         if isinstance(data, dict) and 'data' in data and isinstance(data['data'], list):
             return [model.get("id") for model in data["data"] if isinstance(model, dict) and "id" in model]
         else:
@@ -186,26 +206,59 @@ def determine_request_type(model_name):
     else:
         return "unknown"
 
-def select_key(request_type):
+def select_key(request_type, model_name):
     """
-    根据请求类型选择合适的 KEY。
+    根据请求类型和模型名称选择合适的 KEY，并实现轮询和重试机制。
     """
     if request_type == "free":
-        # 免费请求：使用 2、3、4 类 KEY
         available_keys = free_keys_global + unverified_keys_global + valid_keys_global
     elif request_type == "paid":
-        # 付费请求：使用 3、4 类 KEY
         available_keys = unverified_keys_global + valid_keys_global
     else:
-        # 未知请求：使用所有 KEY
         available_keys = free_keys_global + unverified_keys_global + valid_keys_global
 
     if not available_keys:
         return None
 
-    # 简单的轮询策略选择 KEY
-    key = available_keys[int(time.time() * 1000) % len(available_keys)]
-    return key
+    # 获取当前模型的索引，如果没有则初始化为 0
+    current_index = model_key_indices.get(model_name, 0)
+
+    # 轮询并重试
+    for _ in range(len(available_keys)):
+        key = available_keys[current_index % len(available_keys)]
+        current_index += 1
+        
+        # 检查 KEY 是否有效
+        if key_is_valid(key, request_type):
+            # 更新模型索引并返回 KEY
+            model_key_indices[model_name] = current_index
+            return key
+        else:
+            logging.warning(f"KEY {key} 无效或达到限制，尝试下一个 KEY")
+
+    # 所有 KEY 都尝试过，重置索引并返回 None
+    model_key_indices[model_name] = 0
+    return None
+
+def key_is_valid(key, request_type):
+    """
+    检查 KEY 是否有效，根据不同的请求类型进行不同的检查。
+    """
+    if request_type == "invalid":
+        return False  # 无效 KEY 始终返回 False
+    
+    credit_summary = get_credit_summary(key)
+    if credit_summary is None:
+        return False
+    
+    total_balance = credit_summary.get("total_balance", 0)
+    
+    if request_type == "free":
+        return True # 免费 KEY，只要能获取到信息，就认为是有效的
+    elif request_type == "paid" or request_type == "unverified":
+        return total_balance > 0 # 付费 KEY 或未实名 KEY，需要余额大于 0 才是有效的
+    else:
+        return False # 未知类型，返回 False
 
 def check_authorization(request):
     """
@@ -227,50 +280,50 @@ def check_authorization(request):
 
     return True
 
-# 创建一个后台调度器
 scheduler = BackgroundScheduler()
-
-# 添加定时任务，每小时执行一次 load_keys 函数
 scheduler.add_job(load_keys, 'interval', hours=1)
-# 添加定时任务，每10分钟执行一次 refresh_models 函数
 scheduler.add_job(refresh_models, 'interval', minutes=10)
 
 @app.route('/')
 def index():
-    """
-    处理根路由的访问请求。
-    """
     return "<h1>Welcome to SiliconFlow</h1>"
 
 @app.route('/check_tokens', methods=['POST'])
 def check_tokens():
     """
     处理前端发送的 Token 检测请求。
+    使用线程池并发处理每个 token。
     """
     tokens = request.json.get('tokens', [])
     test_model = os.environ.get("TEST_MODEL", "Pro/google/gemma-2-9b-it")
 
-    results = []
-    for token in tokens:
-        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:
-                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": "无法使用指定模型"})
+    with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
+        future_to_token = {executor.submit(process_key, token, test_model): token for token in tokens}
+
+        results = []
+        for future in concurrent.futures.as_completed(future_to_token):
+            token = future_to_token[future]
+            try:
+                key_type = future.result()
+                credit_summary = get_credit_summary(token)
+                balance = credit_summary.get("total_balance", 0) if credit_summary else 0
+                if key_type == "invalid":
+                    results.append({"token": token, "type": "无效 KEY", "balance": balance, "message": "无法获取额度信息"})
+                elif key_type == "free":
+                    results.append({"token": token, "type": "免费 KEY", "balance": balance, "message": "额度不足"})
+                elif key_type == "unverified":
+                    results.append({"token": token, "type": "未实名 KEY", "balance": balance, "message": "无法使用指定模型"})
+                elif key_type == "valid":
+                    results.append({"token": token, "type": "有效 KEY", "balance": balance, "message": "可以使用指定模型"})
+            except Exception as exc:
+                logging.error(f"处理 Token {token} 生成异常: {exc}")
 
     return jsonify(results)
 
 @app.route('/handsome/v1/chat/completions', methods=['POST'])
 def handsome_chat_completions():
     """
-    处理 /handsome/v1/chat/completions 路由的请求，添加鉴权。
+    处理 /handsome/v1/chat/completions 路由的请求，添加鉴权，并实现 KEY 的轮询和重试机制。
     """
     if not check_authorization(request):
         return jsonify({"error": "Unauthorized"}), 401
@@ -281,17 +334,16 @@ def handsome_chat_completions():
 
     model_name = data['model']
     request_type = determine_request_type(model_name)
-    api_key = select_key(request_type)
+    api_key = select_key(request_type, model_name)
 
     if not api_key:
-        return jsonify({"error": "No available API key for this request type"}), 400
+        return jsonify({"error": "No available API key for this request type or all keys have reached their limits"}), 429
 
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json"
     }
 
-    # 转发请求到真正的 API
     try:
         response = requests.post(
             TEST_MODEL_ENDPOINT,
@@ -300,8 +352,7 @@ def handsome_chat_completions():
             stream=data.get("stream", False),
             timeout=60
         )
-      
-        # 检查是否是429错误
+
         if response.status_code == 429:
             return jsonify(response.json()), 429
 
@@ -315,13 +366,9 @@ def handsome_chat_completions():
 
 @app.route('/handsome/v1/models', methods=['GET'])
 def list_models():
-    """
-    处理 /handsome/v1/models 路由的请求，返回模型列表。
-    """
     if not check_authorization(request):
         return jsonify({"error": "Unauthorized"}), 401
 
-    # 返回模型列表和免费模型列表
     return jsonify({
         "data": [{"id": model, "object": "model"} for model in all_models],
         "free_models": free_models
@@ -330,26 +377,32 @@ def list_models():
 def get_billing_info():
     """
     获取所有KEY的额度信息。
+    使用线程池并发处理每个 key。
     """
+    keys = valid_keys_global + unverified_keys_global
     total_balance = 0
-    for key in valid_keys_global + unverified_keys_global:
-        credit_summary = get_credit_summary(key)
-        if credit_summary:
-            total_balance += credit_summary.get("total_balance", 0)
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
+        futures = [executor.submit(get_credit_summary, key) for key in keys]
+
+        for future in concurrent.futures.as_completed(futures):
+            try:
+                credit_summary = future.result()
+                if credit_summary:
+                    total_balance += credit_summary.get("total_balance", 0)
+            except Exception as exc:
+                logging.error(f"获取额度信息生成异常: {exc}")
+
     return total_balance
 
 @app.route('/handsome/v1/dashboard/billing/usage', methods=['GET'])
 def billing_usage():
-    """
-    处理 /handsome/v1/dashboard/billing/usage 路由的请求，返回用量信息（修改后，始终返回 0）。
-    """
     if not check_authorization(request):
         return jsonify({"error": "Unauthorized"}), 401
 
     end_date = datetime.now()
     start_date = end_date - timedelta(days=30)
 
-    # 构造用量数据（修改后，始终返回 0）
     daily_usage = []
     current_date = start_date
     while current_date <= end_date:
@@ -367,9 +420,6 @@ def billing_usage():
 
 @app.route('/handsome/v1/dashboard/billing/subscription', methods=['GET'])
 def billing_subscription():
-    """
-    处理 /handsome/v1/dashboard/billing/subscription 路由的请求，返回订阅信息。
-    """
     if not check_authorization(request):
         return jsonify({"error": "Unauthorized"}), 401
 
@@ -381,7 +431,7 @@ def billing_subscription():
         "canceled": False,
         "canceled_at": None,
         "delinquent": None,
-        "access_until": int(datetime(9999, 12, 31).timestamp()),  # 设置一个较远的未来时间
+        "access_until": int(datetime(9999, 12, 31).timestamp()),
         "soft_limit": 0,
         "hard_limit": total_balance,
         "system_hard_limit": total_balance,
@@ -401,25 +451,19 @@ def billing_subscription():
     })
 
 if __name__ == '__main__':
-    # 打印所有环境变量，方便调试
     logging.info(f"环境变量：{os.environ}")
 
-    # 初始化全局的 key 列表
     invalid_keys_global = []
     free_keys_global = []
     unverified_keys_global = []
     valid_keys_global = []
 
-    # 启动调度器
     scheduler.start()
 
-    # 手动触发一次 load_keys 任务
     load_keys()
     logging.info("首次加载 keys 已手动触发执行")
 
-    # 手动触发一次 refresh_models 任务
     refresh_models()
     logging.info("首次刷新模型列表已手动触发执行")
 
-    # 启动 Flask 应用，监听所有 IP 的 7860 端口（Hugging Face Space 默认端口）
     app.run(debug=False, host='0.0.0.0', port=int(os.environ.get('PORT', 7860)))