Update app.py

app.py

@@ -5,11 +5,8 @@ import requests
 import json
 import random
 import uuid
-import concurrent.fu
-import base64
-import io
+import concurrent.futures
 import threading
-from PIL import Imagetures
 from datetime import datetime, timedelta
 from apscheduler.schedulers.background import BackgroundScheduler
 from flask import Flask, request, jsonify, Response, stream_with_context
@@ -559,15 +556,11 @@ def handsome_chat_completions():
         return jsonify({"error": "Invalid request data"}), 400
 
     model_name = data['model']
-    
     request_type = determine_request_type(
         model_name,
         text_models,
-        free_text_models,
-        image_models,
-        free_image_models
+        free_text_models
     )
-    
     api_key = select_key(request_type, model_name)
 
     if not api_key:
@@ -580,248 +573,91 @@ def handsome_chat_completions():
                 )
             }
         ), 429
-    
+
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json"
     }
 
-    if model_name in image_models or model_name in free_image_models:
-        # Handle image model
-        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()
-        
-        siliconflow_data = {
-            "model": model_name,
-            "prompt": user_content,
-            "image_size": "1024x1024",
-            "batch_size": 1,
-            "num_inference_steps": 20,
-            "guidance_scale": 7.5,
-            "negative_prompt": None,
-            "seed": None,
-            "prompt_enhancement": False,
-        }
-        
-        try:
-            start_time = time.time()
-            response = requests.post(
-                "https://api.siliconflow.cn/v1/images/generations",
-                headers=headers,
-                json=siliconflow_data,
-                timeout=120
-            )
-
-            if response.status_code == 429:
-                return jsonify(response.json()), 429
-
-            response.raise_for_status()
-            end_time = time.time()
-            response_json = response.json()
-            total_time = end_time - start_time
-            
-            try:
-                images = response_json.get("images", [])
-                openai_images = []
-                for image_url in images:
-                    openai_images.append({"url": image_url})
-
-                response_data = {
-                    "created": int(time.time()),
-                    "data": openai_images
-                }
-
-            except (KeyError, ValueError, IndexError) as e:
-                logging.error(
-                    f"解析响应 JSON 失败: {e}, "
-                    f"完整内容: {response_json}"
-                )
-                response_data = {
-                    "created": int(time.time()),
-                    "data": []
-                }
-
-            logging.info(
-                f"使用的key: {api_key}, "
-                f"总共用时: {total_time:.4f}秒, "
-                f"使用的模型: {model_name}, "
-                f"用户的内容: {user_content}"
-            )
-
-            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}")
-            return jsonify({"error": str(e)}), 500
-
-    else:
-        # Handle text model
-        try:
-            start_time = time.time()
-            response = requests.post(
-                TEST_MODEL_ENDPOINT,
-                headers=headers,
-                json=data,
-                stream=data.get("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}"
-                                )
+    try:
+        start_time = time.time()
+        response = requests.post(
+            TEST_MODEL_ENDPOINT,
+            headers=headers,
+            json=data,
+            stream=data.get("stream", False),
+            timeout=60
+        )
 
-                    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}"
-                    )
+        if response.status_code == 429:
+            return jsonify(response.json()), 429
 
-                    with data_lock:
-                        request_timestamps.append(time.time())
-                        token_counts.append(prompt_tokens+completion_tokens)
+        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
 
-                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()
+                first_token_time = (
+                    first_chunk_time - start_time
+                    if first_chunk_time else 0
+                )
                 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 = ""
+                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}"
+                            )
 
                 user_content = ""
                 messages = data.get("messages", [])
@@ -836,7 +672,8 @@ def handsome_chat_completions():
                                     item.get("type") == "text"
                                 ):
                                     user_content += (
-                                        item.get("text", "") + " "
+                                        item.get("text", "") +
+                                        " "
                                     )
 
                 user_content = user_content.strip()
@@ -852,24 +689,91 @@ def handsome_chat_completions():
                     f"使用的key: {api_key}, "
                     f"提示token: {prompt_tokens}, "
                     f"输出token: {completion_tokens}, "
-                    f"首字用时: 0, "
+                    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())
-                    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)
+                    token_counts.append(prompt_tokens+completion_tokens)
 
-                return jsonify(response_json)
+            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
 
-        except requests.exceptions.RequestException as e:
-            logging.error(f"请求转发异常: {e}")
-            return jsonify({"error": str(e)}), 500
+            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)
+
+    except requests.exceptions.RequestException as e:
+        logging.error(f"请求转发异常: {e}")
+        return jsonify({"error": str(e)}), 500
 
 @app.route('/handsome/v1/models', methods=['GET'])
 def list_models():
@@ -878,7 +782,6 @@ def list_models():
 
     detailed_models = []
     
-    # 添加文本模型
     for model in text_models:
         detailed_models.append({
             "id": model,
@@ -905,7 +808,6 @@ def list_models():
             "parent": None
         })
 
-    # 添加 embedding 模型
     for model in embedding_models:
         detailed_models.append({
             "id": model,
@@ -932,59 +834,6 @@ def list_models():
             "parent": None
         })
 
-    # 添加图像模型
-    for model in image_models:
-         detailed_models.append({
-            "id": model,
-            "object": "model",
-            "created": 1678888888,
-            "owned_by": "openai",
-            "permission": [
-                {
-                    "id": f"modelperm-{uuid.uuid4().hex}",
-                    "object": "model_permission",
-                    "created": 1678888888,
-                    "allow_create_engine": False,
-                    "allow_sampling": True,
-                    "allow_logprobs": False,
-                    "allow_search_indices": False,
-                    "allow_view": True,
-                    "allow_fine_tuning": False,
-                    "organization": "*",
-                    "group": None,
-                    "is_blocking": False
-                }
-            ],
-            "root": model,
-            "parent": None
-        })
-    
-    for model in free_image_models:
-         detailed_models.append({
-            "id": model,
-            "object": "model",
-            "created": 1678888888,
-            "owned_by": "openai",
-            "permission": [
-                {
-                    "id": f"modelperm-{uuid.uuid4().hex}",
-                    "object": "model_permission",
-                    "created": 1678888888,
-                    "allow_create_engine": False,
-                    "allow_sampling": True,
-                    "allow_logprobs": False,
-                    "allow_search_indices": False,
-                    "allow_view": True,
-                    "allow_fine_tuning": False,
-                    "organization": "*",
-                    "group": None,
-                    "is_blocking": False
-                }
-            ],
-            "root": model,
-            "parent": None
-        })
-
     return jsonify({
         "success": True,
         "data": detailed_models
@@ -1154,6 +1003,10 @@ def handsome_embeddings():
     except requests.exceptions.RequestException as e:
         return jsonify({"error": str(e)}), 500
 
+import base64
+import io
+from PIL import Image
+
 @app.route('/handsome/v1/images/generations', methods=['POST'])
 def handsome_images_generations():
     if not check_authorization(request):