Update app.py

app.py

@@ -94,22 +94,49 @@ async def chat_endpoint(request: ChatRequest):
     }
 
     async def event_generator():
-        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:
-                    data = content.replace("\n", "\\n")
-                    yield f"data: {data}\n\n"
+        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)
+                    
+                    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"
 
     return StreamingResponse(
         event_generator(), 
         media_type="text/event-stream",
         headers={
-            "Cache-Control": "no-cache",
+            "Cache-Control": "no-cache, no-transform",
             "Connection": "keep-alive",
+            "Content-Type": "text/event-stream",
             "X-Accel-Buffering": "no",
+            "Transfer-Encoding": "chunked",
         },
     )