Update app.py

app.py

@@ -1,28 +1,24 @@
-from flask import Flask, request, jsonify, Response
-import requests
 import json
-import time
-
-app = Flask(__name__)
+from datetime import datetime
+import aiohttp
+from aiohttp import web
 
 @app.route('/')
 def index():
-    return "text-to-image with siliconflow", 200
-    
-@app.route('/ai/v1/chat/completions', methods=['POST'])
-def handle_request():
+    return "flux2api with siliconflow", 200
+
+async def handle_chat_completions(request):
     try:
-        body = request.json
+        body = await request.json()
         model = body.get('model')
         messages = body.get('messages')
         stream = body.get('stream', False)
 
         if not model or not messages or len(messages) == 0:
-            return jsonify({"error": "Bad Request: Missing required fields"}), 400
+            return web.json_response({'error': 'Bad Request: Missing required fields'}, status=400)
 
         prompt = messages[-1]['content']
-        new_url = f'https://api.siliconflow.cn/v1/{model}/text-to-image'
-
+        new_url = f"https://api.siliconflow.cn/v1/{model}/text-to-image"
         new_request_body = {
             "prompt": prompt,
             "image_size": "1024x1024",
@@ -31,20 +27,20 @@ def handle_request():
             "guidance_scale": 1
         }
 
-        headers = {
-            'accept': 'application/json',
-            'content-type': 'application/json',
-            'Authorization': request.headers.get('Authorization')
-        }
-
-        response = requests.post(new_url, headers=headers, json=new_request_body)
-        response_body = response.json()
+        async with aiohttp.ClientSession() as session:
+            async with session.post(new_url, json=new_request_body, headers={
+                'accept': 'application/json',
+                'content-type': 'application/json',
+                'Authorization': request.headers.get('Authorization')
+            }) as response:
+                response_body = await response.json()
 
         image_url = response_body['images'][0]['url']
-        unique_id = int(time.time() * 1000)
-        current_timestamp = unique_id // 1000
+        unique_id = int(datetime.now().timestamp() * 1000)
+        current_timestamp = int(unique_id / 1000)
 
         if stream:
+            # 构造流式响应
             response_payload = {
                 "id": unique_id,
                 "object": "chat.completion.chunk",
@@ -61,8 +57,11 @@ def handle_request():
                 ]
             }
             data_string = json.dumps(response_payload)
-            return Response(f"data: {data_string}\n\n", content_type='text/event-stream')
+            return web.Response(text=f"data: {data_string}\n\n", 
+                                status=200, 
+                                content_type='text/event-stream')
         else:
+            # 构造非流响应
             response_payload = {
                 "id": unique_id,
                 "object": "chat.completion",
@@ -85,11 +84,14 @@ def handle_request():
                     "total_tokens": len(prompt) + len(image_url)
                 }
             }
-            data_string = json.dumps(response_payload)
-            return Response(f"{data_string}\n\n", content_type='text/event-stream')
+            return web.json_response(response_payload)
+
+    except Exception as error:
+        return web.json_response({'error': f"Internal Server Error: {str(error)}"}, status=500)
+
+app = web.Application()
+app.router.add_post('/ai/v1/chat/completions', handle_chat_completions)
 
-    except Exception as e:
-        return jsonify({"error": f"Internal Server Error: {str(e)}"}), 500
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=8000, debug=True)
+    web.run_app(app, port=8000)