import os import time import logging import requests from apscheduler.schedulers.background import BackgroundScheduler from flask import Flask, request, jsonify, Response, stream_with_context # 配置日志记录 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') 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 密钥获取额度信息。 """ headers = { "Authorization": f"Bearer {api_key}", "Content-Type": "application/json" } try: response = requests.get(API_ENDPOINT, headers=headers) response.raise_for_status() data = response.json().get("data", {}) total_balance = data.get("totalBalance", 0) return {"total_balance": float(total_balance)} except requests.exceptions.RequestException as e: logging.error(f"获取额度信息失败，API Key：{api_key}，错误信息：{e}") return None except (KeyError, TypeError) as e: logging.error(f"解析额度信息失败，API Key：{api_key}，错误信息：{e}") return None except ValueError as e: logging.error(f"total_balance 无法转换为浮点数，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) # 检查是否是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，并根据额度和模型可用性进行分类，然后记录到日志中。 """ 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: 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}") # 更新全局的 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(): """ 处理根路由的访问请求。 """ return "

Welcome to SiliconFlow

" @app.route('/check_tokens', methods=['POST']) 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: 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": "无法使用指定模型"}) 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() # 手动触发一次 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)))