Update app.py

app.py

@@ -64,54 +64,13 @@ 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:
-        # 如果不是流式请求，可以考虑不使用生成器，或者根据 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
+        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'))
     except Exception as 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
+        logging("get models error. %s" ,e)
 
 
 if __name__ == '__main__':