MedReason-RAG / src /query_expander.py
Tasawar-prog1's picture
Upload 12 files
0ad6f75 verified
Raw
History Blame Contribute Delete
3.57 kB
from typing import Dict, List
# Curated medical knowledge expansion — symptom → likely differential conditions
_EXPANSIONS: Dict[str, List[str]] = {
"chest pain": ["angina pectoris", "myocardial infarction", "aortic dissection", "pulmonary embolism", "GERD"],
"shortness of breath": ["heart failure", "pneumonia", "COPD exacerbation", "pulmonary embolism", "asthma"],
"dyspnea": ["heart failure", "pneumonia", "pulmonary embolism", "COPD", "anemia"],
"fever": ["bacterial infection", "viral infection", "sepsis", "malignancy", "autoimmune disease"],
"headache": ["migraine", "subarachnoid hemorrhage", "meningitis", "hypertensive urgency", "tension headache"],
"dizziness": ["BPPV", "vestibular neuritis", "orthostatic hypotension", "cardiac arrhythmia", "anemia"],
"abdominal pain": ["appendicitis", "cholecystitis", "pancreatitis", "bowel obstruction", "peptic ulcer disease"],
"edema": ["heart failure", "nephrotic syndrome", "hepatic cirrhosis", "deep vein thrombosis", "hypoalbuminemia"],
"syncope": ["vasovagal syncope", "cardiac arrhythmia", "orthostatic hypotension", "pulmonary embolism", "AVNRT"],
"weight loss": ["malignancy", "hyperthyroidism", "diabetes mellitus", "inflammatory bowel disease", "tuberculosis"],
"night sweats": ["lymphoma", "tuberculosis", "HIV", "menopause", "brucellosis"],
"fatigue": ["anemia", "hypothyroidism", "heart failure", "depression", "chronic fatigue syndrome"],
"palpitations": ["atrial fibrillation", "SVT", "ventricular tachycardia", "hyperthyroidism", "anxiety"],
"hemoptysis": ["tuberculosis", "lung cancer", "pulmonary embolism", "bronchiectasis", "pneumonia"],
"hematuria": ["urinary tract infection", "renal cell carcinoma", "bladder cancer", "nephrolithiasis", "glomerulonephritis"],
"jaundice": ["hepatitis", "biliary obstruction", "hemolytic anemia", "cirrhosis", "pancreatic cancer"],
}
class QueryExpander:
"""Expands medical entities into PubMed-ready queries using clinical knowledge."""
def expand(self, entities: Dict[str, List[str]]) -> List[str]:
queries: List[str] = []
symptoms = entities.get("symptoms", [])
labs = entities.get("labs", [])
diagnoses = entities.get("diagnoses", [])
# Query 1: Primary symptom-based differential
if symptoms:
symptom_terms = " OR ".join(f'"{s}"[Title/Abstract]' for s in symptoms[:3])
queries.append(f"({symptom_terms}) AND (differential diagnosis OR clinical features)")
# Query 2: Lab-based evidence
if labs:
lab_terms = " OR ".join(f'"{l}"[Title/Abstract]' for l in labs[:3])
queries.append(f"({lab_terms}) AND (diagnosis OR sensitivity specificity)")
# Query 3: Expanded conditions from symptom knowledge
expanded: List[str] = []
for symptom in symptoms:
expanded.extend(_EXPANSIONS.get(symptom.lower(), []))
if diagnoses:
expanded.extend(diagnoses[:3])
if expanded:
unique_expanded = list(dict.fromkeys(expanded))[:6]
cond_terms = " OR ".join(f'"{c}"' for c in unique_expanded)
queries.append(f"({cond_terms}) AND (systematic review OR meta-analysis OR clinical trial)")
# Fallback: broad search from all entities
if not queries:
all_terms = symptoms + labs + diagnoses
if all_terms:
fallback = " OR ".join(f'"{t}"' for t in all_terms[:5])
queries.append(f"({fallback})")
return queries[:3]