Update app.py

app.py

@@ -54,10 +54,15 @@ def get_credit_summary(api_key):
         logging.error(f"解析额度信息失败，API Key：{api_key}，错误信息：{e}")
         return None
     except ValueError as e:
-        logging.error(f"total_balance 无法转换为浮点数，API Key：{api_key}，错误信息：{e}")
+        logging.error(
+            f"total_balance 无法转换为浮点数，"
+            f"API Key：{api_key}，错误信息：{e}"
+        )
         return None
 
-FREE_MODEL_TEST_KEY = "sk-bmjbjzleaqfgtqfzmcnsbagxrlohriadnxqrzfocbizaxukw"
+FREE_MODEL_TEST_KEY = (
+    "sk-bmjbjzleaqfgtqfzmcnsbagxrlohriadnxqrzfocbizaxukw"
+)
 
 def test_model_availability(api_key, model_name):
     """
@@ -68,37 +73,50 @@ def test_model_availability(api_key, model_name):
         "Content-Type": "application/json"
     }
     try:
-        response = requests.post(TEST_MODEL_ENDPOINT,
-                                 headers=headers,
-                                 json={
-                                     "model": model_name,
-                                     "messages": [{"role": "user", "content": "hi"}],
-                                     "max_tokens": 5,
-                                     "stream": False
-                                 },
-                                 timeout=10)
+        response = requests.post(
+            TEST_MODEL_ENDPOINT,
+            headers=headers,
+            json={
+                "model": model_name,
+                "messages": [{"role": "user", "content": "hi"}],
+                "max_tokens": 5,
+                "stream": False
+            },
+            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} 可用性失败，API Key：{api_key}，错误信息：{e}")
+        logging.error(
+            f"测试模型 {model_name} 可用性失败，"
+            f"API Key：{api_key}，错误信息：{e}"
+        )
         return False
 
 def refresh_models():
     """
     刷新模型列表和免费模型列表。
     """
-    global all_models, free_models, embedding_models, free_embedding_models
+    global all_models, free_models
+    global embedding_models, free_embedding_models
 
     all_models = get_all_models(FREE_MODEL_TEST_KEY, "chat")
     embedding_models = get_all_models(FREE_MODEL_TEST_KEY, "embedding")
     free_models = []
     free_embedding_models = []
 
-    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}
+    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,10 +126,15 @@ def refresh_models():
             except Exception as exc:
                 logging.error(f"模型 {model} 测试生成异常: {exc}")
 
-    with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
+    with concurrent.futures.ThreadPoolExecutor(
+        max_workers=10
+    ) as executor:
         future_to_model = {
-            executor.submit(test_embedding_model_availability, FREE_MODEL_TEST_KEY, model): model for model in
-            embedding_models}
+            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:
@@ -135,42 +158,64 @@ def test_embedding_model_availability(api_key, model_name):
         "Content-Type": "application/json"
     }
     try:
-        response = requests.post(EMBEDDINGS_ENDPOINT,
-                                 headers=headers,
-                                 json={
-                                     "model": model_name,
-                                     "input": ["hi"],
-                                 },
-                                 timeout=10)
+        response = requests.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} 可用性失败，API Key：{api_key}，错误信息：{e}")
+        logging.error(
+            f"测试向量模型 {model_name} 可用性失败，"
+            f"API Key：{api_key}，错误信息：{e}"
+        )
         return False
 
 def load_keys():
     """
