Fix streaming

app.py

@@ -1,4 +1,4 @@
-from flask import Flask, request, Response, stream_with_context
+from flask import Flask, request, Response
 import requests
 import os
 
@@ -6,16 +6,24 @@ app = Flask(__name__)
 
 OLLAMA_API_URL = "http://localhost:11434"
 
-@stream_with_context
 def generate(data):
-    with requests.post(f"{OLLAMA_API_URL}/api/chat", json=data, stream=True) as r:
-        for line in r.iter_lines(decode_unicode=True):
-            yield line
+    try:
+        with requests.post(f"{OLLAMA_API_URL}/api/chat", json=data, stream=True) as r:
+            r.raise_for_status()  # Raise an exception for bad status codes
+            for line in r.iter_lines():
+                if line:
+                    yield line.decode('utf-8') + '\n'
+    except Exception as e:
+        # Yield an error message if something goes wrong
+        yield f'{{"error": "Stream error: {str(e)}"}}\n'
 
 @app.route('/chat', methods=['POST'])
 def chat():
     data = request.get_json(silent=True)
-    return Response(generate(data), mimetype='text/plain')
+    response = Response(generate(data), mimetype='text/plain')
+    response.headers['Cache-Control'] = 'no-cache'
+    response.headers['X-Accel-Buffering'] = 'no'  # Disable buffering in nginx
+    return response
 
 @app.route('/embed', methods=['POST'])
 def embed():
@@ -25,4 +33,4 @@ def embed():
 
 if __name__ == '__main__':
     port = int(os.environ.get('PORT', 7860))
-    app.run(host='0.0.0.0', port=port, debug=False)
+    app.run(host='0.0.0.0', port=port, debug=False, threaded=True)