How to Build RAG Streaming With Inline Citations Over SSE

Why stream a RAG answer

Retrieval-Augmented Generation (RAG) has two latency sources: retrieval (fast) and LLM generation (slow - often 2-10 seconds for a full answer). Buffering both and returning one JSON blob means a blank screen the whole time. Streaming with Server-Sent Events (SSE) fixes this: tokens appear as the model generates them, and sources arrive as a final event once generation is done.

The event protocol looks like this:

data: {"type":"token","content":"Workspaces"}
data: {"type":"token","content":" are"}
data: {"type":"token","content":" isolated..."}
data: {"type":"citations","sources":["isolation","billing"]}
data: [DONE]

The endpoint

import json
import requests
from flask import Flask, request, Response, stream_with_context

app = Flask(__name__)
LLM_URL = "http://llm-service:8080/v1/chat/completions"

@app.post("/api/rag/stream")
def rag_stream():
    question = (request.get_json(silent=True) or {}).get("question", "")
    hits = retrieve(question, k=2)              # your retrieval function
    context = "\n".join(text for _, text in hits)
    sources = [doc_id for doc_id, _ in hits]

    @stream_with_context
    def gen():
        body = {
            "model": "gpt-4o-mini",
            "messages": [
                {"role": "system", "content": f"Answer using this context:\n{context}"},
                {"role": "user",   "content": question},
            ],
            "stream": True,
        }
        with requests.post(LLM_URL, json=body, stream=True, timeout=60) as r:
            for line in r.iter_lines():
                if not line:
                    continue
                decoded = line.decode("utf-8")
                if not decoded.startswith("data: "):
                    continue
                payload = decoded[6:]
                if payload.strip() == "[DONE]":
                    continue
                try:
                    delta = (
                        json.loads(payload)["choices"][0]["delta"]
                        .get("content", "")
                    )
                except Exception:
                    continue
                if delta:
                    yield f"data: {json.dumps({'type':'token','content':delta})}\n\n"
        yield f"data: {json.dumps({'type':'citations','sources':sources})}\n\n"
        yield "data: [DONE]\n\n"

    return Response(gen(), mimetype="text/event-stream")

Why citations come last

Retrieval happens before generation, so you already know which chunks were used. Emitting {"type":"citations"} after the final token is the simplest design: the UI renders the answer as it streams, then appends "[1] [2]" links when the citations event lands. Production systems often interleave citation markers inside the token stream (e.g. [isolation] after the relevant sentence), but tokens-then-sources is the right starting pattern.

Key implementation details

• stream=True on the requests call - without it, the requests library buffers the entire response before returning, defeating the whole point.
• stream_with_context in Flask - keeps the request context alive during generator execution so request.* stays accessible inside gen().
• mimetype="text/event-stream" - tells the client this is an SSE stream. Without this header, browsers and curl -N won't treat it as a stream.
• Consume lines, not chunks - iter_lines() reassembles partial TCP frames so you emit complete data: ... events, not split halves.

Consuming the stream on the client

const es = new EventSource("/api/rag/stream?q=" + encodeURIComponent(question));
// or use fetch + ReadableStream for POST
es.onmessage = (e) => {
  if (e.data === "[DONE]") { es.close(); return; }
  const msg = JSON.parse(e.data);
  if (msg.type === "token")     appendToken(msg.content);
  if (msg.type === "citations") renderSources(msg.sources);
};