Update app.py

app.py

@@ -1,7 +1,8 @@
 import os
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, Response, stream_with_context
 from functools import wraps
 import requests
+import json
 
 app = Flask(__name__)
 
@@ -9,7 +10,7 @@ AZURE_API_KEY = os.environ.get("AZURE_API_KEY")
 AZURE_API_BASE = "https://openai-skyline-jp.openai.azure.com"
 AZURE_API_VERSION = "2023-05-15"
 AZURE_DEPLOYMENT_NAME = "GPT-4"
-API_TOKEN = os.environ.get("API_TOKEN")  # 从环境变量获取 API Token
+API_TOKEN = os.environ.get("API_TOKEN")
 
 def token_required(f):
     @wraps(f)
@@ -28,10 +29,31 @@ def token_required(f):
         return f(*args, **kwargs)
     return decorated
 
+def convert_to_openai_format(azure_response):
+    openai_response = {
+        "id": azure_response.get("id", ""),
+        "object": "chat.completion",
+        "created": azure_response.get("created", 0),
+        "model": azure_response.get("model", "gpt-4"),
+        "choices": [],
+        "usage": azure_response.get("usage", {})
+    }
+    
+    for choice in azure_response.get("choices", []):
+        openai_choice = {
+            "index": choice.get("index", 0),
+            "message": choice.get("message", {}),
+            "finish_reason": choice.get("finish_reason")
+        }
+        openai_response["choices"].append(openai_choice)
+    
+    return {k: v for k, v in openai_response.items() if v is not None}
+
 @app.route('/hf/v1/chat/completions', methods=['POST'])
 @token_required
 def chat_completions():
     openai_data = request.json
+    stream = openai_data.get('stream', False)
 
     azure_url = f"{AZURE_API_BASE}/openai/deployments/{AZURE_DEPLOYMENT_NAME}/chat/completions?api-version={AZURE_API_VERSION}"
 
@@ -40,23 +62,47 @@ def chat_completions():
         "api-key": AZURE_API_KEY
     }
 
-    response = requests.post(azure_url, json=openai_data, headers=headers)
-
-    if response.status_code == 200:
-        azure_response = response.json()
-        
-        openai_response = {
-            "id": azure_response.get("id", ""),
-            "object": "chat.completion",
-            "created": azure_response.get("created", 0),
-            "model": openai_data.get("model", "gpt-4"),
-            "choices": azure_response.get("choices", []),
-            "usage": azure_response.get("usage", {})
-        }
-        
-        return jsonify(openai_response), 200
+    if not stream:
+        response = requests.post(azure_url, json=openai_data, headers=headers)
+
+        if response.status_code == 200:
+            azure_response = response.json()
+            openai_response = convert_to_openai_format(azure_response)
+            return jsonify(openai_response), 200
+        else:
+            return jsonify({"error": "Azure OpenAI API request failed", "details": response.text}), response.status_code
     else:
-        return jsonify({"error": "Azure OpenAI API request failed", "details": response.text}), response.status_code
+        def generate():
+            with requests.post(azure_url, json=openai_data, headers=headers, stream=True) as response:
+                if response.status_code != 200:
+                    yield f"data: {json.dumps({'error': 'Azure OpenAI API request failed', 'details': response.text})}\n\n"
+                    return
+
+                for line in response.iter_lines():
+                    if line:
+                        azure_data = json.loads(line.decode('utf-8').replace('data: ', ''))
+                        
+                        openai_data = {
+                            "id": azure_data.get("id", ""),
+                            "object": azure_data.get("object", ""),
+                            "created": azure_data.get("created", 0),
+                            "model": azure_data.get("model", "gpt-4"),
+                            "choices": [
+                                {
+                                    "index": choice.get("index", 0),
+                                    "delta": choice.get("delta", {}),
+                                    "finish_reason": choice.get("finish_reason")
+                                } for choice in azure_data.get("choices", [])
+                            ]
+                        }
+                        
+                        openai_data = {k: v for k, v in openai_data.items() if v is not None}
+                        
+                        yield f"data: {json.dumps(openai_data)}\n\n"
+
+                yield "data: [DONE]\n\n"
+
+        return Response(stream_with_context(generate()), content_type='text/event-stream')
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)