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,6 +18,7 @@ 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"])
@@ -101,34 +102,78 @@ def chat_completions():
             "session_hash": session_hash,
         }
 
-    def generate():
+    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:
         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)
+        data_stream = 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": "*",
-        },
-    )
+        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)
 
 
 def gen_res_data(data, time_now=0, start=False):
@@ -152,5 +197,4 @@ 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()