-    从环境变量中加载 keys，并根据额度和模型可用性进行分类，然后记录到日志中。
+    从环境变量中加载 keys，进行去重，
+    并根据额度和模型可用性进行分类，
+    然后记录到日志中。
     使用线程池并发处理每个 key。
     """
     keys_str = os.environ.get("KEYS")
-    test_model = os.environ.get("TEST_MODEL", "Pro/google/gemma-2-9b-it")
+    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(',')]
-        logging.info(f"加载的 keys：{keys}")
-
-        with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
-            future_to_key = {executor.submit(process_key, key, test_model): key for key in keys}
+        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=20
+        ) as executor:
+            future_to_key = {
+                executor.submit(
+                    process_key, key, test_model
+                ): key for key in unique_keys
+            }
 
             invalid_keys = []
             free_keys = []
             unverified_keys = []
             valid_keys = []
 
-            for future in concurrent.futures.as_completed(future_to_key):
+            for future in concurrent.futures.as_completed(
+                future_to_key
+            ):
                 key = future_to_key[future]
                 try:
                     key_type = future.result()
@@ -190,7 +235,8 @@ def load_keys():
         logging.info(f"未实名 KEY：{unverified_keys}")
         logging.info(f"有效 KEY：{valid_keys}")
 
-        global invalid_keys_global, free_keys_global, unverified_keys_global, valid_keys_global
+        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
@@ -225,20 +271,36 @@ def get_all_models(api_key, sub_type):
         "Content-Type": "application/json"
     }
     try:
-        response = requests.get(MODELS_ENDPOINT, headers=headers, params={"sub_type": sub_type})
+        response = requests.get(
+            MODELS_ENDPOINT,
+            headers=headers,
+            params={"sub_type": sub_type}
+        )
         response.raise_for_status()
         data = response.json()
-        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]
+        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}")
+        logging.error(
+            f"获取模型列表失败，"
+            f"API Key：{api_key}，错误信息：{e}"
+        )
         return []
     except (KeyError, TypeError) as e:
-        logging.error(f"解析模型列表失败，API Key：{api_key}，错误信息：{e}")
+        logging.error(
+            f"解析模型列表失败，"
+            f"API Key：{api_key}，错误信息：{e}"
+        )
         return []
 
 def determine_request_type(model_name, model_list, free_model_list):
@@ -254,14 +316,23 @@ def determine_request_type(model_name, model_list, free_model_list):
 
 def select_key(request_type, model_name):
     """
-    根据请求类型和模型名称选择合适的 KEY，并实现轮询和重试机制。
+    根据请求类型和模型名称选择合适的 KEY，
+    并实现轮询和重试机制。
     """
     if request_type == "free":
-        available_keys = free_keys_global + unverified_keys_global + valid_keys_global
+        available_keys = (
+            free_keys_global +
+            unverified_keys_global +
+            valid_keys_global
+        )
     elif request_type == "paid":
         available_keys = unverified_keys_global + valid_keys_global
     else:
-        available_keys = free_keys_global + unverified_keys_global + valid_keys_global
+        available_keys = (
+            free_keys_global +
+            unverified_keys_global +
+            valid_keys_global
+        )
 
     if not available_keys:
         return None
@@ -276,14 +347,17 @@ def select_key(request_type, model_name):
             model_key_indices[model_name] = current_index
             return key
         else:
-            logging.warning(f"KEY {key} 无效或达到限制，尝试下一个 KEY")
+            logging.warning(
+                f"KEY {key} 无效或达到限制，尝试下一个 KEY"
+            )
 
     model_key_indices[model_name] = 0
     return None
 
 def key_is_valid(key, request_type):
     """
-    检查 KEY 是否有效，根据不同的请求类型进行不同的检查。
+    检查 KEY 是否有效，
+    根据不同的请求类型进行不同的检查。
     """
     if request_type == "invalid":
         return False
@@ -303,7 +377,8 @@ def key_is_valid(key, request_type):
 
 def check_authorization(request):
     """
