Streaming RAG Pipelines: Overlap Retrieval With First Token Generation

Streaming RAG Pipelines: Overlap Retrieval With First Token Generation

Waiting for full retrieval → full prompt pack → first token makes chat feel broken even when total latency is acceptable. Stream thoughtfully: overlap speculative generation on early chunks with late retrieval, but never emit uncited claims before citation validation catches up. Perceived latency and grounding can coexist if you stage the stream.

⚡ TL;DR: Start generation on a fast primary retrieve; opportunistically merge late reranked chunks into a second phase; stream tokens with a preamble that can be revised; buffer claim JSON until citations validate; use Lambda response streaming / Bedrock ConverseStream. See Lambda + Bedrock streaming and Citation-Required RAG.

Two-phase retrieve + generate

async function* streamingRag(question: string, filters: Filters) {
  // Phase A: cheap/fast retrieve
  const early = await retrieve(question, { ...filters, k: 4, mode: "fast" });

  // Kick late rerank in parallel
  const lateP = retrieve(question, { ...filters, k: 12, mode: "rerank" });

  // Stream a grounded preamble from early context only
  yield { type: "status", message: "Searching code…" };
  const draft = bedrockConverseStream({
    system: SYSTEM_CITATIONS,
    chunks: early,
    question,
    instruction: "Stream a 2-sentence summary. Do not list claims yet.",
  });

  for await (const tok of draft) yield { type: "token", tok };

  const late = await lateP;
  const merged = dedupeMerge(early, late);

  // Phase B: structured claims with full evidence (buffered)
  const claims = await bedrockConverseJson({
    chunks: merged,
    question,
    schema: ANSWER_JSON_SCHEMA,
  });
  const check = validateCitations(claims, merged);
  if (!check.ok) {
    yield { type: "refuse", violations: check.violations };
    return;
  }
  yield { type: "claims", claims };
  yield { type: "sources", chunks: merged };
}

Do not stream free-form answers that invent APIs before sources arrive — users screenshot the wrong draft.

Overlap without lying

Stage User sees Rule
0–300 ms “Searching…” No model claims
Early tokens High-level summary Only facts in early chunks; hedge
Claims event Bullet claims + links Citation-validated only
Sources panel Files/lines Always

If early and late evidence disagree, prefer revising the summary with a follow-up event over silently contradicting mid-stream.

yield { type: "revise_summary", text: reconciledSummary(earlySummary, claims) };

AWS wiring

  • Bedrock: ConverseStream for phase A tokens; non-stream JSON for phase B claims (easier to validate).
  • Lambda Function URLs: stream SSE/NDJSON events; apply backpressure if clients pause.
  • Timeouts: budget early retrieve 150–250 ms; late rerank can finish after tokens started.
# Illustrative budgets
BUDGETS_MS = {
  "retrieve_fast": 200,
  "retrieve_rerank": 800,
  "first_token": 600,
  "claims_json": 2500,
}

Measuring perceived vs true latency

Track time_to_status, time_to_first_token, time_to_valid_claims, and citation_precision. Optimizing TTFT while tanking citation precision is a regression—gate launches on both (eval harness).

Closing checklist

Dos
– Overlap fast retrieve with late rerank
– Stream status + hedged summary early; buffer claims until validated
– Emit explicit revise / refuse events
– Budget each stage; alarm on TTFT and citation precision
– Use NDJSON/SSE event types the UI understands

Donts
– Do not stream uncited API names in phase A
– Do not block first pixel on full rerank
– Do not drop late chunks that contradict the draft without a revise
– Do not use API Gateway for long token streams (prefer Function URLs)
– Do not skip validation because “streaming is hard”

Related reading

Last updated on September 11, 2026


Discover more from CheatCoders

Subscribe to get the latest posts sent to your email.

Comments

No comments yet. Why don’t you start the discussion?

Leave a Reply