Update app.py

app.py

@@ -56,17 +56,21 @@ session = requests_session_with_retries()
 app = Flask(__name__)
 app.wsgi_app = ProxyFix(app.wsgi_app, x_for=1)
 
-text_models = []
-free_text_models = []
-embedding_models = []
-free_embedding_models = []
-image_models = []
-free_image_models = []
-
-invalid_keys_global = []
-free_keys_global = []
-unverified_keys_global = []
-valid_keys_global = []
+models = {
+    "text": [],
+    "free_text": [],
+    "embedding": [],
+    "free_embedding": [],
+    "image": [],
+    "free_image": []
+}
+
+key_status = {
+    "invalid": [],
+    "free": [],
+    "unverified": [],
+    "valid": []
+}
 
 executor = concurrent.futures.ThreadPoolExecutor(max_workers=10000)
 model_key_indices = {}
@@ -112,30 +116,34 @@ FREE_IMAGE_LIST = [
     "stabilityai/stable-diffusion-2-1"
 ]
 
-def test_model_availability(api_key, model_name):
+def test_model_availability(api_key, model_name, model_type="chat"):
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json"
     }
+    
+    if model_type == "image":
+        return model_name in FREE_IMAGE_LIST
+    
     try:
+        endpoint = EMBEDDINGS_ENDPOINT if model_type == "embedding" else TEST_MODEL_ENDPOINT
+        payload = (
+            {"model": model_name, "input": ["hi"]}
+            if model_type == "embedding"
+            else {"model": model_name, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 5, "stream": False}
+        )
+        timeout = 10 if model_type == "embedding" else 5
+        
         response = session.post(
-            TEST_MODEL_ENDPOINT,
+            endpoint,
             headers=headers,
-            json={
-                "model": model_name,
-                "messages": [{"role": "user", "content": "hi"}],
-                "max_tokens": 5,
-                "stream": False
-            },
-            timeout=5
+            json=payload,
+            timeout=timeout
         )
-        if response.status_code == 429 or response.status_code == 200:
-            return True
-        else:
-            return False
+        return response.status_code in [200, 429]
     except requests.exceptions.RequestException as e:
         logging.error(
-            f"测试模型 {model_name} 可用性失败，"
+            f"测试{model_type}模型 {model_name} 可用性失败，"
             f"API Key：{api_key}，错误信息：{e}"
         )
         return False
@@ -176,188 +184,109 @@ def create_base64_markdown_image(image_url):
         return None
 
 def refresh_models():
-    global text_models, free_text_models
-    global embedding_models, free_embedding_models
-    global image_models, free_image_models
-
-    text_models = get_all_models(FREE_MODEL_TEST_KEY, "chat")
-    embedding_models = get_all_models(FREE_MODEL_TEST_KEY, "embedding")
-    image_models = get_all_models(FREE_MODEL_TEST_KEY, "text-to-image")
-    free_text_models = []
-    free_embedding_models = []
-    free_image_models = []
+    global models
+    
+    # 获取各类型模型列表
+    models["text"] = get_all_models(FREE_MODEL_TEST_KEY, "chat")
+    models["embedding"] = get_all_models(FREE_MODEL_TEST_KEY, "embedding")
+    models["image"] = get_all_models(FREE_MODEL_TEST_KEY, "text-to-image")
+    
+    # 重置免费模型列表
+    models["free_text"] = []
+    models["free_embedding"] = []
+    models["free_image"] = []
 
-    ban_models_str = os.environ.get("BAN_MODELS")
+    # 处理禁用模型
     ban_models = []
+    ban_models_str = os.environ.get("BAN_MODELS")
     if ban_models_str:
         try:
             ban_models = json.loads(ban_models_str)
             if not isinstance(ban_models, list):
-                logging.warning(
-                    "环境变量 BAN_MODELS 格式不正确，应为 JSON 数组。"
-                )
+                logging.warning("环境变量 BAN_MODELS 格式不正确，应为 JSON 数组。")
                 ban_models = []
         except json.JSONDecodeError:
-            logging.warning(
-                "环境变量 BAN_MODELS JSON 解析失败，请检查格式。"
-            )
-            ban_models = []
+            logging.warning("环境变量 BAN_MODELS JSON 解析失败，请检查格式。")
     
