feat: SSE streaming pipeline, progress bar, PDF export

app.py

@@ -1,7 +1,7 @@
-import sys, os
+import sys, os, json
 sys.path.append(os.path.dirname(os.path.abspath(__file__)))
-from flask import Flask, render_template, request, jsonify, send_file
-from agent.agent import run_pipeline
+from flask import Flask, render_template, request, jsonify, send_file, Response, stream_with_context
+from agent.agent import run_pipeline, run_query_architect, run_literature_scout, run_evidence_synthesiser, run_citation_builder
 from reportlab.lib.pagesizes import A4
 from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
 from reportlab.lib.units import mm
@@ -12,10 +12,12 @@ import io
 
 app = Flask(__name__)
 
+
 @app.route("/")
 def index():
     return render_template("index.html")
 
+
 @app.route("/query", methods=["POST"])
 def query():
     data = request.get_json()
@@ -33,6 +35,57 @@ def query():
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
+
+@app.route("/stream", methods=["GET"])
+def stream():
+    user_query = request.args.get("query", "").strip()
+    if not user_query:
+        return jsonify({"error": "Empty query"}), 400
+
+    def generate():
+        def emit(event, data):
+            return "event: " + event + "\ndata: " + json.dumps(data) + "\n\n"
+
+        try:
+            # Stage 1
+            yield emit("stage", {"stage": 1, "pct": 10})
+            queries = run_query_architect(user_query)
+            yield emit("queries", {"queries": queries, "pct": 25})
+
+            # Stage 2
+            yield emit("stage", {"stage": 2, "pct": 35})
+            papers = run_literature_scout(queries)
+            yield emit("papers", {"paper_count": len(papers), "pct": 55})
+
+            # Stage 3
+            yield emit("stage", {"stage": 3, "pct": 70})
+            synthesis = run_evidence_synthesiser(user_query, papers)
+            yield emit("synthesis", {"synthesis": synthesis, "pct": 88})
+
+            # Stage 4
+            yield emit("stage", {"stage": 4, "pct": 90})
+            citations = run_citation_builder(papers)
+            yield emit("done", {
+                "synthesis": synthesis,
+                "citations": citations,
+                "paper_count": len(papers),
+                "queries": queries,
+                "pct": 100
+            })
+
+        except Exception as e:
+            yield emit("error", {"message": str(e)})
+
+    return Response(
+        stream_with_context(generate()),
+        mimetype="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "X-Accel-Buffering": "no"
+        }
+    )
+
+
 @app.route("/export-pdf", methods=["POST"])
 def export_pdf():
     data = request.get_json()
@@ -69,7 +122,7 @@ def export_pdf():
     story = []
     story.append(Paragraph("ARIA — Autonomous Research Intelligence Agent", title_style))
     story.append(Paragraph(
-        f"Query: {query}  |  {paper_count} papers retrieved  |  Groq LLaMA-3.1",
+        "Query: " + query + "  |  " + str(paper_count) + " papers retrieved  |  Groq LLaMA-3.1",
         meta_style))
     story.append(HRFlowable(width="100%", thickness=1,
         color=colors.HexColor("#1e2936"), spaceAfter=16))
@@ -116,9 +169,10 @@ def export_pdf():
     doc.build(story)
     buf.seek(0)
     safe_query = "".join(c for c in query[:40] if c.isalnum() or c in " -_").strip()
-    filename = f"ARIA_{safe_query}.pdf".replace(" ", "_")
+    filename = "ARIA_" + safe_query.replace(" ", "_") + ".pdf"
     return send_file(buf, mimetype="application/pdf",
         as_attachment=True, download_name=filename)
 
+
 if __name__ == "__main__":
-    app.run(debug=True, port=5000)
+    app.run(debug=True, port=5000, threaded=True)

templates/index.html

@@ -502,92 +502,90 @@ function toggleCitations() {
         document.getElementById("citationsList").classList.contains("open") ? "▲" : "▼";
 }
 
