feat: PRISMA filter, follow-up questions, rate limit retry, staggered post-pipeline calls

agent/agent.py

@@ -9,6 +9,7 @@ from retrieval.pubmed import fetch_pubmed
 
 
 def get_llm():
+    from langchain_groq import ChatGroq
     return ChatGroq(
         model="llama-3.1-8b-instant",
         temperature=0,
@@ -16,6 +17,21 @@ def get_llm():
     )
 
 
+def llm_invoke_with_retry(llm, prompt, max_retries=5):
+    import time
+    for attempt in range(max_retries):
+        try:
+            return llm.invoke(prompt)
+        except Exception as e:
+            if "429" in str(e) or "rate_limit" in str(e).lower():
+                wait = 10 * (attempt + 1)
+                print(f"[ARIA] Rate limit hit, waiting {wait}s (attempt {attempt+1}/{max_retries})")
+                time.sleep(wait)
+            else:
+                raise
+    raise RuntimeError("Max retries exceeded on rate limit")
+
+
 @tool
 def PubMedSearch(query: str) -> str:
     """Searches PubMed for biomedical literature abstracts."""
@@ -37,7 +53,7 @@ def run_query_architect(user_question):
         "Return ONLY a numbered list 1-5, one query per line, no explanations.\n\n"
         "Question: " + user_question
     )
-    response = llm.invoke(prompt)
+    response = llm_invoke_with_retry(llm, prompt)
     raw_lines = response.content.strip().split("\n")
     queries = []
     for line in raw_lines:
@@ -87,7 +103,7 @@ def run_evidence_synthesiser(user_question, papers):
         "Retrieved Literature:\n" + corpus + "\n\n"
         "Be precise and cite PMIDs throughout."
     )
-    response = llm.invoke(prompt)
+    response = llm_invoke_with_retry(llm, prompt)
     return response.content
 
 
@@ -104,7 +120,6 @@ def run_citation_builder(papers):
     return "\n".join(result_lines)
 
 
-
 def run_confidence_scorer(synthesis):
     llm = get_llm()
     prompt = (
@@ -122,7 +137,7 @@ def run_confidence_scorer(synthesis):
         "Scores: 8-10 = strong evidence, 5-7 = moderate, 1-4 = weak/preliminary.\n\n"
         "Synthesis:\n" + synthesis
     )
-    response = llm.invoke(prompt)
+    response = llm_invoke_with_retry(llm, prompt)
     import json
     text = response.content.strip()
     text = text.replace("```json", "").replace("```", "").strip()
@@ -148,7 +163,7 @@ def run_selective_review(user_question, selected_papers):
         "Question: " + user_question + "\n\n"
         "Selected Papers:\n" + corpus
     )
-    response = llm.invoke(prompt)
+    response = llm_invoke_with_retry(llm, prompt)
     return response.content
 
 
@@ -163,17 +178,17 @@ def run_predictive_model(user_question, synthesis):
         "## Destructive Forecast\n"
         "2-3 sentences: Which current assumptions, treatments, or paradigms does the evidence suggest "
         "may be challenged, overturned, or significantly revised in coming years?\n\n"
-        "IMPORTANT: Always produce both sections even if evidence is limited. Never ask for more input.\n""Be specific and grounded in the evidence. No speculation beyond what the data implies.\n\n"
+        "IMPORTANT: Always produce both sections even if evidence is limited. Never ask for more input.\n"
+        "Be specific and grounded in the evidence. No speculation beyond what the data implies.\n\n"
         "Clinical Question: " + user_question + "\n\n"
         "Synthesis:\n" + synthesis
     )
-    response = llm.invoke(prompt)
+    response = llm_invoke_with_retry(llm, prompt)
     return response.content
 
 
 def run_table_extractor(user_question, synthesis, papers):
     llm = get_llm()
-    # Build a brief paper list for context
     paper_list = []
     for pmid, p in list(papers.items())[:10]:
         paper_list.append("PMID " + pmid + ": " + p.get("title", "N/A") + " (" + p.get("year", "") + ")")
