Update app.py

app.py

@@ -94,49 +94,34 @@ async def chat_endpoint(request: ChatRequest):
     }
 
     async def event_generator():
-        try:
-            # Send initial connection message
-            yield ": connected\n\n"
-            
-            last_yield_time = asyncio.get_event_loop().time()
-            
-            async for event in rag_app.astream_events(inputs, config=config, version="v1"):
-                # Send heartbeat every 15 seconds to keep connection alive
-                current_time = asyncio.get_event_loop().time()
-                if current_time - last_yield_time > 15:
-                    yield ": heartbeat\n\n"
-                    last_yield_time = current_time
-                kind = event["event"]
-                if kind == "on_chat_model_stream":
-                    chunk = event["data"]["chunk"]
-                    content = chunk.content if hasattr(chunk, "content") else str(chunk)
+        # [TRICK 1] Send a "padding" comment to force Nginx to flush the buffer immediately
+        # 2KB of whitespace usually does the trick
+        padding = " " * 2048
+        yield f": {padding}\n\n"
+
+        async for event in rag_app.astream_events(inputs, config=config, version="v1"):
+            kind = event["event"]
+            if kind == "on_chat_model_stream":
+                content = event["data"]["chunk"].content
+                
+                if content:
+                    # We use JSON serialization to be safe with newlines
+                    # But if you prefer your manual replace, that's fine too:
+                    # data = content.replace("\n", "\\n")
                     
-                    if content:
-                        # Escape newlines and special chars for JSON
-                        escaped_content = json.dumps(content)
-                        yield f"data: {escaped_content}\n\n"
-                        
-                        # Update last yield time
-                        last_yield_time = asyncio.get_event_loop().time()
-                        
-                        # CRITICAL: Force immediate flush
-                        await asyncio.sleep(0.01)
-            
-            yield "data: [DONE]\n\n"
-            
-        except Exception as e:
-            error_msg = json.dumps({"error": str(e)})
-            yield f"data: {error_msg}\n\n"
+                    # Ensure we aren't sending empty chunks
+                    if content.strip() != "":
+                        data = content.replace("\n", "\\n")
+                        yield f"data: {data}\n\n"
 
     return StreamingResponse(
         event_generator(), 
         media_type="text/event-stream",
         headers={
-            "Cache-Control": "no-cache, no-transform",
+            "Cache-Control": "no-cache",
             "Connection": "keep-alive",
-            "Content-Type": "text/event-stream",
+            "Content-Encoding": "none",  # [TRICK 2] Disable compression
             "X-Accel-Buffering": "no",
-            "Transfer-Encoding": "chunked",
         },
     )