Update app.py

app.py

@@ -547,8 +547,8 @@ def check_tokens():
     return jsonify(results)
 
 @app.route('/handsome/v1/chat/completions', methods=['POST'])
-def handsome_chat_completions(request, authorization_key):
-    if not check_authorization(request, authorization_key):
+def handsome_chat_completions():
+    if not check_authorization(request):
         return jsonify({"error": "Unauthorized"}), 401
 
     data = request.get_json()
@@ -556,13 +556,13 @@ def handsome_chat_completions(request, authorization_key):
         return jsonify({"error": "Invalid request data"}), 400
 
     model_name = data['model']
-    
+  
     request_type = determine_request_type(
         model_name,
         text_models + image_models,
         free_text_models + free_image_models
     )
-    
+  
     api_key = select_key(request_type, model_name)
 
     if not api_key:
@@ -580,7 +580,7 @@ def handsome_chat_completions(request, authorization_key):
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json"
     }
-    
+  
     if model_name in image_models:
         # Handle image generation
         # Map OpenAI-style parameters to SiliconFlow's parameters
@@ -595,7 +595,7 @@ def handsome_chat_completions(request, authorization_key):
             "seed": data.get("seed"),
             "prompt_enhancement": False,
         }
-        
+      
         # Parameter validation and adjustments
         if siliconflow_data["batch_size"] < 1:
             siliconflow_data["batch_size"] = 1
@@ -606,12 +606,12 @@ def handsome_chat_completions(request, authorization_key):
             siliconflow_data["num_inference_steps"] = 1
         if siliconflow_data["num_inference_steps"] > 50:
             siliconflow_data["num_inference_steps"] = 50
-            
+          
         if siliconflow_data["guidance_scale"] < 0:
             siliconflow_data["guidance_scale"] = 0
         if siliconflow_data["guidance_scale"] > 100:
             siliconflow_data["guidance_scale"] = 100
-        
+      
         if siliconflow_data["image_size"] not in ["1024x1024", "512x1024", "768x512", "768x1024", "1024x576", "576x1024"]:
             siliconflow_data["image_size"] = "1024x1024"
 
@@ -623,7 +623,7 @@ def handsome_chat_completions(request, authorization_key):
                 json=siliconflow_data,
                 timeout=120
             )
-           
+         
             if response.status_code == 429:
                 return jsonify(response.json()), 429
 
@@ -631,10 +631,10 @@ def handsome_chat_completions(request, authorization_key):
             end_time = time.time()
             response_json = response.json()
             total_time = end_time - start_time
-            
+          
             try:
                 images = response_json.get("images", [])
-                
+              
                 # Extract the first URL if available
                 image_url = ""
                 if images and isinstance(images[0], dict) and "url" in images[0]:
@@ -643,7 +643,7 @@ def handsome_chat_completions(request, authorization_key):
                 elif images and isinstance(images[0], str):
                   image_url = images[0]
                   logging.info(f"Extracted image URL: {image_url}")
-                
+              
                 # Construct the expected JSON output - Mimicking OpenAI
                 response_data = {
                    "id": f"chatcmpl-{uuid.uuid4()}",
@@ -656,10 +656,21 @@ def handsome_chat_completions(request, authorization_key):
                         "index": 0,
                         "message": {
                             "role": "assistant",
-                            "content": None,
-                            "image_url": image_url if image_url else None, # Return URL in image_url
+                            "content": None, # set to None as image is in tool_calls
+                            "tool_calls": [
+                                {
+                                    "id": f"call_{uuid.uuid4()}",
+                                    "type": "function",
+                                    "function": {
+                                        "name": "image_generation",
+                                        "arguments": json.dumps({
+                                            "image_url": image_url
+                                        })
+                                    }
+                                }
+                            ]
                         },
-                        "finish_reason": "length",
+                        "finish_reason": "tool_calls",
                      }
                    ],
                    "usage": { # Added usage
@@ -668,9 +679,7 @@ def handsome_chat_completions(request, authorization_key):
                         "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}, "
@@ -688,7 +697,6 @@ def handsome_chat_completions(request, authorization_key):
                         "message": {
                             "role": "assistant",
                             "content": "Failed to process image data",
-                            "image_url": None,
                         },
                         "finish_reason": "stop",
                       }
@@ -706,6 +714,10 @@ def handsome_chat_completions(request, authorization_key):
                 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,69 +730,191 @@ def handsome_chat_completions(request, authorization_key):
                 TEST_MODEL_ENDPOINT,
                 headers=headers,
                 json=data,
-                stream=False,
+                stream=data.get("stream", False),
                 timeout=60
             )
             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:
-                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", "") +
-                                    " "
+
+            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}"
                                 )
 
-            user_content = user_content.strip()
+                    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}"
+                    )
 
-            user_content_replaced = user_content.replace(
-                '\n', '\\n'
-            ).replace('\r', '\\n')
-            response_content_replaced = response_content.replace(
-                '\n', '\\n'
-            ).replace('\r', '\\n')
+                    with data_lock:
+                        request_timestamps.append(time.time())
+                        token_counts.append(prompt_tokens+completion_tokens)
 
-            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)
+                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
+
+                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}")