@@ -200,11 +215,56 @@ def run_table_extractor(user_question, synthesis, papers):
         "Papers:\n" + papers_str + "\n\n"
         "Synthesis:\n" + synthesis[:1500]
     )
-    response = llm.invoke(prompt)
+    response = llm_invoke_with_retry(llm, prompt)
     import json
     text = response.content.strip().replace("```json", "").replace("```", "").strip()
     return json.loads(text)
 
+
+def run_prisma_filter(user_question, papers):
+    llm = get_llm()
+    import json
+    paper_list = []
+    for pmid, p in papers.items():
+        paper_list.append(
+            "PMID " + pmid + ": " + p.get("title", "N/A") + "\n" +
+            p.get("abstract", "")[:200]
+        )
+    corpus = "\n\n".join(paper_list)
+    prompt = (
+        "You are a systematic review methodologist applying PRISMA screening criteria.\n"
+        "For each paper, decide if it should be INCLUDED or EXCLUDED for answering this clinical question.\n"
+        "Return ONLY valid JSON, no markdown, no explanation.\n"
+        "Format:\n"
+        "{\n"
+        '  "decisions": [\n'
+        '    {"pmid": "12345678", "decision": "included", "reason": "one sentence"},\n'
+        '    {"pmid": "87654321", "decision": "excluded", "reason": "one sentence"}\n'
+        '  ]\n'
+        "}\n\n"
+        "Inclusion criteria: directly relevant to the clinical question, has empirical data or clinical findings.\n"
+        "Exclusion criteria: off-topic, editorial, commentary without data, animal studies if human data exists.\n\n"
+        "Clinical Question: " + user_question + "\n\n"
+        "Papers:\n" + corpus
+    )
+    response = llm_invoke_with_retry(llm, prompt)
+    text = response.content.strip().replace("```json", "").replace("```", "").strip()
+    data = json.loads(text)
+    result = {}
+    for d in data["decisions"]:
+        pmid = d["pmid"]
+        if pmid in papers:
+            result[pmid] = {
+                **papers[pmid],
+                "included": d["decision"] == "included",
+                "reason": d["reason"]
+            }
+    for pmid in papers:
+        if pmid not in result:
+            result[pmid] = {**papers[pmid], "included": True, "reason": "Not reviewed"}
+    return result
+
+
 def run_pipeline(user_question):
     print("[1/4] Query Architect: generating search queries...")
     queries = run_query_architect(user_question)
@@ -237,4 +297,4 @@ if __name__ == "__main__":
     print("\n=== SYNTHESIS ===")
     print(result["synthesis"])
     print("\n=== REFERENCES ===")
-    print(result["citations"])
+    print(result["citations"])

app.py

@@ -1,7 +1,8 @@
-import sys, os, json
+import sys, os, json, time
 sys.path.append(os.path.dirname(os.path.abspath(__file__)))
 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 agent.agent import (run_pipeline, run_query_architect, run_literature_scout,
+                         run_evidence_synthesiser, run_citation_builder, llm_invoke_with_retry)
 from reportlab.lib.pagesizes import A4
 from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
 from reportlab.lib.units import mm
@@ -55,22 +56,46 @@ def stream():
             # Stage 2
             yield emit("stage", {"stage": 2, "pct": 35})
             papers = run_literature_scout(queries)
-            yield emit("papers", {"paper_count": len(papers), "pct": 55})
+            yield emit("papers", {"paper_count": len(papers), "pct": 50})
+
+            # PRISMA filter
+            yield emit("stage", {"stage": 3, "pct": 55})
+            from agent.agent import run_prisma_filter
+            filtered = run_prisma_filter(user_query, papers)
+            included = {pmid: p for pmid, p in filtered.items() if p["included"]}
+            yield emit("prisma", {
+                "filtered": {
+                    pmid: {"title": p.get("title", ""), "included": p["included"], "reason": p["reason"]}
+                    for pmid, p in filtered.items()
+                },
+                "included_count": len(included),
+                "excluded_count": len(filtered) - len(included),
+                "pct": 65
+            })
+            time.sleep(12)
 
