File size: 12,414 Bytes
bcceea4
 
1ccdc0b
 
bcceea4
 
d1f5505
1ccdc0b
 
 
bcceea4
7dd2c52
 
 
bcceea4
7dd2c52
 
bcceea4
 
d62c791
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bcceea4
 
 
1ccdc0b
 
 
bcceea4
 
 
 
 
 
 
 
 
 
 
 
 
 
1ccdc0b
d62c791
bcceea4
 
 
 
 
 
 
 
 
 
904dfee
bcceea4
58d2397
 
 
904dfee
 
 
 
 
58d2397
 
bcceea4
 
 
 
 
 
58d2397
bcceea4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d62c791
bcceea4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
59cb995
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d62c791
59cb995
 
 
 
 
9083f85
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d62c791
9083f85
 
61b12b4
 
 
 
 
 
 
 
 
 
 
 
d62c791
 
61b12b4
 
 
d62c791
61b12b4
 
526a8e8
 
 
 
 
67a1218
 
 
526a8e8
 
 
 
 
 
 
 
 
 
67a1218
526a8e8
 
 
 
67a1218
 
 
526a8e8
 
 
 
 
 
 
67a1218
526a8e8
 
 
 
d62c791
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bcceea4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1ccdc0b
fe7e528
 
bcceea4
 
 
fe7e528
1ccdc0b
bcceea4
 
 
 
 
d62c791
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
import sys, os, re
from concurrent.futures import ThreadPoolExecutor, as_completed
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from langchain_groq import ChatGroq
from langchain_core.tools import tool
from langgraph.prebuilt import create_react_agent
from retrieval.pubmed import fetch_pubmed


