Update app.py

app.py

@@ -64,13 +64,54 @@ def chat():
     }
     data = request.get_json()
     stream_flag = data.get('stream', False)
+
+    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"}'
+
+
     try:
-        response = requests.post(API_BASE_URL+"/chat/completions",headers=headers,json=data,stream=stream_flag)
-        return Response(stream_with_context(response.iter_content(chunk_size=1024)),
-                        status=response.status_code,
-                        content_type=response.headers.get('content-type'))
+        # 如果不是流式请求，可以考虑不使用生成器，或者根据 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:
+            # 对于流式请求，使用生成器
+            # 注意：在生成器中处理异常时，如果已经开始发送数据，状态码和头部就不能更改了。
+            # 所以最好是在生成器开始之前捕获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("Initial upstream API request failed: %s", e)
+        return jsonify({"success": False, "message": f"Upstream API request failed: {e}"}), 500
     except Exception as e:
-        logging("get models error. %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__':