Update app.py

app.py

@@ -546,9 +546,197 @@ def check_tokens():
 
     return jsonify(results)
 
-@app.route('/handsome/v1/chat/completions', methods=['POST'])
-def handsome_chat_completions():
-    if not check_authorization(request):
+import os
+import time
+import logging
+import requests
+import json
+import uuid
+from flask import Flask, request, jsonify
+
+API_ENDPOINT = "https://api.siliconflow.cn/v1/user/info"
+TEST_MODEL_ENDPOINT = "https://api.siliconflow.cn/v1/chat/completions"
+EMBEDDINGS_ENDPOINT = "https://api.siliconflow.cn/v1/embeddings"
+
+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 = []
+model_key_indices = {}
+
+request_timestamps = []
+token_counts = []
+data_lock = None
+
+FREE_MODEL_TEST_KEY = (
+    "sk-bmjbjzleaqfgtqfzmcnsbagxrlohriadnxqrzfocbizaxukw"
+)
+
+FREE_IMAGE_LIST = [
+    "stabilityai/stable-diffusion-3-5-large",
+    "black-forest-labs/FLUX.1-schnell",
+    "stabilityai/stable-diffusion-3-medium",
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    "stabilityai/stable-diffusion-2-1"
+]
+
+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
+
+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": 5,
+                "stream": False
+            },
+            timeout=5
+        )
+        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):
+    """
+    测试指定的图像模型是否在 FREE_IMAGE_LIST 中。
+    如果在列表中，返回 True，否则返回 False。
+    """
+    return model_name in FREE_IMAGE_LIST
+
+def determine_request_type(model_name, model_list, free_model_list):
+    """
+    根据用户请求的模型判断请求类型。
+    """
+    if model_name in free_model_list:
+        return "free"
+    elif model_name in model_list:
+        return "paid"
+    else:
+        return "unknown"
+
+def select_key(request_type, model_name):
+    """
+    根据请求类型和模型名称选择合适的 KEY，
+    并实现轮询和重试机制。
+    """
+    if request_type == "free":
+        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
+        )
+
+    if not available_keys:
+        return None
+
+    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
+
+        if key_is_valid(key, request_type):
+            model_key_indices[model_name] = current_index
+            return key
+        else:
+            logging.warning(
+                f"KEY {key} 无效或达到限制，尝试下一个 KEY"
+            )
+
+    model_key_indices[model_name] = 0
+    return None
+
+def key_is_valid(key, request_type):
+    """
+    检查 KEY 是否有效，
+    根据不同的请求类型进行不同的检查。
+    """
+    if request_type == "invalid":
+        return 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
+    elif request_type == "paid" or request_type == "unverified":
+        return total_balance > 0
+    else:
+        return False
+
+def check_authorization(request, authorization_key):
+    """
+    检查请求头中的 Authorization 字段
+    是否匹配环境变量 AUTHORIZATION_KEY。
+    """
+    if not authorization_key:
+        logging.warning("环境变量 AUTHORIZATION_KEY 未设置，请设置后重试。")
+        return False
+
+    auth_header = request.headers.get('Authorization')
+    if not auth_header:
+        logging.warning("请求头中缺少 Authorization 字段。")
+        return False
+
+    if auth_header != f"Bearer {authorization_key}":
+        logging.warning(f"无效的 Authorization 密钥：{auth_header}")
+        return False
+
+    return True
+
+def handsome_chat_completions(request, authorization_key):
+    if not check_authorization(request, authorization_key):
         return jsonify({"error": "Unauthorized"}), 401
 
     data = request.get_json()
@@ -656,7 +844,8 @@ def handsome_chat_completions():
                         "index": 0,
                         "message": {
                             "role": "assistant",
-                            "content": image_url if image_url else "Failed to generate image", # Directly return the URL in content
+                            "content": None,
+                            "image_url": image_url if image_url else None, # Return URL in image_url
                         },
                         "finish_reason": "length",
                      }
@@ -667,7 +856,9 @@ def handsome_chat_completions():
                         "total_tokens": 0
                     }
                 }
-
+                if not image_url:
+                    response_data["choices"][0]["message"]["content"] = "Failed to generate image"
+                
             except (KeyError, ValueError, IndexError) as e:
                 logging.error(
                     f"解析响应 JSON 失败: {e}, "
@@ -685,6 +876,7 @@ def handsome_chat_completions():
                         "message": {
                             "role": "assistant",
                             "content": "Failed to process image data",
+                            "image_url": None,
                         },
                         "finish_reason": "stop",
                       }
