Update main.py

main.py

@@ -1,5 +1,5 @@
 import gevent.pywsgi
-from gevent import monkey; monkey.patch_all()
+from gevent import monkey;monkey.patch_all()
 from flask import Flask, request, Response, jsonify
 import argparse
 import requests
@@ -18,7 +18,6 @@ parser.add_argument("--port", type=int, help="Set the port.(default: 7860)", def
 args = parser.parse_args()
 
 base_url = os.getenv('MODEL_BASE_URL')
-streaming_enabled = os.getenv('STREAMING_ENABLED', 'false').lower() == 'true'
 
 @app.route('/api/v1/models', methods=["GET", "POST"])
 @app.route('/v1/models', methods=["GET", "POST"])
@@ -102,78 +101,34 @@ def chat_completions():
             "session_hash": session_hash,
         }
 
-    if streaming_enabled:
-        def generate():
-            response = requests.post(f"{base_url}/queue/join", json=json_prompt)
-            url = f"{base_url}/queue/data?session_hash={session_hash}"
-            data = requests.get(url, stream=True)
-
-            time_now = int(time.time())
-
-            for line in data.iter_lines():
-                if line:
-                    decoded_line = line.decode("utf-8")
-                    json_line = json.loads(decoded_line[6:])
-                    if json_line["msg"] == "process_starts":
-                        res_data = gen_res_data({}, time_now=time_now, start=True)
-                        yield f"data: {json.dumps(res_data)}\n\n"
-                    elif json_line["msg"] == "process_generating":
-                        res_data = gen_res_data(json_line, time_now=time_now)
-                        yield f"data: {json.dumps(res_data)}\n\n"
-                    elif json_line["msg"] == "process_completed":
-                        yield "data: [DONE]"
-
-        return Response(
-            generate(),
-            mimetype="text/event-stream",
-            headers={
-                "Access-Control-Allow-Origin": "*",
-                "Access-Control-Allow-Headers": "*",
-            },
-        )
-    else:
+    def generate():
         response = requests.post(f"{base_url}/queue/join", json=json_prompt)
         url = f"{base_url}/queue/data?session_hash={session_hash}"
-        data_stream = requests.get(url, stream=True)
+        data = requests.get(url, stream=True)
 
         time_now = int(time.time())
-        collected_data = []
-
-        try:
-            for line in data_stream.iter_lines():
-                if line:
-                    try:
-                        decoded_line = line.decode("utf-8")
-                        json_line = json.loads(decoded_line[6:])  # 确保这里的切片索引是正确的
-                        collected_data.append(json_line)
-                        if json_line.get("msg") == "process_completed":
-                            break
-                    except json.JSONDecodeError:
-                        print("Error decoding JSON from response")
-                        continue  # 跳过不能解码的行
-
-            if not collected_data:
-                return jsonify({"error": "No data received from the server", "word_count": 0})
-
-            # 检查数据是否有效并尝试生成最终响应
-            try:
-                res_data = gen_res_data(collected_data[-1], time_now=time_now) if collected_data else {"error": "No valid data processed", "word_count": 0}
-                print("Response data:", res_data)  # 打印输出res_data
-            except KeyError as e:
-                res_data = {"error": f"Missing key in data: {str(e)}", "word_count": 0}
-                print("KeyError:", res_data)  # 错误情况下也打印res_data
-            except IndexError as e:
-                res_data = {"error": f"Index error in data: {str(e)}", "word_count": 0}
-                print("IndexError:", res_data)
-            except Exception as e:
-                res_data = {"error": f"An unexpected error occurred: {str(e)}", "word_count": 0}
-                print("Exception:", res_data)
-
-        except Exception as e:
-            res_data = {"error": f"An error occurred while processing stream: {str(e)}", "word_count": 0}
-            print("Stream processing exception:", res_data)
-
-        return jsonify(res_data)
+
+        for line in data.iter_lines():
+            if line:
+                decoded_line = line.decode("utf-8")
+                json_line = json.loads(decoded_line[6:])
+                if json_line["msg"] == "process_starts":
+                    res_data = gen_res_data({}, time_now=time_now, start=True)
+                    yield f"data: {json.dumps(res_data)}\n\n"
+                elif json_line["msg"] == "process_generating":
+                    res_data = gen_res_data(json_line, time_now=time_now)
+                    yield f"data: {json.dumps(res_data)}\n\n"
+                elif json_line["msg"] == "process_completed":
+                    yield "data: [DONE]"
+
+    return Response(
+        generate(),
+        mimetype="text/event-stream",
+        headers={
+            "Access-Control-Allow-Origin": "*",
+            "Access-Control-Allow-Headers": "*",
+        },
+    )
 
 
 def gen_res_data(data, time_now=0, start=False):
@@ -197,4 +152,5 @@ def gen_res_data(data, time_now=0, start=False):
 
 
 if __name__ == "__main__":
+    # app.run(host=args.host, port=args.port, debug=True)
     gevent.pywsgi.WSGIServer((args.host, args.port), app).serve_forever()