Real-Time RAG: WebSocket Architectures for Instant Responses

TL;DR

Real-time RAG combines WebSocket for bidirectional communication, SSE for unidirectional streaming, and an event-driven pipeline for instant document updates. Result: responses in p50 < 800ms (TTFT) compared to 2-5s with a traditional architecture. This guide covers the 3 architectural patterns, a complete FastAPI + WebSocket example, and latency benchmarks by approach.

Why Traditional RAG Is Not Fast Enough

The Latency Problem in RAG

A traditional RAG pipeline (HTTP request-response) has inherent latency:

User -> [HTTP Request]
  -> Query embedding (50-100ms)
  -> Vector search (20-50ms)
  -> Reranking (100-200ms)
  -> LLM generation (2000-5000ms)
  -> [Complete HTTP Response]
User receives EVERYTHING at once after 2-5 seconds

The user waits for the entire response. With streaming, they see the first words in < 1 second.

Impact on User Experience

Perception is everything. Even if total time is similar, streaming radically changes the experience.

The 3 Architectural Patterns

Pattern 1: HTTP Polling (avoid)

DEVELOPERpython
# Anti-pattern: HTTP polling
# The client queries the server regularly

# Client (JavaScript)
"""
setInterval(async () => {
    const response = await fetch('/api/chat/status/' + taskId);
    if (response.data.status === 'complete') {
        displayAnswer(response.data.answer);
    }
}, 500); // Poll every 500ms
"""

# Problems:
# - Bandwidth waste (unnecessary requests)
# - Latency = polling interval (500ms minimum)
# - High server load
# - No streaming possible

Pattern 2: Server-Sent Events (SSE)

DEVELOPERpython
# Good for unidirectional streaming
# The server pushes events to the client

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import asyncio

app = FastAPI()

async def rag_stream(query: str):
    """RAG pipeline with SSE streaming"""

    # Phase 1: Retrieval (send progress signal)
    yield f"data: {json.dumps({'type': 'status', 'message': 'Searching...'})}\n\n"

    chunks = await vector_search(query, top_k=10)
    reranked = await rerank(query, chunks, top_n=5)

    yield f"data: {json.dumps({'type': 'status', 'message': 'Generating...'})}\n\n"

    # Phase 2: Generation with streaming
    context = build_context(reranked)
    async for token in llm_stream(query, context):
        yield f"data: {json.dumps({'type': 'token', 'content': token})}\n\n"

    # Phase 3: Sources
    sources = [{"title": c.title, "url": c.url} for c in reranked]
    yield f"data: {json.dumps({'type': 'sources', 'data': sources})}\n\n"

    yield f"data: {json.dumps({'type': 'done'})}\n\n"

@app.get("/api/chat/stream")
async def chat_stream(query: str):
    return StreamingResponse(
        rag_stream(query),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
        }
    )

SSE advantages: Simple, HTTP/2 compatible, automatic reconnection. SSE limitations: Unidirectional (server -> client only), not ideal for chat.

Pattern 3: WebSocket (recommended for chat)

DEVELOPERpython
# Ideal for bidirectional RAG chat
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import json
import asyncio

app = FastAPI()

class RAGWebSocketHandler:
    def __init__(self):
        self.retriever = VectorRetriever()
        self.reranker = CohereReranker()
        self.llm = StreamingLLM()

    async def handle_message(self, websocket: WebSocket, data: dict):
        query = data.get("message", "")
        conversation_id = data.get("conversation_id")

        # Signal: processing started
        await websocket.send_json({
            "type": "thinking",
            "message": "Analyzing your question..."
        })

        # Step 1: Parallel retrieval
        retrieval_start = time.time()
        chunks = await self.retriever.search(query, top_k=20)
        reranked = await self.reranker.rerank(query, chunks, top_n=5)

        await websocket.send_json({
            "type": "retrieval_done",
            "latency_ms": int((time.time() - retrieval_start) * 1000),
            "chunks_found": len(reranked)
        })

        # Step 2: Streaming generation
        context = self.build_context(reranked)
        full_response = ""

        async for token in self.llm.stream(query, context):
            full_response += token
            await websocket.send_json({
                "type": "token",
                "content": token
            })

        # Step 3: Send metadata
        await websocket.send_json({
            "type": "complete",
            "sources": [
                {"title": c.title, "score": c.score, "url": c.url}
                for c in reranked
            ],
            "usage": {
                "input_tokens": count_tokens(context),
                "output_tokens": count_tokens(full_response),
            }
        })

handler = RAGWebSocketHandler()

@app.websocket("/ws/chat")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()

    try:
        while True:
            data = await websocket.receive_json()

            if data.get("type") == "message":
                await handler.handle_message(websocket, data)
            elif data.get("type") == "ping":
                await websocket.send_json({"type": "pong"})

    except WebSocketDisconnect:
        print("Client disconnected")