-    text_models = [model for model in text_models if model not in ban_models]
-    embedding_models = [model for model in embedding_models if model not in ban_models]
-    image_models = [model for model in image_models if model not in ban_models]
-
-    with concurrent.futures.ThreadPoolExecutor(
-        max_workers=10000
-    ) as executor:
-        future_to_model = {
-            executor.submit(
-                test_model_availability,
-                FREE_MODEL_TEST_KEY,
-                model
-            ): model for model in text_models
-        }
-        for future in concurrent.futures.as_completed(future_to_model):
-            model = future_to_model[future]
-            try:
-                is_free = future.result()
-                if is_free:
-                    free_text_models.append(model)
-            except Exception as exc:
-                logging.error(f"模型 {model} 测试生成异常: {exc}")
-
-    with concurrent.futures.ThreadPoolExecutor(
-        max_workers=10000
-    ) as executor:
-        future_to_model = {
-            executor.submit(
-                test_embedding_model_availability,
-                FREE_MODEL_TEST_KEY, model
-            ): model for model in embedding_models
-        }
-        for future in concurrent.futures.as_completed(future_to_model):
-            model = future_to_model[future]
-            try:
-                is_free = future.result()
-                if is_free:
-                    free_embedding_models.append(model)
-            except Exception as exc:
-                logging.error(f"模型 {model} 测试生成异常: {exc}")
-                
-    with concurrent.futures.ThreadPoolExecutor(
-        max_workers=10000
-    ) as executor:
-        future_to_model = {
-            executor.submit(
-                test_image_model_availability,
-                FREE_MODEL_TEST_KEY, model
-            ): model for model in image_models
-        }
-        for future in concurrent.futures.as_completed(future_to_model):
-            model = future_to_model[future]
-            try:
-                is_free = future.result()
-                if is_free:
-                    free_image_models.append(model)
-            except Exception as exc:
-                logging.error(f"模型 {model} 测试生成异常: {exc}")
-
-    logging.info(f"所有文本模型列表：{text_models}")
-    logging.info(f"免费文本模型列表：{free_text_models}")
-    logging.info(f"所有向量模型列表：{embedding_models}")
-    logging.info(f"免费向量模型列表：{free_embedding_models}")
-    logging.info(f"所有生图模型列表：{image_models}")
-    logging.info(f"免费生图模型列表：{free_image_models}")
-
-def test_embedding_model_availability(api_key, model_name):
-    headers = {
-        "Authorization": f"Bearer {api_key}",
-        "Content-Type": "application/json"
-    }
-    try:
-        response = session.post(
-            EMBEDDINGS_ENDPOINT,
-            headers=headers,
-            json={
-                "model": model_name,
-                "input": ["hi"],
-            },
-            timeout=10
-        )
-        if response.status_code == 429 or response.status_code == 200:
-            return True
-        else:
-            return False
-    except requests.exceptions.RequestException as e:
-        logging.error(
-            f"测试向量模型 {model_name} 可用性失败，"
-            f"API Key：{api_key}，错误信息：{e}"
-        )
-        return False
-        
-def test_image_model_availability(api_key, model_name):
-    return model_name in FREE_IMAGE_LIST
-
-def load_keys():
-    keys_str = os.environ.get("KEYS")
-    test_model = os.environ.get(
-        "TEST_MODEL",
-        "Pro/google/gemma-2-9b-it"
-    )
-
-    if keys_str:
-        keys = [key.strip() for key in keys_str.split(',')]
-        unique_keys = list(set(keys))
-        keys_str = ','.join(unique_keys) 
-        os.environ["KEYS"] = keys_str
-
-        logging.info(f"加载的 keys：{unique_keys}")
-
-        with concurrent.futures.ThreadPoolExecutor(
-            max_workers=10000
-        ) as executor:
-            future_to_key = {
+    # 过滤禁用模型
+    models["text"] = [model for model in models["text"] if model not in ban_models]
+    models["embedding"] = [model for model in models["embedding"] if model not in ban_models]
+    models["image"] = [model for model in models["image"] if model not in ban_models]
+
+    # 使用统一的测试函数测试各类型模型
+    model_types = [
+        ("text", "chat"),
+        ("embedding", "embedding"),
+        ("image", "image")
+    ]
+    
+    for model_type, test_type in model_types:
+        with concurrent.futures.ThreadPoolExecutor(max_workers=10000) as executor:
+            future_to_model = {
                 executor.submit(
-                    process_key, key, test_model
-                ): key for key in unique_keys
+                    test_model_availability,
+                    FREE_MODEL_TEST_KEY,
+                    model,
+                    test_type
+                ): model for model in models[model_type]
             }
-
-            invalid_keys = []
-            free_keys = []
-            unverified_keys = []
-            valid_keys = []
-
-            for future in concurrent.futures.as_completed(
-                future_to_key
-            ):
-                key = future_to_key[future]
+            
+            for future in concurrent.futures.as_completed(future_to_model):
+                model = future_to_model[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)
+                    is_free = future.result()
+                    if is_free:
+                        models[f"free_{model_type}"].append(model)
                 except Exception as exc:
-                    logging.error(f"处理 KEY {key} 生成异常: {exc}")
+                    logging.error(f"{model_type}模型 {model} 测试生成异常: {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}")
+    # 记录日志
+    for model_type in ["text", "embedding", "image"]:
+        logging.info(f"所有{model_type}模型列表：{models[model_type]}")
+        logging.info(f"免费{model_type}模型列表：{models[f'free_{model_type}']}")
 
-        global invalid_keys_global, free_keys_global
-        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:
+def load_keys():
+    global key_status
+    # 重置key状态
+    for status in key_status:
+        key_status[status] = []
+        
+    keys_str = os.environ.get("KEYS")
+    if not keys_str:
         logging.warning("环境变量 KEYS 未设置。")
+        return
+
+    test_model = os.environ.get("TEST_MODEL", "Pro/google/gemma-2-9b-it")
+    unique_keys = list(set(key.strip() for key in keys_str.split(',')))
+    os.environ["KEYS"] = ','.join(unique_keys)
+
+    logging.info(f"加载的 keys：{unique_keys}")
+
+    def process_key_with_logging(key):
+        try:
+            key_type = process_key(key, test_model)
+            if key_type in key_status:
+                key_status[key_type].append(key)
+            return key_type
+        except Exception as exc:
+            logging.error(f"处理 KEY {key} 生成异常: {exc}")
+            return "invalid"
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=10000) as executor:
+        futures = [executor.submit(process_key_with_logging, key) for key in unique_keys]
+        concurrent.futures.wait(futures)
+
+    # 记录每种状态的keys
+    for status, keys in key_status.items():
+        logging.info(f"{status.capitalize()} KEYS: {keys}")
+    
+    # 更新全局变量
+    global invalid_keys_global, free_keys_global, unverified_keys_global, valid_keys_global
+    invalid_keys_global = key_status["invalid"]
+    free_keys_global = key_status["free"]
+    unverified_keys_global = key_status["unverified"]
+    valid_keys_global = key_status["valid"]
 
 def process_key(key, test_model):
     credit_summary = get_credit_summary(key)
@@ -525,7 +454,14 @@ def list_models():
 
     detailed_models = []
     
-    for model in chain(text_models, embedding_models, image_models):
+    # 合并所有类型的模型
+    all_models = chain(
+        models["text"], 
+        models["embedding"], 
+        models["image"]
+    )
+    
+    for model in all_models:
         detailed_models.append({
             "id": model,
             "object": "model",
@@ -596,14 +532,14 @@ def handsome_embeddings():
     data = request.get_json()
     if not data or 'model' not in data:
         return jsonify({"error": "Invalid request data"}), 400
-    if data['model'] not in embedding_models:
+    if data['model'] not in models["embedding"]:
         return jsonify({"error": "Invalid model"}), 400
 
     model_name = data['model']
     request_type = determine_request_type(
         model_name,
-        embedding_models,
-        free_embedding_models
+        models["embedding"],
+        models["free_embedding"]
     )
     api_key = select_key(request_type, model_name)
 
@@ -677,15 +613,15 @@ def handsome_images_generations():
     data = request.get_json()
     if not data or 'model' not in data:
         return jsonify({"error": "Invalid request data"}), 400
-    if data['model'] not in image_models:
+    if data['model'] not in models["image"]:
         return jsonify({"error": "Invalid model"}), 400
 
     model_name = data.get('model')
     
     request_type = determine_request_type(
         model_name,
-        image_models,
-        free_image_models
+        models["image"],
+        models["free_image"]
     )
     
     api_key = select_key(request_type, model_name)
@@ -848,15 +784,15 @@ def handsome_chat_completions():
     data = request.get_json()
     if not data or 'model' not in data:
         return jsonify({"error": "Invalid request data"}), 400
-    if data['model'] not in text_models and data['model'] not in image_models:
+    if data['model'] not in models["text"] and data['model'] not in models["image"]:
         return jsonify({"error": "Invalid model"}), 400
 
     model_name = data['model']
     
     request_type = determine_request_type(
         model_name,
-        text_models + image_models,
-        free_text_models + free_image_models
+        models["text"] + models["image"],
+        models["free_text"] + models["free_image"]
     )
     
     api_key = select_key(request_type, model_name)
@@ -877,7 +813,7 @@ def handsome_chat_completions():
         "Content-Type": "application/json"
     }
     
-    if model_name in image_models:
+    if model_name in models["image"]:
         user_content = ""
         messages = data.get("messages", [])
         for message in messages:
@@ -1175,7 +1111,7 @@ def handsome_chat_completions():
                 headers=headers,
                 json=data,
                 stream=data.get("stream", False),
-                timeout=60
+                timeout=600
             )
 
             if response.status_code == 429:
@@ -1388,4 +1324,4 @@ if __name__ == '__main__':
         debug=False,
         host='0.0.0.0',
         port=int(os.environ.get('PORT', 7860))
-    )
+    )