def get_llm():
    from langchain_openai import ChatOpenAI
    return ChatOpenAI(
        model="Qwen/Qwen2.5-72B-Instruct",
        temperature=0,
        openai_api_key=os.environ.get("VLLM_API_KEY", "EMPTY"),
        openai_api_base=os.environ.get("VLLM_BASE_URL", "http://localhost:8000/v1"),
    )

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."""
    results = fetch_pubmed(query, max_results=5)
    if not results:
        return "No abstracts found for this query."
    out = []
    for r in results:
        out.append("[PMID " + r["pmid"] + "]\n" + r["abstract"])
    return "\n\n".join(out)


def run_query_architect(user_question):
    llm = get_llm()
    prompt = (
        "You are a biomedical librarian expert in PubMed search strategy.\n"
        "Given this clinical question, generate exactly 5 distinct PubMed search queries using "
        "MeSH terminology and clinical keywords to maximise literature coverage.\n"
        "Return ONLY a numbered list 1-5, one query per line, no explanations.\n\n"
        "Question: " + user_question
    )
    response = llm_invoke_with_retry(llm, prompt)
    raw_lines = response.content.strip().split("\n")
    queries = []
    for line in raw_lines:
        clean = re.sub(r"^[\d]+[\.)\s]+", "", line.strip())
        if clean:
            queries.append(clean)
    return queries[:5]


def run_literature_scout(queries):
    from retrieval.pubmed import fetch_europepmc
    all_papers = {}
    import time
    for q in queries:
        time.sleep(0.4)
        for r in fetch_pubmed(q, max_results=5):
            if r["pmid"] not in all_papers:
                all_papers[r["pmid"]] = r
        time.sleep(0.4)
        for r in fetch_europepmc(q, max_results=3):
            if r["pmid"] not in all_papers:
                all_papers[r["pmid"]] = r
    return all_papers


def run_evidence_synthesiser(user_question, papers):
    llm = get_llm()
    parts = []
    for pmid, p in list(papers.items())[:6]:
        title = p.get("title", "N/A")
        abstract = p["abstract"]
        parts.append("[PMID " + pmid + "]\nTitle: " + title + "\n" + abstract)
    corpus = "\n\n".join(parts)
    prompt = (
        "You are a senior biomedical researcher writing a structured evidence synthesis.\n"
        "Answer the clinical question using ONLY this structure:\n\n"
        "## Background\n"
        "Brief context (2-3 sentences).\n\n"
        "## Key Findings\n"
        "Most important findings with PMID citations inline.\n\n"
        "## Level of Evidence\n"
        "Rate: Strong / Moderate / Preliminary. Justify briefly.\n\n"
        "## Conflicting Evidence\n"
        "Any contradictions across studies.\n\n"
        "## Research Gaps\n"
        "What the literature does not answer.\n\n"
        "## Clinical Implications\n"
        "What this means for practice or future research.\n\n"
        "Clinical Question: " + user_question + "\n\n"
        "Retrieved Literature:\n" + corpus + "\n\n"
        "Be precise and cite PMIDs throughout."
    )
    response = llm_invoke_with_retry(llm, prompt)
    return response.content


def run_citation_builder(papers):
    result_lines = []
    for i, (pmid, p) in enumerate(papers.items(), 1):
        title = p.get("title", "Title unavailable")
        authors = p.get("authors", "Authors unavailable")
        journal = p.get("journal", "Journal unavailable")
        year = p.get("year", "n.d.")
        result_lines.append(
            str(i) + ". " + authors + " (" + year + "). " + title + ". " + journal + ". PMID: " + pmid
        )
    return "\n".join(result_lines)


def run_confidence_scorer(synthesis):
    llm = get_llm()
    prompt = (
        "You are a critical appraiser of biomedical evidence.\n"
        "Given this synthesis, score each section for evidence quality.\n"
        "Return ONLY valid JSON, no markdown, no explanation:\n"
        "{\n"
        '  "Background": {"score": 8, "rationale": "one sentence"},\n'
        '  "Key Findings": {"score": 7, "rationale": "one sentence"},\n'
        '  "Level of Evidence": {"score": 6, "rationale": "one sentence"},\n'
        '  "Conflicting Evidence": {"score": 5, "rationale": "one sentence"},\n'
        '  "Research Gaps": {"score": 7, "rationale": "one sentence"},\n'
        '  "Clinical Implications": {"score": 6, "rationale": "one sentence"}\n'
        "}\n\n"
        "Scores: 8-10 = strong evidence, 5-7 = moderate, 1-4 = weak/preliminary.\n\n"
        "Synthesis:\n" + synthesis
    )
    response = llm_invoke_with_retry(llm, prompt)
    import json
    text = response.content.strip()
    text = text.replace("```json", "").replace("```", "").strip()
    return json.loads(text)


def run_selective_review(user_question, selected_papers):
    llm = get_llm()
    parts = []
    for pmid, p in selected_papers.items():
        title = p.get("title", "N/A")
        abstract = p.get("abstract", "")[:400]
        authors = p.get("authors", "")
        year = p.get("year", "")
        parts.append("[PMID " + pmid + "] " + authors + " (" + year + "). " + title + "\n" + abstract)
    corpus = "\n\n".join(parts)
    prompt = (
        "You are an academic writer producing a literature review paragraph.\n"
        "Write a single cohesive academic paragraph (200-300 words) that synthesises "
        "the following selected papers in relation to this question.\n"
        "Cite papers inline by PMID in parentheses e.g. (PMID: 12345678).\n"
        "Write in formal academic prose. No bullet points. No headings.\n\n"
        "Question: " + user_question + "\n\n"
        "Selected Papers:\n" + corpus
    )
    response = llm_invoke_with_retry(llm, prompt)
    return response.content


def run_predictive_model(user_question, synthesis):
    llm = get_llm()
    prompt = (
        "You are a biomedical futurist analyzing research trends.\n"
        "Based on this evidence synthesis, provide two short forecasts:\n\n"
        "## Constructive Forecast\n"
        "2-3 sentences: What directions does the current evidence suggest the field is moving toward? "
        "What findings are likely to be confirmed or expanded in future research?\n\n"
        "## 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"
        "Clinical Question: " + user_question + "\n\n"
        "Synthesis:\n" + synthesis
    )
    response = llm_invoke_with_retry(llm, prompt)
    return response.content


def run_table_extractor(user_question, synthesis, papers):
    llm = get_llm()
    paper_list = []
    for pmid, p in list(papers.items())[:10]:
        authors = p.get("authors", "")
        first_author = authors.split(",")[0].strip() if authors else "Unknown"
        paper_list.append("PMID " + pmid + ": " + first_author + " - " + p.get("title", "N/A") + " (" + p.get("year", "") + ")")
    papers_str = "\n".join(paper_list)
    prompt = (
        "You are a biomedical data extractor.\n"
        "From this synthesis and paper list, extract a comparison table.\n"
        "Return ONLY valid JSON, no markdown, no explanation.\n"
        "Format:\n"
        "{\n"
        '  "title": "Comparison of Methods/Treatments",\n'
        '  "columns": ["Study (PMID)", "Method/Treatment", "Key Metric", "Outcome", "Year"],\n'
        '  "rows": [\n'
        '    ["Smith et al. (PMID: 12345)", "CNN", "Accuracy: 95%", "Positive", "2024"],\n'
        '    ...\n'
        '  ]\n'
        "}\n\n"
        "Rules:\n"
        "- Maximum 8 rows, NO duplicate rows\n"
        "- Each study should appear at most once\n"
        "- Use the actual first author surname from the paper list, never write 'Author'\n"
        "- Only include rows where you have concrete data from the synthesis\n"
        "- If no specific metrics exist, use concise descriptive outcomes\n"
        "- Never invent data\n\n"
        "Clinical Question: " + user_question + "\n\n"
        "Papers:\n" + papers_str + "\n\n"
        "Synthesis:\n" + synthesis[:1500]
    )
    response = llm.invoke(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)
    print("      Generated " + str(len(queries)) + " queries")
    print("[2/4] Literature Scout: fetching PubMed in parallel...")
    papers = run_literature_scout(queries)
    print("      Retrieved " + str(len(papers)) + " unique papers")
    print("[3/4] Evidence Synthesiser: building structured synthesis...")
    synthesis = run_evidence_synthesiser(user_question, papers)
    print("[4/4] Citation Builder: formatting references...")
    citations = run_citation_builder(papers)
    return {
        "question": user_question,
        "queries": queries,
        "paper_count": len(papers),
        "synthesis": synthesis,
        "citations": citations,
        "papers": papers
    }


def build_agent():
    llm = get_llm()
    return create_react_agent(llm, [PubMedSearch])


if __name__ == "__main__":
    question = "What are the most effective ML methods for epilepsy seizure detection from EEG signals?"
    result = run_pipeline(question)
    print("\n=== SYNTHESIS ===")
    print(result["synthesis"])
    print("\n=== REFERENCES ===")
    print(result["citations"])