+function renderResults(data) {
+    const meta = document.getElementById("metaBar");
+    meta.innerHTML = `<div class="badge green">${data.paper_count} papers retrieved</div><div class="badge">${data.queries ? data.queries.length : 0} PubMed queries</div><div class="badge">LLaMA-3.1-8B via Groq</div>`;
+    const sections = parseSynthesis(data.synthesis);
+    const synthEl = document.getElementById("synthesis");
+    synthEl.innerHTML = "";
+    SECTIONS.forEach(s => {
+        if (!sections[s.label]) return;
+        const div = document.createElement("div");
+        div.className = "section open";
+        div.innerHTML = `<div class="section-header" onclick="toggleSection(this.parentElement)"><span class="section-tag">${s.label}</span><span class="section-chevron">▲</span></div><div class="section-body">${sections[s.label].replace(/\n/g, "<br>")}</div>`;
+        synthEl.appendChild(div);
+    });
+    const citList = document.getElementById("citationsList");
+    citList.innerHTML = "";
+    citList.classList.add("open");
+    data.citations.split("\n").forEach(line => {
+        if (!line.trim()) return;
+        const pmidMatch = line.match(/PMID:\s*(\d+)/);
+        const pmid = pmidMatch ? pmidMatch[1] : null;
+        const div = document.createElement("div");
+        div.className = "citation-item";
+        div.innerHTML = pmid ? line.replace(/PMID:\s*\d+/, `<span class="citation-pmid">PMID: ${pmid}</span>`) : line;
+        citList.appendChild(div);
+    });
+    document.getElementById("results").classList.add("visible");
+}
+
+function updateProgress(pct) {
+    document.getElementById("progressFill").style.width = pct + "%";
+    document.getElementById("progressPct").textContent = pct + "%";
+}
+
 async function submitQuery() {
     const q = document.getElementById("queryInput").value.trim();
     if (!q) return;
-
     const btn = document.getElementById("submitBtn");
     btn.disabled = true;
-
     document.getElementById("results").classList.remove("visible");
     document.getElementById("errorBox").classList.remove("visible");
     document.getElementById("pipeline").classList.add("visible");
     setStage(1);
-
-    // Simulate stage progression while waiting
-    const delays = [0, 3000, 8000, 14000];
-    delays.forEach((d, i) => setTimeout(() => setStage(i + 1), d));
-
-    try {
-        const res = await fetch("/query", {
-            method: "POST",
-            headers: { "Content-Type": "application/json" },
-            body: JSON.stringify({ query: q })
-        });
-        const data = await res.json();
+    updateProgress(5);
+    lastQuery = q;
+
+    const es = new EventSource("/stream?query=" + encodeURIComponent(q));
+
+    es.addEventListener("stage", e => {
+        const d = JSON.parse(e.data);
+        setStage(d.stage);
+        updateProgress(d.pct);
+    });
+    es.addEventListener("queries", e => {
+        const d = JSON.parse(e.data);
+        updateProgress(d.pct);
+    });
+    es.addEventListener("papers", e => {
+        const d = JSON.parse(e.data);
+        updateProgress(d.pct);
+    });
+    es.addEventListener("synthesis", e => {
+        const d = JSON.parse(e.data);
+        updateProgress(d.pct);
+    });
+    es.addEventListener("done", e => {
+        const data = JSON.parse(e.data);
         lastResult = data;
-        lastQuery = q;
-
-        setStage(5); // all done
-
-        if (data.error) {
-            document.getElementById("errorBox").textContent = "Error: " + data.error;
-            document.getElementById("errorBox").classList.add("visible");
-            return;
-        }
-
-        // Meta bar
-        const meta = document.getElementById("metaBar");
-        meta.innerHTML = `
-            <div class="badge green">${data.paper_count} papers retrieved</div>
-            <div class="badge">${data.queries ? data.queries.length : 0} PubMed queries</div>
-            <div class="badge">LLaMA-3.1-8B via Groq</div>
-        `;
-
-        // Synthesis sections
-        const sections = parseSynthesis(data.synthesis);
-        const synthEl = document.getElementById("synthesis");
-        synthEl.innerHTML = "";
-        SECTIONS.forEach(s => {
-            if (!sections[s.label]) return;
-            const div = document.createElement("div");
-            div.className = "section open";
-            div.innerHTML = `
-                <div class="section-header" onclick="toggleSection(this.parentElement)">
-                    <span class="section-tag">${s.label}</span>
-                    <span class="section-chevron">▲</span>
-                </div>
-                <div class="section-body">${sections[s.label].replace(/\n/g, '<br>')}</div>
-            `;
-            synthEl.appendChild(div);
-        });
-
-        // Citations
-        const citList = document.getElementById("citationsList");
-        citList.innerHTML = "";
-        citList.classList.add("open");
-        data.citations.split("\n").forEach(line => {
-            if (!line.trim()) return;
-            const pmidMatch = line.match(/PMID:\s*(\d+)/);
-            const pmid = pmidMatch ? pmidMatch[1] : null;
-            const div = document.createElement("div");
-            div.className = "citation-item";
-            div.innerHTML = pmid
-                ? line.replace(/PMID:\s*\d+/, `<span class="citation-pmid">PMID: ${pmid}</span>`)
-                : line;
-            citList.appendChild(div);
-        });
-
-        document.getElementById("results").classList.add("visible");
-
-    } catch (err) {
-        document.getElementById("errorBox").textContent = "Network error: " + err.message;
+        updateProgress(100);
+        setStage(5);
+        es.close();
+        renderResults(data);
+        btn.disabled = false;
+    });
+    es.addEventListener("error", e => {
+        es.close();
+        let msg = "Pipeline error";
+        try { msg = JSON.parse(e.data).message; } catch(err) {}
+        document.getElementById("errorBox").textContent = "Error: " + msg;
         document.getElementById("errorBox").classList.add("visible");
-    } finally {
         btn.disabled = false;
-    }
+    });
 }
 
+
 document.getElementById("queryInput").addEventListener("keydown", e => {
     if (e.key === "Enter" && !e.shiftKey) { e.preventDefault(); submitQuery(); }
 });