-            # Stage 3
-            yield emit("stage", {"stage": 3, "pct": 70})
-            synthesis = run_evidence_synthesiser(user_query, papers)
+            # Stage 4 - synthesise on included papers only
+            yield emit("stage", {"stage": 4, "pct": 70})
+            synthesis = run_evidence_synthesiser(user_query, included)
             yield emit("synthesis", {"synthesis": synthesis, "pct": 88})
 
-            # Stage 4
-            yield emit("stage", {"stage": 4, "pct": 90})
-            citations = run_citation_builder(papers)
+            # Stage 5
+            yield emit("stage", {"stage": 5, "pct": 90})
+            citations = run_citation_builder(included)
             yield emit("done", {
                 "synthesis": synthesis,
                 "citations": citations,
-                "paper_count": len(papers),
+                "paper_count": len(included),
                 "queries": queries,
-                "papers": {pmid: {"title": p.get("title",""), "abstract": p.get("abstract",""), "authors": p.get("authors",""), "journal": p.get("journal",""), "year": p.get("year","")} for pmid, p in papers.items()},
+                "papers": {
+                    pmid: {
+                        "title": p.get("title", ""),
+                        "abstract": p.get("abstract", ""),
+                        "authors": p.get("authors", ""),
+                        "journal": p.get("journal", ""),
+                        "year": p.get("year", "")
+                    } for pmid, p in included.items()
+                },
                 "pct": 100
             })
 
@@ -175,7 +200,6 @@ def export_pdf():
         as_attachment=True, download_name=filename)
 
 
-
 @app.route("/score", methods=["POST"])
 def score():
     data = request.get_json()
@@ -224,22 +248,26 @@ import json as _json
 from datetime import datetime
 SESSIONS_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "sessions.json")
 
+
 def load_sessions():
     try:
         return _json.load(open(SESSIONS_FILE))
     except:
         return []
 
+
 def save_session(entry):
     sessions = load_sessions()
     sessions.insert(0, entry)
     sessions = sessions[:20]
     _json.dump(sessions, open(SESSIONS_FILE, "w"), indent=2)
 
+
 @app.route("/sessions", methods=["GET"])
 def get_sessions():
     return jsonify({"sessions": load_sessions()})
 
+
 @app.route("/sessions/save", methods=["POST"])
 def save_session_route():
     data = request.get_json()
@@ -273,6 +301,7 @@ def extract_table():
         traceback.print_exc()
         return jsonify({"error": str(e)}), 500
 
+
 @app.route("/followup", methods=["POST"])
 def followup():
     data = request.get_json()
@@ -298,10 +327,11 @@ def followup():
             f"Papers:\n{corpus}\n\n"
             f"Follow-up Question: {question}"
         )
-        response = llm.invoke(prompt)
+        response = llm_invoke_with_retry(llm, prompt)
         return jsonify({"answer": response.content})
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
+
 if __name__ == "__main__":
-    app.run(debug=True, port=5000, threaded=True)
+    app.run(debug=True, port=5000, threaded=True)

templates/index.html

@@ -1058,9 +1058,9 @@ async function submitQuery() {
         setStage(5);
         es.close();
         renderResults(data);
-        scoreResults(data.synthesis);
-        setTimeout(() => runPredictiveModel(lastQuery, data.synthesis), 4000);
-        setTimeout(() => buildTable(lastQuery, data.synthesis, data.papers || {}), 6000);
+	setTimeout(() => scoreResults(data.synthesis), 2000);
+	setTimeout(() => runPredictiveModel(lastQuery, data.synthesis), 20000);
+	setTimeout(() => buildTable(lastQuery, data.synthesis, data.papers || {}), 38000);
         saveSession(data, q);
         btn.disabled = false;
     });