Detailed Comparison of the 3 Approaches

Comparison Table

When to Use What?

DEVELOPERpython
# Decision framework
def choose_architecture(requirements):
    if requirements.get("bidirectional"):
        return "WebSocket"  # Chat, collaboration

    if requirements.get("streaming") and not requirements.get("bidirectional"):
        return "SSE"  # Notifications, dashboards

    if requirements.get("simple") and not requirements.get("real_time"):
        return "HTTP"  # Classic REST API

# Typical use cases
use_cases = {
    "WebSocket": [
        "Interactive RAG chatbot",
        "Team chat with AI",
        "Real-time collaboration",
        "Gaming / interactive applications",
    ],
    "SSE": [
        "LLM response streaming (like ChatGPT)",
        "Document update notifications",
        "Real-time dashboards",
        "Long task progress",
    ],
    "HTTP": [
        "Batch processing API",
        "Webhooks",
        "Third-party integrations",
        "One-off requests",
    ],
}

Event-Driven Architecture for Live Updates

The Document Update Problem

When a document is modified, how do you make it available instantly?

Document modified
  -> Webhook received (50ms)
  -> Re-chunking (200-500ms)
  -> Re-embedding (100-300ms per chunk)
  -> Upsert to Qdrant (50ms)
  -> Available for next queries
  Total: 500ms - 2s

Complete Event-Driven Architecture

DEVELOPERpython
# Event-driven update pipeline
import asyncio
from datetime import datetime

class EventDrivenRAGPipeline:
    def __init__(self):
        self.event_bus = AsyncEventBus()
        self.chunker = SmartChunker()
        self.embedder = BatchEmbedder()
        self.vector_db = QdrantClient()

    async def on_document_updated(self, event: DocumentEvent):
        """Webhook received: document modified"""
        doc = event.document

        # Step 1: Re-chunking
        chunks = await self.chunker.chunk(doc.content, doc.metadata)

        # Step 2: Delete old chunks
        await self.vector_db.delete(
            filter={"document_id": doc.id}
        )

        # Step 3: Batch embeddings
        embeddings = await self.embedder.embed_batch(
            [c.text for c in chunks]
        )

        # Step 4: Upsert to Qdrant
        points = [
            {
                "id": chunk.id,
                "vector": embedding,
                "payload": {
                    "text": chunk.text,
                    "document_id": doc.id,
                    "updated_at": datetime.utcnow().isoformat(),
                    **chunk.metadata,
                }
            }
            for chunk, embedding in zip(chunks, embeddings)
        ]
        await self.vector_db.upsert(points)

        # Step 5: Notify connected clients
        await self.event_bus.publish("document_updated", {
            "document_id": doc.id,
            "chunks_updated": len(chunks),
        })

    # Webhook endpoint
    @app.post("/webhooks/document")
    async def document_webhook(self, payload: dict):
        event = DocumentEvent.from_webhook(payload)
        # Async processing (does not block the webhook)
        asyncio.create_task(self.on_document_updated(event))
        return {"status": "accepted"}

Real-Time Client Notification

DEVELOPERpython
# Notify connected users that the knowledge base was updated
class ConnectionManager:
    def __init__(self):
        self.active_connections: dict[str, WebSocket] = {}

    async def broadcast_update(self, document_id: str):
        """Inform clients that new data is available"""
        message = {
            "type": "knowledge_updated",
            "document_id": document_id,
            "message": "The knowledge base has been updated.",
            "timestamp": datetime.utcnow().isoformat()
        }

        for ws in self.active_connections.values():
            try:
                await ws.send_json(message)
            except Exception:
                pass  # Disconnected client

Latency Optimization: Parallel Retrieval

Sequential vs Parallel Pipeline

DEVELOPERpython
# Sequential pipeline (slow)
async def sequential_rag(query: str):
    # Total: 50 + 50 + 150 + 100 = 350ms before generation
    embedding = await embed_query(query)       # 50ms
    chunks = await vector_search(embedding)     # 50ms
    reranked = await rerank(query, chunks)      # 150ms
    history = await get_conversation_history()   # 100ms
    return await generate(query, reranked, history)

# Parallel pipeline (fast)
async def parallel_rag(query: str):
    # Total: max(50+50+150, 100) = 250ms before generation
    # That is 100ms saved (29% faster)

    # Launch independent tasks in parallel
    embedding_task = asyncio.create_task(embed_query(query))
    history_task = asyncio.create_task(get_conversation_history())

    # Wait for embedding (needed for search)
    embedding = await embedding_task

    # Launch vector search
    chunks = await vector_search(embedding)

    # Reranking and history in parallel
    rerank_task = asyncio.create_task(rerank(query, chunks))
    history = await history_task  # Probably already finished

    reranked = await rerank_task

    return await generate(query, reranked, history)

