Update app.py

app.py

@@ -3,7 +3,7 @@ import time
 import logging
 import requests
 from apscheduler.schedulers.background import BackgroundScheduler
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, Response, stream_with_context
 
 # 配置日志记录
 logging.basicConfig(level=logging.INFO,
@@ -11,9 +11,14 @@ logging.basicConfig(level=logging.INFO,
 
 API_ENDPOINT = "https://api.siliconflow.cn/v1/user/info"
 TEST_MODEL_ENDPOINT = "https://api.siliconflow.cn/v1/chat/completions"
+MODELS_ENDPOINT = "https://api.siliconflow.cn/v1/models"
 
 app = Flask(__name__)
 
+# 全局变量，用于存储模型列表和免费模型列表
+all_models = []
+free_models = []
+
 def get_credit_summary(api_key):
     """
     使用 API 密钥获取额度信息。
@@ -27,7 +32,6 @@ def get_credit_summary(api_key):
         response.raise_for_status()
         data = response.json().get("data", {})
         total_balance = data.get("totalBalance", 0)
-        # 将 total_balance 转换为浮点数
         return {"total_balance": float(total_balance)}
     except requests.exceptions.RequestException as e:
         logging.error(f"获取额度信息失败，API Key：{api_key}，错误信息：{e}")
@@ -57,7 +61,6 @@ def test_model_availability(api_key, model_name):
                                      "stream": False
                                  },
                                  timeout=10)
-        response.raise_for_status()
         # 检查是否是429错误
         if response.status_code == 429:
             return True
@@ -105,15 +108,104 @@ def load_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
+        unverified_keys_global = unverified_keys
+        valid_keys_global = valid_keys
 
     else:
         logging.warning("环境变量 KEYS 未设置。")
 
+def get_all_models(api_key):
+    """
+    获取所有模型列表。
+    """
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    try:
+        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:
+            logging.error("获取模型列表失败：响应数据格式不正确")
+            return []
+    except requests.exceptions.RequestException as e:
+        logging.error(f"获取模型列表失败，API Key：{api_key}，错误信息：{e}")
+        return []
+    except (KeyError, TypeError) as e:
+        logging.error(f"解析模型列表失败，API Key：{api_key}，错误信息：{e}")
+        return []
+
+def refresh_models():
+    """
+    刷新模型列表和免费模型列表。
+    """
+    global all_models, free_models
+    
+    # 使用 valid_keys_global 中的第一个 key 获取完整模型列表
+    if valid_keys_global:
+        all_models = get_all_models(valid_keys_global[0])
+    else:
+        logging.warning("没有有效的key，无法获取完整模型列表。")
+        all_models = []
+
+    # 使用 free_keys_global 中的第一个 key 获取免费模型列表
+    if free_keys_global:
+        free_models = get_all_models(free_keys_global[0])
+    else:
+        logging.warning("没有免费的key，无法获取免费模型列表。")
+        free_models = []
+
+    logging.info(f"所有模型列表：{all_models}")
+    logging.info(f"免费模型列表：{free_models}")
+
+def determine_request_type(model_name):
+    """
+    根据用户请求的模型判断请求类型。
+    """
+    if model_name in free_models:
+        return "free"
+    elif model_name in all_models:
+        return "paid"
+    else:
+        return "unknown"
+
+def select_key(request_type):
+    """
+    根据请求类型选择合适的 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
+
 # 创建一个后台调度器
 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():
@@ -147,10 +239,59 @@ def check_tokens():
 
     return jsonify(results)
 
+@app.route('/handsome/v1/chat/completions', methods=['POST'])
+def handsome_chat_completions():
+    """
+    处理 /handsome/v1/chat/completions 路由的请求。
+    """
+    data = request.get_json()
+    if not data or 'model' not in data:
+        return jsonify({"error": "Invalid request data"}), 400
+
+    model_name = data['model']
+    request_type = determine_request_type(model_name)
+    api_key = select_key(request_type)
+
+    if not api_key:
+        return jsonify({"error": "No available API key for this request type"}), 400
+
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+
+    # 转发请求到真正的 API
+    try:
+        response = requests.post(
+            TEST_MODEL_ENDPOINT,
+            headers=headers,
+            json=data,
+            stream=data.get("stream", False),
+            timeout=60
+        )
+        
+        # 检查是否是429错误
+        if response.status_code == 429:
+            return jsonify(response.json()), 429
+
+        if data.get("stream", False):
+            return Response(stream_with_context(response.iter_content(chunk_size=1024)), content_type=response.headers['Content-Type'])
+        else:
+            response.raise_for_status()
+            return jsonify(response.json())
+    except requests.exceptions.RequestException as e:
+        return jsonify({"error": str(e)}), 500
+
 if __name__ == '__main__':
     # 打印所有环境变量，方便调试
     logging.info(f"环境变量：{os.environ}")
 
+    # 初始化全局的 key 列表
+    invalid_keys_global = []
+    free_keys_global = []
+    unverified_keys_global = []
+    valid_keys_global = []
+
     # 启动调度器
     scheduler.start()
 
@@ -158,5 +299,9 @@ if __name__ == '__main__':
     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)))