-    检查请求头中的 Authorization 字段是否匹配环境变量 AUTHORIZATION_KEY。
+    检查请求头中的 Authorization 字段
+    是否匹配环境变量 AUTHORIZATION_KEY。
     """
     authorization_key = os.environ.get("AUTHORIZATION_KEY")
     if not authorization_key:
@@ -323,7 +398,8 @@ def check_authorization(request):
 
 scheduler = BackgroundScheduler()
 scheduler.add_job(load_keys, 'interval', hours=1)
-scheduler.add_job(refresh_models, 'interval', minutes=10)
+scheduler.remove_all_jobs()
+scheduler.add_job(refresh_models, 'interval', hours=1)
 
 @app.route('/')
 def index():
@@ -332,10 +408,19 @@ def index():
 @app.route('/check_tokens', methods=['POST'])
 def check_tokens():
     tokens = request.json.get('tokens', [])
-    test_model = os.environ.get("TEST_MODEL", "Pro/google/gemma-2-9b-it")
-
-    with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
-        future_to_token = {executor.submit(process_key, token, test_model): token for token in tokens}
+    test_model = os.environ.get(
+        "TEST_MODEL",
+        "Pro/google/gemma-2-9b-it"
+    )
+
+    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):
@@ -343,20 +428,50 @@ def check_tokens():
             try:
                 key_type = future.result()
                 credit_summary = get_credit_summary(token)
-                balance = credit_summary.get("total_balance", 0) if credit_summary else 0
+                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": "无法获取额度信息"})
+                        {
+                            "token": token,
+                            "type": "无效 KEY",
+                            "balance": balance,
+                            "message": "无法获取额度信息"
+                        }
+                    )
                 elif key_type == "free":
-                    results.append({"token": token, "type": "免费 KEY", "balance": balance, "message": "额度不足"})
+                    results.append(
+                        {
+                            "token": token,
+                            "type": "免费 KEY",
+                            "balance": balance,
+                            "message": "额度不足"
+                        }
+                    )
                 elif key_type == "unverified":
                     results.append(
-                        {"token": token, "type": "未实名 KEY", "balance": balance, "message": "无法使用指定模型"})
+                        {
+                            "token": token,
+                            "type": "未实名 KEY",
+                            "balance": balance,
+                            "message": "无法使用指定模型"
+                        }
+                    )
                 elif key_type == "valid":
                     results.append(
-                        {"token": token, "type": "有效 KEY", "balance": balance, "message": "可以使用指定模型"})
+                        {
+                            "token": token,
+                            "type": "有效 KEY",
+                            "balance": balance,
+                            "message": "可以使用指定模型"
+                        }
+                    )
             except Exception as exc:
-                logging.error(f"处理 Token {token} 生成异常: {exc}")
+                logging.error(
+                    f"处理 Token {token} 生成异常: {exc}"
+                )
 
     return jsonify(results)
 
@@ -370,12 +485,23 @@ def handsome_chat_completions():
         return jsonify({"error": "Invalid request data"}), 400
 
     model_name = data['model']
-    request_type = determine_request_type(model_name, all_models, free_models)
+    request_type = determine_request_type(
+        model_name,
+        all_models,
+        free_models
+    )
     api_key = select_key(request_type, model_name)
 
     if not api_key:
         return jsonify(
-            {"error": "No available API key for this request type or all keys have reached their limits"}), 429
+            {
+                "error": (
+                    "No available API key for this "
+                    "request type or all keys have "
+                    "reached their limits"
+                )
+            }
+        ), 429
 
     headers = {
         "Authorization": f"Bearer {api_key}",
@@ -407,7 +533,10 @@ def handsome_chat_completions():
                         yield chunk
 
                 end_time = time.time()
-                first_token_time = first_chunk_time - start_time if first_chunk_time else 0
+                first_token_time = (
+                    first_chunk_time - start_time
+                    if first_chunk_time else 0
+                )
                 total_time = end_time - start_time
 
                 prompt_tokens = 0
@@ -421,19 +550,43 @@ def handsome_chat_completions():
                         try:
                             response_json = json.loads(line)
 
-                            if "usage" in response_json and "completion_tokens" in response_json["usage"]:
-                                completion_tokens = response_json["usage"]["completion_tokens"]
-
-                            if "choices" in response_json and len(response_json["choices"]) > 0 and "delta" in \
-                                    response_json["choices"][0] and "content" in response_json["choices"][0][
-                                "delta"]:
-                                response_content += response_json["choices"][0]["delta"]["content"]
-
-                            if "usage" in response_json and "prompt_tokens" in response_json["usage"]:
-                                prompt_tokens = response_json["usage"]["prompt_tokens"]
-
-                        except (KeyError, ValueError, IndexError) as e:
-                            logging.error(f"解析流式响应单行 JSON 失败: {e}, 行内容: {line}")
+                            if (
+                                "usage" in response_json and
+                                "completion_tokens" in response_json["usage"]
+                            ):
+                                completion_tokens = response_json[
+                                    "usage"
+                                ]["completion_tokens"]
+
+                            if (
+                                "choices" in response_json and
+                                len(response_json["choices"]) > 0 and
+                                "delta" in response_json["choices"][0] and
+                                "content" in response_json[
+                                    "choices"
+                                ][0]["delta"]
+                            ):
+                                response_content += response_json[
+                                    "choices"
+                                ][0]["delta"]["content"]
+
+                            if (
+                                "usage" in response_json and
+                                "prompt_tokens" in response_json["usage"]
+                            ):
+                                prompt_tokens = response_json[
+                                    "usage"
+                                ]["prompt_tokens"]
+
+                        except (
+                            KeyError,
+                            ValueError,
+                            IndexError
+                        ) as e:
+                            logging.error(
+                                f"解析流式响应单行 JSON 失败: {e}, "
+                                f"行内容: {line}"
+                            )
 
                 user_content = ""
                 messages = data.get("messages", [])
@@ -443,19 +596,39 @@ def handsome_chat_completions():
                             user_content += message["content"] + " "
                         elif isinstance(message["content"], list):
                             for item in message["content"]:
-                                if isinstance(item, dict) and item.get("type") == "text":
-                                    user_content += item.get("text", "") + " "
+                                if (
+                                    isinstance(item, dict) and
+                                    item.get("type") == "text"
+                                ):
+                                    user_content += (
+                                        item.get("text", "") +
+                                        " "
+                                    )
 
                 user_content = user_content.strip()
 
-                user_content_replaced = user_content.replace('\n', '\\n').replace('\r', '\\n')
-                response_content_replaced = response_content.replace('\n', '\\n').replace('\r', '\\n')
+                user_content_replaced = user_content.replace(
+                    '\n', '\\n'
+                ).replace('\r', '\\n')
+                response_content_replaced = response_content.replace(
+                    '\n', '\\n'
+                ).replace('\r', '\\n')
 
                 logging.info(
-                    f"使用的key: {api_key}, 提示token: {prompt_tokens}, 输出token: {completion_tokens}, 首字用时: {first_token_time:.4f}秒, 总共用时: {total_time:.4f}秒, 使用的模型: {model_name}, 用户的内容: {user_content_replaced}, 输出的内容: {response_content_replaced}"
+                    f"使用的key: {api_key}, "
+                    f"提示token: {prompt_tokens}, "
+                    f"输出token: {completion_tokens}, "
+                    f"首字用时: {first_token_time:.4f}秒, "
+                    f"总共用时: {total_time:.4f}秒, "
+                    f"使用的模型: {model_name}, "
+                    f"用户的内容: {user_content_replaced}, "
+                    f"输出的内容: {response_content_replaced}"
                 )
 
-            return Response(stream_with_context(generate()), content_type=response.headers['Content-Type'])
+            return Response(
+                stream_with_context(generate()),
+                content_type=response.headers['Content-Type']
+            )
         else:
             response.raise_for_status()
             end_time = time.time()
@@ -464,10 +637,17 @@ def handsome_chat_completions():
 
             try:
                 prompt_tokens = response_json["usage"]["prompt_tokens"]
-                completion_tokens = response_json["usage"]["completion_tokens"]
-                response_content = response_json["choices"][0]["message"]["content"]
+                completion_tokens = response_json[
+                    "usage"
+                ]["completion_tokens"]
+                response_content = response_json[
+                    "choices"
+                ][0]["message"]["content"]
             except (KeyError, ValueError, IndexError) as e:
-                logging.error(f"解析非流式响应 JSON 失败: {e}, 完整内容: {response_json}")
+                logging.error(
+                    f"解析非流式响应 JSON 失败: {e}, "
+                    f"完整内容: {response_json}"
+                )
                 prompt_tokens = 0
                 completion_tokens = 0
                 response_content = ""
@@ -480,16 +660,32 @@ def handsome_chat_completions():
                         user_content += message["content"] + " "
                     elif isinstance(message["content"], list):
                         for item in message["content"]:
-                            if isinstance(item, dict) and item.get("type") == "text":
-                                user_content += item.get("text", "") + " "
+                            if (
+                                isinstance(item, dict) and
+                                item.get("type") == "text"
+                            ):
+                                user_content += (
+                                    item.get("text", "") + " "
+                                )
 
             user_content = user_content.strip()
 
-            user_content_replaced = user_content.replace('\n', '\\n').replace('\r', '\\n')
-            response_content_replaced = response_content.replace('\n', '\\n').replace('\r', '\\n')
+            user_content_replaced = user_content.replace(
+                '\n', '\\n'
+            ).replace('\r', '\\n')
+            response_content_replaced = response_content.replace(
+                '\n', '\\n'
+            ).replace('\r', '\\n')
 
             logging.info(
-                f"使用的key: {api_key}, 提示token: {prompt_tokens}, 输出token: {completion_tokens}, 首字用时: 0, 总共用时: {total_time:.4f}秒, 使用的模型: {model_name}, 用户的内容: {user_content_replaced}, 输出的内容: {response_content_replaced}"
+                f"使用的key: {api_key}, "
+                f"提示token: {prompt_tokens}, "
+                f"输出token: {completion_tokens}, "
+                f"首字用时: 0, "
+                f"总共用时: {total_time:.4f}秒, "
+                f"使用的模型: {model_name}, "
+                f"用户的内容: {user_content_replaced}, "
+                f"输出的内容: {response_content_replaced}"
             )
             return jsonify(response_json)
 
@@ -538,14 +734,21 @@ def get_billing_info():
     keys = valid_keys_global + unverified_keys_global
     total_balance = 0
 
-    with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
-        futures = [executor.submit(get_credit_summary, key) for key in keys]
+    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)
+                    total_balance += credit_summary.get(
+                        "total_balance",
+                        0
+                    )
             except Exception as exc:
                 logging.error(f"获取额度信息生成异常: {exc}")
 
@@ -574,7 +777,10 @@ def billing_usage():
         "total_usage": 0
     })
 
-@app.route('/handsome/v1/dashboard/billing/subscription', methods=['GET'])
+@app.route(
+    '/handsome/v1/dashboard/billing/subscription',
+    methods=['GET']
+)
 def billing_subscription():
     if not check_authorization(request):
         return jsonify({"error": "Unauthorized"}), 401
@@ -616,12 +822,23 @@ def handsome_embeddings():
         return jsonify({"error": "Invalid request data"}), 400
 
     model_name = data['model']
-    request_type = determine_request_type(model_name, embedding_models, free_embedding_models)
+    request_type = determine_request_type(
+        model_name,
+        embedding_models,
+        free_embedding_models
+    )
     api_key = select_key(request_type, model_name)
 
     if not api_key:
         return jsonify(
-            {"error": "No available API key for this request type or all keys have reached their limits"}), 429
+            {
+                "error": (
+                    "No available API key for this "
+                    "request type or all keys have "
+                    "reached their limits"
+                )
+            }
+        ), 429
 
     headers = {
         "Authorization": f"Bearer {api_key}",
@@ -634,7 +851,7 @@ def handsome_embeddings():
             EMBEDDINGS_ENDPOINT,
             headers=headers,
             json=data,
-            timeout=60
+            timeout=120
         )
 
         if response.status_code == 429:
@@ -649,12 +866,18 @@ def handsome_embeddings():
             prompt_tokens = response_json["usage"]["prompt_tokens"]
             embedding_data = response_json["data"]
         except (KeyError, ValueError, IndexError) as e:
-            logging.error(f"解析响应 JSON 失败: {e}, 完整内容: {response_json}")
+            logging.error(
+                f"解析响应 JSON 失败: {e}, "
+                f"完整内容: {response_json}"
+            )
             prompt_tokens = 0
             embedding_data = []
 
         logging.info(
-            f"使用的key: {api_key}, 提示token: {prompt_tokens}, 总共用时: {total_time:.4f}秒, 使用的模型: {model_name}"
+            f"使用的key: {api_key}, "
+            f"提示token: {prompt_tokens}, "
+            f"总共用时: {total_time:.4f}秒, "
+            f"使用的模型: {model_name}"
         )
 
         return jsonify({
@@ -679,12 +902,20 @@ if __name__ == '__main__':
     unverified_keys_global = []
     valid_keys_global = []
 
+    # 启动时先调用一次 load_keys()
+    load_keys()
+    logging.info("程序启动时首次加载 keys 已执行")
+
     scheduler.start()
 
-    load_keys()
+    # load_keys()
     logging.info("首次加载 keys 已手动触发执行")
 
     refresh_models()
     logging.info("首次刷新模型列表已手动触发执行")
 
-    app.run(debug=False, host='0.0.0.0', port=int(os.environ.get('PORT', 7860)))
+    app.run(
+        debug=False,
+        host='0.0.0.0',
+        port=int(os.environ.get('PORT', 7860))
+    )