Advanced Optimizations

DEVELOPERpython
# Speculative retrieval: start generation before reranking
async def speculative_rag(query: str):
    embedding = await embed_query(query)
    chunks = await vector_search(embedding, top_k=20)

    # Start generation with raw top-3 (without reranking)
    # while reranking runs
    quick_context = build_context(chunks[:3])
    gen_task = asyncio.create_task(
        generate_stream(query, quick_context)
    )

    # Reranking in parallel
    reranked = await rerank(query, chunks, top_n=5)

    # If reranking results differ significantly
    if reranked_differs_significantly(chunks[:3], reranked[:3]):
        gen_task.cancel()
        # Restart with correct context
        return generate_stream(query, build_context(reranked))

    # Otherwise, continue with current generation
    return gen_task

Complete Example: FastAPI + WebSocket + RAG Streaming

Complete Backend

DEVELOPERpython
# app/main.py
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
import json
import asyncio
import time

app = FastAPI(title="Real-time RAG API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

class RealtimeRAG:
    """Real-time RAG pipeline with WebSocket"""

    def __init__(self):
        self.embedder = OpenAIEmbedder()
        self.qdrant = QdrantClient(url="http://localhost:6333")
        self.reranker = CohereReranker()
        self.llm = AsyncOpenAI()

    async def process_query(self, ws: WebSocket, query: str, conv_id: str):
        start = time.perf_counter()
        metrics = {}

        try:
            # Phase 1: Retrieval
            await ws.send_json({"type": "phase", "phase": "retrieval"})

            t0 = time.perf_counter()
            embedding = await self.embedder.embed(query)
            results = await self.qdrant.search(
                collection_name="knowledge_base",
                query_vector=embedding,
                limit=20,
            )
            metrics["retrieval_ms"] = int((time.perf_counter() - t0) * 1000)

            # Phase 2: Reranking
            await ws.send_json({"type": "phase", "phase": "reranking"})

            t0 = time.perf_counter()
            texts = [r.payload["text"] for r in results]
            reranked = await self.reranker.rerank(query, texts, top_n=5)
            metrics["reranking_ms"] = int((time.perf_counter() - t0) * 1000)

            # Phase 3: Streaming generation
            await ws.send_json({"type": "phase", "phase": "generation"})

            context = "\n\n".join([r["text"] for r in reranked])
            prompt = f"Context:\n{context}\n\nQuestion: {query}\n\nAnswer:"

            full_response = ""
            t0 = time.perf_counter()
            first_token = True

            stream = await self.llm.chat.completions.create(
                model="gpt-4o",
                messages=[
                    {"role": "system", "content": "Answer based on the provided context."},
                    {"role": "user", "content": prompt}
                ],
                stream=True,
            )

            async for chunk in stream:
                if chunk.choices[0].delta.content:
                    token = chunk.choices[0].delta.content
                    full_response += token

                    if first_token:
                        metrics["ttft_ms"] = int((time.perf_counter() - t0) * 1000)
                        first_token = False

                    await ws.send_json({
                        "type": "token",
                        "content": token
                    })

            metrics["generation_ms"] = int((time.perf_counter() - t0) * 1000)
            metrics["total_ms"] = int((time.perf_counter() - start) * 1000)

            # Phase 4: Completion
            await ws.send_json({
                "type": "complete",
                "sources": [
                    {"text": r["text"][:100], "score": r["score"]}
                    for r in reranked
                ],
                "metrics": metrics,
            })

        except Exception as e:
            await ws.send_json({
                "type": "error",
                "message": str(e)
            })

rag = RealtimeRAG()

@app.websocket("/ws/chat/{conversation_id}")
async def websocket_chat(websocket: WebSocket, conversation_id: str):
    await websocket.accept()

    await websocket.send_json({
        "type": "connected",
        "conversation_id": conversation_id
    })

    try:
        while True:
            data = await websocket.receive_json()

            if data["type"] == "message":
                await rag.process_query(
                    websocket,
                    data["content"],
                    conversation_id
                )
            elif data["type"] == "ping":
                await websocket.send_json({"type": "pong"})

    except WebSocketDisconnect:
        pass

JavaScript Client

DEVELOPERjavascript
// client.js - WebSocket client for real-time RAG
class RAGWebSocketClient {
    constructor(conversationId) {
        this.ws = new WebSocket(`wss://api.example.com/ws/chat/${conversationId}`);
        this.responseContainer = document.getElementById('response');
        this.setupHandlers();
    }

    setupHandlers() {
        this.ws.onmessage = (event) => {
            const data = JSON.parse(event.data);

            switch (data.type) {
                case 'phase':
                    this.showPhase(data.phase);
                    break;
                case 'token':
                    this.appendToken(data.content);
                    break;
                case 'complete':
                    this.showSources(data.sources);
                    this.showMetrics(data.metrics);
                    break;
                case 'error':
                    this.showError(data.message);
                    break;
            }
        };
    }

    sendMessage(text) {
        this.responseContainer.innerHTML = '';
        this.ws.send(JSON.stringify({
            type: 'message',
            content: text
        }));
    }

    appendToken(token) {
        this.responseContainer.textContent += token;
    }
}

Latency Benchmarks by Approach

Test Protocol

DEVELOPERpython
# Benchmark on 1000 queries, 10K document base
benchmark = {
    "queries": 1000,
    "documents": 10_000,
    "vector_db": "Qdrant",
    "llm": "GPT-4o",
    "embedding": "text-embedding-3-small",
    "server": "4 vCPU, 8GB RAM",
}

p50/p95 Results

Latency Breakdown (WebSocket, p50)

Total: 620ms TTFT
  |- Query embedding    :  45ms (7%)
  |- Qdrant search      :  35ms (6%)
  |- Cohere reranking   : 140ms (23%)
  |- Network overhead   :  20ms (3%)
  +- LLM TTFT           : 380ms (61%)

The bottleneck is the LLM. That is why caching is so effective: it completely eliminates LLM time for repeated queries.

WebSocket Scalability Management

The Persistent Connection Challenge

DEVELOPERpython
# Scalable architecture with Redis pub/sub
# aioredis was merged into redis-py in 4.2: import it via redis.asyncio
from redis import asyncio as aioredis

class ScalableWebSocketManager:
    """Manages WebSockets across multiple server instances"""

    def __init__(self):
        self.redis = aioredis.from_url("redis://localhost:6379")
        self.local_connections: dict[str, WebSocket] = {}

    async def register(self, user_id: str, ws: WebSocket):
        self.local_connections[user_id] = ws
        # Subscribe to user's Redis channel
        pubsub = self.redis.pubsub()
        await pubsub.subscribe(f"user:{user_id}")

        asyncio.create_task(self._listen_redis(pubsub, ws))

    async def _listen_redis(self, pubsub, ws: WebSocket):
        """Relay Redis messages to WebSocket"""
        async for message in pubsub.listen():
            if message["type"] == "message":
                await ws.send_text(message["data"])

    async def broadcast_to_user(self, user_id: str, data: dict):
        """Send to a user (even on another server)"""
        await self.redis.publish(
            f"user:{user_id}",
            json.dumps(data)
        )

Load Balancing with Sticky Sessions

DEVELOPERnginx
# nginx.conf for WebSocket with sticky sessions
upstream rag_backend {
    ip_hash;  # IP-based sticky sessions
    server backend1:8000;
    server backend2:8000;
    server backend3:8000;
}

server {
    location /ws/ {
        proxy_pass http://rag_backend;
        proxy_http_version 1.1;
        proxy_set_header Upgrade $http_upgrade;
        proxy_set_header Connection "upgrade";
        proxy_set_header Host $host;
        proxy_read_timeout 86400;  # 24h for long connections
    }
}

FAQ

WebSocket or SSE for a RAG chatbot?

WebSocket is recommended for a chatbot because communication is bidirectional: the user sends messages, the server streams responses. SSE works if you only need server -> client streaming (e.g., notifications). In practice, most modern RAG chatbots use WebSocket.

How to handle WebSocket reconnection?

Implement an exponential backoff on the client side with conversation state persistence. On reconnection, send the last received message_id so the server can resume where it left off. This is exactly what the Ailog widget does natively.

What is the impact on server consumption?

Each WebSocket connection consumes ~50KB of RAM. A server with 8GB can handle ~100,000 idle connections. In practice, with RAG processing, expect ~5,000-10,000 simultaneous connections per server. Use Redis pub/sub for horizontal scaling.

Do CDNs support WebSocket?

Yes, most modern CDNs (Cloudflare, AWS CloudFront, Fastly) support WebSocket. However, some enterprise proxies may block them. Plan an automatic SSE fallback for those cases.

How to secure WebSocket connections?

Use WSS (WebSocket Secure) in production, with JWT token authentication in the initial handshake. Validate the token before accepting the connection. Implement per-user and per-IP rate limiting to prevent abuse.

Conclusion

Real-time RAG architecture with WebSocket delivers an unmatched user experience:

• TTFT < 800ms at p50 (vs 2s+ with classic HTTP)
• Native streaming of responses token by token
• Live updates of documents via event-driven pipeline
• Bidirectional communication for true interactive chat

The key is combining WebSocket for transport, parallelism for retrieval, and caching for frequent queries.

Ailog natively integrates WebSocket in its widget and API. Deploy a real-time RAG chatbot in 5 minutes at app.ailog.fr.

---

See also: Streaming RAG responses | RAG caching strategies | RAG agents orchestration