Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import random
 import requests
 from huggingface_hub import InferenceClient
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, Response, stream_with_context
 from flask_cors import CORS
 import json
 
@@ -16,7 +16,7 @@ def get_random_api_key():
     else:
         raise ValueError("API keys not found. Please set the KEYS environment variable.")
 
-def generate_story(prompt, style, stream=False):
+def generate_story(prompt, style):
     try:
         client = InferenceClient(api_key=get_random_api_key())
         
@@ -25,20 +25,14 @@ def generate_story(prompt, style, stream=False):
             {"role": "user", "content": prompt}
         ]
         
-        completion = client.chat.completions.create(model="Qwen/Qwen2.5-Coder-32B-Instruct", messages=messages, temperature=0.7, max_tokens=1200, stream=stream)
-        if stream:
-            story = ""
-            for chunk in completion:
-                if chunk.choices:
-                    story += chunk.choices[0].delta.content or ""
-                    yield story
-        else:
-            story = completion.choices[0].message.content
-            yield story
+        completion = client.chat.completions.create(model="Qwen/Qwen2.5-Coder-32B-Instruct", messages=messages, temperature=0.7, max_tokens=1200, stream=True)
+        for chunk in completion:
+            if chunk.choices:
+                yield chunk.choices[0].delta.content or ""
     except Exception as e:
         yield f"Ошибка генерации: {e}"
 
-def edit_story(original_story, edited_prompt, stream=False):
+def edit_story(original_story, edited_prompt):
     if not original_story:
         yield f"Сначала сгенерируйте историю!"
         return
@@ -53,20 +47,14 @@ def edit_story(original_story, edited_prompt, stream=False):
             {"role": "user", "content": edited_prompt},
             {"role": "assistant", "content": original_story}
         ]
-        completion = client.chat.completions.create(model="Qwen/Qwen2.5-Coder-32B-Instruct", messages=messages, temperature=0.7, max_tokens=155000, stream=stream)
-        if stream:
-            edited_story = ""
-            for chunk in completion:
-                if chunk.choices:
-                    edited_story += chunk.choices[0].delta.content or ""
-                    yield edited_story
-        else:
-            edited_story = completion.choices[0].message.content
-            yield edited_story
+        completion = client.chat.completions.create(model="Qwen/Qwen2.5-Coder-32B-Instruct", messages=messages, temperature=0.7, max_tokens=155000, stream=True)
+        for chunk in completion:
+            if chunk.choices:
+                yield chunk.choices[0].delta.content or ""
     except Exception as e:
         yield f"Ошибка редактирования: {e}"
 
-def next_story_func(original_story, next_prompt, continuation_type="Продолжение", stream=False):
+def next_story_func(original_story, next_prompt, continuation_type="Продолжение"):
     if not original_story:
         yield f"Сначала сгенерируйте историю!"
         return
@@ -85,19 +73,22 @@ def next_story_func(original_story, next_prompt, continuation_type="Продол
             {"role": "user", "content": continuation_prompt},
             {"role": "assistant", "content": original_story}
         ]
-        completion = client.chat.completions.create(model="Qwen/Qwen2.5-Coder-32B-Instruct", messages=messages, temperature=0.7, max_tokens=1200, stream=stream)
-        if stream:
-            next_story = ""
-            for chunk in completion:
-                if chunk.choices:
-                    next_story += chunk.choices[0].delta.content or ""
-                    yield next_story
-        else:
-            next_story = completion.choices[0].message.content
-            yield next_story
+        completion = client.chat.completions.create(model="Qwen/Qwen2.5-Coder-32B-Instruct", messages=messages, temperature=0.7, max_tokens=1200, stream=True)
+        for chunk in completion:
+            if chunk.choices:
+                yield chunk.choices[0].delta.content or ""
     except Exception as e:
         yield f"Ошибка продления: {e}"
 
+def stream_sse(generator):
+    def event_stream():
+       try:
+           for chunk in generator:
+              yield f"data: {json.dumps({'message': chunk})}\n\n"
+       except Exception as e:
+           yield f"data: {json.dumps({'error': True, 'message': str(e)})}\n\n"
+    return Response(stream_with_context(event_stream()), content_type='text/event-stream')
+
 @app.route('/generate', methods=['POST'])
 def api_generate_story():
     data = request.get_json()
@@ -105,17 +96,10 @@ def api_generate_story():
         return jsonify({"error": True, "message": "Missing 'input' in request body"}), 400
     
     prompt = data['input']
-    style = data.get('style', 'Приключенческая')  # Default style
-    stream = data.get('stream', False)
+    style = data.get('style', 'Приключенческая')
     
-    try:
-        story_generator = generate_story(prompt, style, stream)
-        if stream:
-            return jsonify({"error": False, "message": "".join(story_generator)})
-        else:
-            return jsonify({"error": False, "message": next(story_generator)})
-    except Exception as e:
-        return jsonify({"error": True, "message": str(e)}), 500
+    story_generator = generate_story(prompt, style)
+    return stream_sse(story_generator)
 
 @app.route('/edit', methods=['POST'])
 def api_edit_story():
@@ -125,16 +109,9 @@ def api_edit_story():
     
     original_story = data['original']
     edited_prompt = data['input']
-    stream = data.get('stream', False)
     
-    try:
-        edited_story_generator = edit_story(original_story, edited_prompt, stream)
-        if stream:
-           return jsonify({"error": False, "message": "".join(edited_story_generator)})
-        else:
-            return jsonify({"error": False, "message": next(edited_story_generator)})
-    except Exception as e:
-        return jsonify({"error": True, "message": str(e)}), 500
+    edited_story_generator = edit_story(original_story, edited_prompt)
+    return stream_sse(edited_story_generator)
 
 @app.route('/continue', methods=['POST'])
 def api_continue_story():
@@ -144,18 +121,11 @@ def api_continue_story():
     
     original_story = data['original']
     next_prompt = data['input']
-    continuation_type = data.get('type', 'Продолжение') # Default continuation type
-    stream = data.get('stream', False)
+    continuation_type = data.get('type', 'Продолжение')
     
-    try:
-        next_story_generator = next_story_func(original_story, next_prompt, continuation_type, stream)
-        if stream:
-           return jsonify({"error": False, "message": "".join(next_story_generator)})
-        else:
-           return jsonify({"error": False, "message": next(next_story_generator)})
-    except Exception as e:
-        return jsonify({"error": True, "message": str(e)}), 500
+    next_story_generator = next_story_func(original_story, next_prompt, continuation_type)
+    return stream_sse(next_story_generator)
+
 
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=7860, debug=True)
-