@@ -702,10 +894,6 @@ def handsome_chat_completions():
                 f"使用的模型: {model_name}"
             )
 
-            with data_lock:
-                request_timestamps.append(time.time())
-                token_counts.append(0)  # Image generation doesn't use tokens
-
             return jsonify(response_data)
         except requests.exceptions.RequestException as e:
             logging.error(f"请求转发异常: {e}")
@@ -718,191 +906,69 @@ def handsome_chat_completions():
                 TEST_MODEL_ENDPOINT,
                 headers=headers,
                 json=data,
-                stream=data.get("stream", False),
+                stream=False,
                 timeout=60
             )
             if response.status_code == 429:
                 return jsonify(response.json()), 429
-
-            if data.get("stream", False):
-                def generate():
-                    first_chunk_time = None
-                    full_response_content = ""
-                    for chunk in response.iter_content(chunk_size=1024):
-                        if chunk:
-                            if first_chunk_time is None:
-                                first_chunk_time = time.time()
-                            full_response_content += chunk.decode("utf-8")
-                            yield chunk
-
-                    end_time = time.time()
-                    first_token_time = (
-                        first_chunk_time - start_time
-                        if first_chunk_time else 0
-                    )
-                    total_time = end_time - start_time
-
-                    prompt_tokens = 0
-                    completion_tokens = 0
-                    response_content = ""
-                    for line in full_response_content.splitlines():
-                        if line.startswith("data:"):
-                            line = line[5:].strip()
-                            if line == "[DONE]":
-                                continue
-                            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}, "
-                                    f"行内容: {line}"
+            
+            response.raise_for_status()
+            end_time = time.time()
+            response_json = response.json()
+            total_time = end_time - start_time
+            try:
+                prompt_tokens = response_json["usage"]["prompt_tokens"]
+                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}, "
+                    f"完整内容: {response_json}"
+                )
+                prompt_tokens = 0
+                completion_tokens = 0
+                response_content = ""
+            user_content = ""
+            messages = data.get("messages", [])
+            for message in messages:
+                if message["role"] == "user":
+                    if isinstance(message["content"], str):
+                        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", "") +
+                                    " "
                                 )
 
-                    user_content = ""
-                    messages = data.get("messages", [])
-                    for message in messages:
-                        if message["role"] == "user":
-                            if isinstance(message["content"], str):
-                                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", "") +
-                                            " "
-                                        )
-
-                    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')
-
-                    logging.info(
-                        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}"
-                    )
-
-                    with data_lock:
-                        request_timestamps.append(time.time())
-                        token_counts.append(prompt_tokens+completion_tokens)
+            user_content = user_content.strip()
 
-                return Response(
-                    stream_with_context(generate()),
-                    content_type=response.headers['Content-Type']
-                )
-            else:
-                response.raise_for_status()
-                end_time = time.time()
-                response_json = response.json()
-                total_time = end_time - start_time
+            user_content_replaced = user_content.replace(
+                '\n', '\\n'
+            ).replace('\r', '\\n')
+            response_content_replaced = response_content.replace(
+                '\n', '\\n'
+            ).replace('\r', '\\n')
 
-                try:
-                    prompt_tokens = response_json["usage"]["prompt_tokens"]
-                    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}, "
-                        f"完整内容: {response_json}"
-                    )
-                    prompt_tokens = 0
-                    completion_tokens = 0
-                    response_content = ""
-
-                user_content = ""
-                messages = data.get("messages", [])
-                for message in messages:
-                    if message["role"] == "user":
-                        if isinstance(message["content"], str):
-                            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", "") +
-                                        " "
-                                    )
-
-                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')
-
-                logging.info(
-                    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}"
-                )
-                with data_lock:
-                    request_timestamps.append(time.time())
-                    if "prompt_tokens" in response_json["usage"] and "completion_tokens" in response_json["usage"]:
-                        token_counts.append(response_json["usage"]["prompt_tokens"] + response_json["usage"]["completion_tokens"])
-                    else:
-                        token_counts.append(0)
-
-                return jsonify(response_json)
+            logging.info(
+                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)
 
         except requests.exceptions.RequestException as e:
             logging.error(f"请求转发异常: {e}")