feat: streaming response via SSE with token-by-token display

api.py

@@ -1,9 +1,10 @@
+import json
 from contextlib import asynccontextmanager
 from fastapi import FastAPI, UploadFile, File, HTTPException
 from fastapi.staticfiles import StaticFiles
-from fastapi.responses import HTMLResponse
+from fastapi.responses import HTMLResponse, StreamingResponse
 from pydantic import BaseModel
-from src.rag import answer
+from src.rag import answer, retrieve
 from src.reranker import _get_model
 from src.uploader import ingest_pdf, get_uploaded_files
 
@@ -68,6 +69,30 @@ def documents():
     return get_uploaded_files()
 
 
+@app.post("/stream")
+async def stream(req: QueryRequest):
+    def generate():
+        contexts, sources = retrieve(req.question, top_k=req.top_k)
+        # Gửi sources trước
+        yield f"data: {json.dumps({'type': 'sources', 'sources': sources})}\n\n"
+        # Stream từng token
+        from src.rag import build_prompt, client, LLM_MODEL
+        prompt = build_prompt(req.question, contexts)
+        stream_resp = client.chat.completions.create(
+            model=LLM_MODEL,
+            messages=[{"role": "user", "content": prompt}],
+            stream=True,
+        )
+        for chunk in stream_resp:
+            token = chunk.choices[0].delta.content
+            if token:
+                yield f"data: {json.dumps({'type': 'token', 'text': token})}\n\n"
+        yield "data: {\"type\": \"done\"}\n\n"
+
+    return StreamingResponse(generate(), media_type="text/event-stream",
+                             headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})
+
+
 @app.get("/health")
 def health():
     return {"status": "ok"}

src/rag.py

@@ -31,8 +31,8 @@ CÂU HỎI: {question}
 TRẢ LỜI:"""
 
 
-def answer(question: str, top_k: int = 5) -> dict:
-    """RAG pipeline: rewrite -> embed -> hybrid retrieve -> rerank -> generate."""
+def retrieve(question: str, top_k: int = 5):
+    """Rewrite -> embed -> hybrid retrieve -> rerank. Returns (contexts, sources)."""
     rewritten = rewrite_query(question)
     if rewritten != question:
         print(f"Query rewritten: {rewritten}")
@@ -40,25 +40,30 @@ def answer(question: str, top_k: int = 5) -> dict:
     query_vec = embed_query(rewritten)
     candidates = vector_query(query_vec, query_text=rewritten, top_k=top_k * 3)
     uploaded = search_uploaded(query_vec, rewritten, top_k=top_k * 2)
-    all_candidates = candidates + uploaded
-    contexts = rerank(question, all_candidates, top_k=top_k)
+    contexts = rerank(question, candidates + uploaded, top_k=top_k)
+
+    sources = [
+        {
+            "source": c["metadata"]["source"],
+            "page": c["metadata"]["page"],
+            "rrf_score": round(c["score"], 3),
+            "rerank_score": round(c["rerank_score"], 3),
+        }
+        for c in contexts
+    ]
+    return contexts, sources
+
 
+def answer(question: str, top_k: int = 5) -> dict:
+    """RAG pipeline: rewrite -> embed -> hybrid retrieve -> rerank -> generate."""
+    contexts, sources = retrieve(question, top_k)
     prompt = build_prompt(question, contexts)
     response = client.chat.completions.create(
         model=LLM_MODEL,
         messages=[{"role": "user", "content": prompt}],
     )
-
     return {
         "answer": response.choices[0].message.content,
         "contexts": [c["text"] for c in contexts],
-        "sources": [
-            {
-                "source": c["metadata"]["source"],
-                "page": c["metadata"]["page"],
-                "rrf_score": round(c["score"], 3),
-                "rerank_score": round(c["rerank_score"], 3),
-            }
-            for c in contexts
-        ],
+        "sources": sources,
     }

static/index.html

@@ -281,14 +281,42 @@
       btn.disabled = true;
 
       try {
-        const res = await fetch('/query', {
+        const res = await fetch('/stream', {
           method: 'POST',
           headers: { 'Content-Type': 'application/json' },
           body: JSON.stringify({ question: q, top_k: 5 })
         });
-        const data = await res.json();
-        item.querySelector('.bubble-a').innerHTML =
-          `<div class="answer-text">${highlightCitations(data.answer)}</div>${renderSources(data.sources)}`;
+
+        const bubble = item.querySelector('.bubble-a');
+        bubble.innerHTML = '<div class="answer-text"></div>';
+        const answerEl = bubble.querySelector('.answer-text');
+        let fullText = '';
+        let sources = [];
+
+        const reader = res.body.getReader();
+        const decoder = new TextDecoder();
+        let buffer = '';
+
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          buffer += decoder.decode(value, { stream: true });
+          const lines = buffer.split('\n');
+          buffer = lines.pop();
+          for (const line of lines) {
+            if (!line.startsWith('data: ')) continue;
+            const msg = JSON.parse(line.slice(6));
+            if (msg.type === 'sources') {
+              sources = msg.sources;
+            } else if (msg.type === 'token') {
+              fullText += msg.text;
+              answerEl.innerHTML = highlightCitations(fullText);
+              item.scrollIntoView({ behavior: 'smooth', block: 'end' });
+            } else if (msg.type === 'done') {
+              bubble.innerHTML = `<div class="answer-text">${highlightCitations(fullText)}</div>${renderSources(sources)}`;
+            }
+          }
+        }
       } catch (e) {
         item.querySelector('.bubble-a').innerHTML = `<span class="error-text">Lỗi: ${e.message}</span>`;
       }