Update app.py

app.py

@@ -56,54 +56,63 @@ def getModels():
         logging.error("An unexpected error occurred in getModels: %s", e)
         return jsonify({"success": False, "message": str(e)}), 500
 
-@app.route("/v1/chat/completions", methods=['POST'])
+@app.route("/v1/chat/completions",methods=['POST'])
 def chat():
-    api_key = getAPI_KEY()
-    if not api_key:
-        return jsonify({"success": False, "message": "API Key not available"}), 500
-
     headers = {
-        "Authorization": "Bearer " + api_key,
-        "Content-Type": "application/json"
+        "Authorization":"Bearer "+getAPI_KEY(),
+        "Content-Type":"application/json"
     }
     data = request.get_json()
+    stream_flag = data.get('stream', True)
+
+    def generate():
+        try:
+            with requests.post(API_BASE_URL+"/chat/completions", headers=headers, json=data, stream=stream_flag) as response:
+                response.raise_for_status() # 检查上游请求是否成功
+                for chunk in response.iter_content(chunk_size=1024):
+                    yield chunk
+        except requests.exceptions.RequestException as e:
+            logging.error("Request to upstream API failed: %s", e)
+            # 在这里处理上游请求失败，例如可以 yield 一个错误消息或者抛出异常
+            # 但请注意，一旦开始 yield 数据，就不能改变 HTTP 状态码和头部了
+            yield b'{"error": "Upstream API request failed"}' # 作为 JSON 错误返回
+        except Exception as e:
+            logging.error("Unexpected error during streaming: %s", e)
+            yield b'{"error": "Internal server error during streaming"}'
 
-    # 重点在这里：根据客户端请求中的 'stream' 字段决定是否进行流式转发
-    # 如果客户端没有提供 'stream' 字段，我们假设它需要非流式响应（或者默认值取决于上游API的约定）
-    # 但为了明确支持非流式，这里我们设为 False
-    client_wants_stream = data.get('stream', False) # 客户端请求中 stream 的值
 
     try:
-        # 使用 requests.post 发送请求到上游 API
-        # 上游 API 的 'stream' 参数应该与客户端请求的 'stream' 字段保持一致
-        upstream_response = requests.post(
-            API_BASE_URL + "/chat/completions",
-            headers=headers,
-            json=data, # 客户端请求的 payload，包括 stream 字段
-            stream=client_wants_stream, # 控制 requests 是否以流式接收上游响应
-            timeout= (600 if client_wants_stream else 60) # 流式请求可以有更长的超时
-        )
-        upstream_response.raise_for_status() # 检查上游 API 响应的 HTTP 状态码
-
-        # 根据客户端是否想要流式响应来处理
-        if client_wants_stream:
-            # 流式响应：使用 stream_with_context 逐块发送
-            return Response(stream_with_context(upstream_response.iter_content(chunk_size=1024)),
-                            status=upstream_response.status_code,
-                            content_type=upstream_response.headers.get('content-type', 'application/json'))
+        # 如果不是流式请求，可以考虑不使用生成器，或者根据 stream_flag 来判断
+        if not stream_flag:
+            # 对于非流式请求，直接返回完整响应
+            response = requests.post(API_BASE_URL+"/chat/completions", headers=headers, json=data, stream=False)
+            response.raise_for_status()
+            return Response(response.content,
+                            status=response.status_code,
+                            content_type=response.headers.get('content-type'))
         else:
-            # 非流式响应：直接返回完整的响应内容
-            # 这里 upstream_response.content 会等待所有数据接收完毕
-            response_headers = {'Content-Type': upstream_response.headers.get('content-type', 'application/json')}
-            return (upstream_response.content, upstream_response.status_code, response_headers)
+            # 对于流式请求，使用生成器
+            # 注意：在生成器中处理异常时，如果已经开始发送数据，状态码和头部就不能更改了。
+            # 所以最好是在生成器开始之前捕获requests.post的异常。
+            initial_response = requests.post(API_BASE_URL+"/chat/completions", headers=headers, json=data, stream=True)
+            initial_response.raise_for_status() # 检查初始请求是否成功
+
+            return Response(generate_from_response(initial_response),
+                            status=initial_response.status_code,
+                            content_type=initial_response.headers.get('content-type'))
 
     except requests.exceptions.RequestException as e:
-        logging.error("Chat completion request error to upstream API: %s", e)
+        logging.error("Initial upstream API request failed: %s", e)
         return jsonify({"success": False, "message": f"Upstream API request failed: {e}"}), 500
     except Exception as e:
-        logging.error("An unexpected error occurred in chat completion: %s", e)
+        logging.error("Error setting up chat completion: %s", e)
         return jsonify({"success": False, "message": str(e)}), 500
 
+def generate_from_response(upstream_response):
+    # 这是一个辅助函数，用于将上游响应的迭代器包装成一个生成器
+    for chunk in upstream_response.iter_content(chunk_size=1024):
+        yield chunk
+
 
 if __name__ == '__main__':
     print("Starting Flask app...")