Update app.py

app.py

@@ -632,118 +632,47 @@ def handsome_chat_completions():
                 def generate():
                     first_chunk_time = None
                     full_response_content = ""
-                    try:
-                        response.raise_for_status()
-                        end_time = time.time()
-                        response_json = response.json()
-                        total_time = end_time - start_time
+                    for chunk in response.iter_content(chunk_size=1024):
+                        if chunk:
+                            if first_chunk_time is None:
+                                first_chunk_time = time.time()
+                            try:
+                                chunk_json = json.loads(chunk.decode("utf-8"))
+                                if "images" in chunk_json and chunk_json["images"]:
+                                    image_url = chunk_json["images"][0].get("url", "")
+                                    if image_url:
+                                        full_response_content += json.dumps({"url":image_url})
+                                        yield f"data: {json.dumps({'url': image_url})}\n\n".encode('utf-8')
+                                    else:
+                                        full_response_content += json.dumps({"content":"Failed to generate image"})
+                                        yield f"data: {json.dumps({'content': 'Failed to generate image'})}\n\n".encode('utf-8')
+                                else:
+                                    full_response_content += json.dumps({"content":"Failed to generate image"})
+                                    yield f"data: {json.dumps({'content': 'Failed to generate image'})}\n\n".encode('utf-8')
+                            except json.JSONDecodeError:
+                                logging.error(f"Failed to decode chunk JSON: {chunk.decode('utf-8')}")
+                                full_response_content += json.dumps({"content":"Failed to generate image"})
+                                yield f"data: {json.dumps({'content': 'Failed to generate image'})}\n\n".encode('utf-8')
+                                
+                    end_time = time.time()
+                    first_token_time = (
+                        first_chunk_time - start_time
+                        if first_chunk_time else 0
+                    )
+                    total_time = end_time - start_time
+                    
                     
-                        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]:
-                            image_url = images[0]["url"]
-                            logging.info(f"Extracted image URL: {image_url}")
-                        elif images and isinstance(images[0], str):
-                            image_url = images[0]
-                            logging.info(f"Extracted image URL: {image_url}")
-                        
-                        if image_url:
-                            chunk_data = {
-                                "id": f"chatcmpl-{uuid.uuid4()}",
-                                "object": "chat.completion.chunk",
-                                "created": int(time.time()),
-                                "model": model_name,
-                                "choices": [
-                                    {
-                                        "index": 0,
-                                        "delta": {
-                                            "role": "assistant",
-                                            "content": image_url
-                                        },
-                                        "finish_reason": None
-                                    }
-                                ]
-                            }
-                            yield f"data: {json.dumps(chunk_data)}\n\n".encode('utf-8')
-                            full_response_content = image_url
-                        else:
-                            chunk_data = {
-                                "id": f"chatcmpl-{uuid.uuid4()}",
-                                "object": "chat.completion.chunk",
-                                "created": int(time.time()),
-                                "model": model_name,
-                                "choices": [
-                                    {
-                                        "index": 0,
-                                        "delta": {
-                                            "role": "assistant",
-                                            "content": "Failed to generate image"
-                                        },
-                                        "finish_reason": None
-                                    }
-                                ]
-                            }
-                            yield f"data: {json.dumps(chunk_data)}\n\n".encode('utf-8')
-                            full_response_content = "Failed to generate image"
-                        
-                        end_chunk_data = {
-                            "id": f"chatcmpl-{uuid.uuid4()}",
-                            "object": "chat.completion.chunk",
-                            "created": int(time.time()),
-                            "model": model_name,
-                            "choices": [
-                                {
-                                    "index": 0,
-                                    "delta": {},
-                                    "finish_reason": "stop"
-                                }
-                            ]
-                        }
-                        yield f"data: {json.dumps(end_chunk_data)}\n\n".encode('utf-8')
-
-                        with data_lock:
-                            request_timestamps.append(time.time())
-                            token_counts.append(0)  # Image generation doesn't use tokens
-                    except requests.exceptions.RequestException as e:
-                        logging.error(f"请求转发异常: {e}")
-                        error_chunk_data = {
-                            "id": f"chatcmpl-{uuid.uuid4()}",
-                            "object": "chat.completion.chunk",
-                            "created": int(time.time()),
-                            "model": model_name,
-                            "choices": [
-                                {
-                                    "index": 0,
-                                    "delta": {
-                                        "role": "assistant",
-                                        "content": f"Error: {str(e)}"
-                                    },
-                                    "finish_reason": None
-                                }
-                            ]
-                        }
-                        yield f"data: {json.dumps(error_chunk_data)}\n\n".encode('utf-8')
-                        end_chunk_data = {
-                            "id": f"chatcmpl-{uuid.uuid4()}",
-                            "object": "chat.completion.chunk",
-                            "created": int(time.time()),
-                            "model": model_name,
-                            "choices": [
-                                {
-                                    "index": 0,
-                                    "delta": {},
-                                    "finish_reason": "stop"
-                                }
-                            ]
-                        }
-                        yield f"data: {json.dumps(end_chunk_data)}\n\n".encode('utf-8')
-                        
                     logging.info(
                         f"使用的key: {api_key}, "
+                        f"首字用时: {first_token_time:.4f}秒, "
+                        f"总共用时: {total_time:.4f}秒, "
                         f"使用的模型: {model_name}"
                     )
+
+                    with data_lock:
+                        request_timestamps.append(time.time())
+                        token_counts.append(0)  # Image generation doesn't use tokens
+                    
                     yield "data: [DONE]\n\n".encode('utf-8')
                 return Response(stream_with_context(generate()), content_type='text/event-stream')
             else:
@@ -1015,7 +944,7 @@ def handsome_chat_completions():
         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():
     if not check_authorization(request):