Better streaming

app.py

@@ -220,6 +220,14 @@ async def stream_response(query, history, game_context=None, mode: Optional[str]
         async for chunk in stream_response_local(query, history, game_context, mode):
             yield chunk
 
+def _text_event(text: str) -> str:
+    """Encode model text as one SSE event without exposing framing characters."""
+    payload = json.dumps({"text": text}, ensure_ascii=False, separators=(",", ":"))
+    return f"data: {payload}\n\n"
+
+def _done_event() -> str:
+    return 'data: {"done":true}\n\n'
+
 def _format_game_context_for_prompt(game_context: Optional[Union[str, Dict[str, Any]]]) -> str:
     """Return a concise text snippet to inject into the system prompt from game context.
 
@@ -388,13 +396,13 @@ async def stream_response_gemini(query, history, game_context=None, mode: Option
                 if not first_token_logged:
                     logger.info(f"First token time (Gemini): {time.time() - start_time:.2f}s")
                     first_token_logged = True
-                yield f"data: {chunk.text}\n\n"
-        yield "data: [DONE]\n\n"
+                yield _text_event(chunk.text)
+        yield _done_event()
 
     except Exception as e:
         logger.error(f"Gemini API error: {str(e)}")
-        yield f"data: Sorry, I encountered an error with Gemini API: {str(e)}\n\n"
-        yield "data: [DONE]\n\n"
+        yield _text_event(f"Sorry, I encountered an error with Gemini API: {str(e)}")
+        yield _done_event()
 
 
 async def stream_response_local(query, history, game_context=None, mode: Optional[str] = None):
@@ -441,8 +449,8 @@ async def stream_response_local(query, history, game_context=None, mode: Optiona
         history_tokens = [len(generator.tokenize(msg["content"].encode('utf-8'), add_bos=False, special=True)) for msg in history]
     except Exception as e:
         logger.error(f"Tokenization error: {str(e)}")
-        yield f"data: Sorry, I encountered a tokenization error: {str(e)}\n\n"
-        yield "data: [DONE]\n\n"
+        yield _text_event(f"Sorry, I encountered a tokenization error: {str(e)}")
+        yield _done_event()
         return
 
     total_tokens = system_tokens + query_tokens + sum(history_tokens) + len(history) * 10 + 10
@@ -470,12 +478,12 @@ async def stream_response_local(query, history, game_context=None, mode: Optiona
                     if not first_token_logged:
                         logger.info(f"First token time (local): {time.time() - start_time:.2f}s")
                         first_token_logged = True
-                    yield f"data: {token}\n\n"
-            yield "data: [DONE]\n\n"
+                    yield _text_event(token)
+            yield _done_event()
         except Exception as e:
             logger.error(f"Generation error: {str(e)}")
-            yield f"data: Sorry, I encountered an error during generation: {str(e)}\n\n"
-            yield "data: [DONE]\n\n"
+            yield _text_event(f"Sorry, I encountered an error during generation: {str(e)}")
+            yield _done_event()
 
 class QueryRequest(BaseModel):
     query: str