Build a RAG Chatbot from Scratch — Part 3: Streaming Responses with SSE

# app/generation_streaming.py
from openai import OpenAI
from app.config import settings
from app.retrieval import retrieve
from app.prompts import build_rag_prompt

client = OpenAI(api_key=settings.openai_api_key)

def answer_question_stream(query: str):
    """RAG pipeline with streaming response — yields tokens as they arrive."""
    chunks = retrieve(query)

    if not chunks:
        yield "No relevant documents found."
        return

    context, system_prompt = build_rag_prompt(query, chunks)

    stream = client.chat.completions.create(
        model=settings.model,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {query}"}
        ],
        temperature=0.3,
        max_tokens=1000,
        stream=True
    )

    for chunk in stream:
        content = chunk.choices[0].delta.content
        if content:
            yield content

# app/api/chat_streaming.py
from fastapi import APIRouter
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from app.generation_streaming import answer_question_stream
from app.retrieval import retrieve
import json

router = APIRouter(prefix="/chat", tags=["chat"])

class Question(BaseModel):
    query: str

@router.post("/stream")
async def chat_stream(question: Question):
    async def event_stream():
        # First, send the sources
        chunks = retrieve(question.query)
        sources = [{"source": c["metadata"]["source"], "similarity": c["similarity"]} for c in chunks]
        yield f"event: sources\ndata: {json.dumps(sources)}\n\n"

        # Then stream the answer tokens
        for token in answer_question_stream(question.query):
            yield f"data: {json.dumps({'token': token})}\n\n"

        # Signal completion
        yield "data: [DONE]\n\n"

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no"  # Disable nginx buffering
        }
    )

// chat-client.js
async function streamChat(query) {
  const response = await fetch('/chat/stream', {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({ query })
  });

  const reader = response.body.getReader();
  const decoder = new TextDecoder();
  let answerEl = document.getElementById('answer');

  while (true) {
    const { done, value } = await reader.read();
    if (done) break;

    const text = decoder.decode(value);
    const lines = text.split('\n');

    for (const line of lines) {
      if (line.startsWith('data: ') && !line.includes('[DONE]')) {
        const data = JSON.parse(line.slice(6));
        if (data.token) {
          answerEl.textContent += data.token;
        }
      } else if (line.startsWith('event: sources')) {
        // Handle sources event
      }
    }
  }
}

SYSTEM_PROMPT_STREAMING = """Answer the question based ONLY on the provided context.
Cite sources inline using the format [1].

Example: "The vacation policy allows 20 days off [1]. Sick leave is 10 days [2]."

If the answer isn't in the context, say exactly: "I don't have enough information to answer this."
"""

from app.core.database import SessionLocal
from app.models import QueryLog
from datetime import datetime, timezone

def log_query(query: str, tokens_used: int, sources_count: int):
    db = SessionLocal()
    log = QueryLog(
        query=query,
        tokens_used=tokens_used,
        sources_count=sources_count,
        created_at=datetime.now(timezone.utc)
    )
    db.add(log)
    db.commit()
    db.close()

# Stream a response (curl can show SSE)
curl -N -X POST http://localhost:8000/chat/stream \
  -H "Content-Type: application/json" \
  -d '{"query": "What is RAG?"}'

# Output (tokens arrive progressively):
# event: sources
# data: [{"source":"rag_paper.pdf","similarity":0.95}]
#
# data: {"token":"RAG"}
# data: {"token":" stands"}
# data: {"token":" for"}
# data: {"token":" Retrieval-Augmented"}
# ...
# data: [DONE]