Update app.py

app.py

@@ -412,7 +412,11 @@ _ANATOMY_OR_HISTORY = re.compile(
     re.I
 )
 _PAPERS_INTENT = re.compile(r"\b(key\s+papers|suggest\s+papers|landmark|seminal|important|top\s+papers)\b", re.I)
-CITE_TRIGGER = re.compile(r"\b(cite|citations?|references?)\b", re.I)
+CITE_TRIGGER = re.compile(
+    r"\b(?:and\s+(?:cite|citations?|references?|studies?|papers?))\.?\s*$",
+    re.IGNORECASE,
+)
+
 
 # ================== PUBMED & RETRIEVAL ==================
 def fetch_pubmed_chunks(query_or_pmid: str, max_papers: int = 3) -> List[Dict[str, Any]]:
@@ -757,65 +761,110 @@ def direct_llm_fallback(question: str) -> str:
 
 # ================== PUBLIC API ==================
 def ask(question: str) -> str:
-    q = question.strip()
+    q = (question or "").strip()
+
+    # --- PMID short-circuit ---------------------------------------------------
     m = re.search(r"pmid[:\s]*(\d+)", q, re.IGNORECASE)
     if m:
         pmid = m.group(1)
         chunks = fetch_pubmed_chunks(pmid, max_papers=1)
         return "\n".join(c.get("text", "") for c in chunks) or "Sorry, no abstract found."
 
-    if _PAPERS_INTENT.search(q) or CITE_TRIGGER.search(q):
-        core_q = re.sub(CITE_TRIGGER, "", q, count=1, flags=re.I).strip().rstrip(".")
-        core_q = re.sub(_PAPERS_INTENT, "", core_q, flags=re.I).strip()
-        if not core_q:
-            core_q = "CT/QCT structural rigidity femur hip finite element"
-        compact = _compact_terms(core_q)
-        pm_query = (
-            f'(({compact}) AND (hip[TiAb] OR femur[TiAb] OR femoral[TiAb])) AND '
-            '("Finite Element Analysis"[MeSH Terms] OR finite element[TiAb] OR QCT[TiAb] OR CT[TiAb] OR rigidity[TiAb]) '
-            'AND ("2000"[DP] : "2025"[DP])'
-        )
-        cits = fetch_pubmed_citations(pm_query, max_results=5)
-        if not cits:
-            lab = detect_lab(core_q)
-            pm_query = build_lab_query(core_q, lab=lab)
-            cits = fetch_pubmed_citations(pm_query, max_results=5)
-        if not cits:
-            cits = _fallback_cits_for("EA")
-        # Provide a short explanation + citations
-        explanation = _answer_from_chunks(retrieve_context(core_q, top_k=3), core_q) or direct_llm_fallback(core_q)
-        explanation = _post_clean(explanation)
-        if cits:
-            explanation += "\n\nCitations:\n" + "\n".join(cits)
-        return _ensure_min_answer(explanation)
-
-    det = deterministic_definitions_text(q)
-    if det:
+    # --- Normalize query & detect 'and cite' / 'key papers' intent ------------
+    wants_cite = bool(CITE_TRIGGER.search(q) or _PAPERS_INTENT.search(q))
+    core_q = CITE_TRIGGER.sub("", q).strip().rstrip(".")
+    core_q = _PAPERS_INTENT.sub("", core_q).strip()
+
+    # --- Deterministic definitions/workflows (used in both paths) -------------
+    det_text = deterministic_definitions_text(core_q)
+
+    # ==========================================================================
+    # Citation path
+    # ==========================================================================
+    if wants_cite:
+        if det_text:
+            explanation = det_text
+            lq = core_q.lower()
+            used_term = None
+            if ("torsion" in lq) or ("gj" in lq):
+                used_term = "GJ"
+                pm_query = (
+                    '(torsion[TiAb] OR "polar moment"[TiAb] OR GJ[TiAb]) AND '
+                    '("Bone and Bones"[MeSH] OR Femur[TiAb]) AND '
+                    '("Finite Element Analysis"[MeSH] OR QCT[TiAb] OR CT[TiAb]) AND '
+                    '("2000"[DP] : "2025"[DP])'
+                )
+            elif ("bending" in lq) or ("ei" in lq):
+                used_term = "EI"
+                pm_query = (
+                    '(bending[TiAb] OR "second moment"[TiAb] OR EI[TiAb]) AND '
+                    '("Bone and Bones"[MeSH] OR Femur[TiAb]) AND '
+                    '("Finite Element Analysis"[MeSH] OR QCT[TiAb] OR CT[TiAb]) AND '
+                    '("2000"[DP] : "2025"[DP])'
+                )
+            else:
+                used_term = "EA"
+                pm_query = (
+                    '("axial rigidity"[TiAb] OR EA[TiAb] OR "axial stiffness"[TiAb]) AND '
+                    '("Bone and Bones"[MeSH] OR Femur[TiAb]) AND '
+                    '("Finite Element Analysis"[MeSH] OR QCT[TiAb] OR CT[TiAb]) AND '
+                    '("2000"[DP] : "2025"[DP])'
+                )
+            citations = fetch_pubmed_citations(pm_query, max_results=5)
+            if not citations and used_term:
+                citations = _fallback_cits_for(used_term)
+        else:
+            # Grounded explanation + progressively broader citation search
+            explanation = _answer_from_chunks(retrieve_context(core_q, top_k=5), core_q)
+
+            # Broad, de-biased search first
+            compact = _compact_terms(core_q)
+            pm_query = (
+                f'({compact}) AND ("Bone and Bones"[MeSH] OR Femur[TiAb] OR Hip[TiAb] '
+                f'OR Rigidity[TiAb] OR "Tomography, X-Ray Computed"[MeSH] OR "Finite Element Analysis"[MeSH]) '
+                f'NOT (heart[TiAb] OR cardiac[TiAb] OR brain[TiAb] OR skull[TiAb] OR EGFR[TiAb]) '
+                f'AND ("2000"[DP] : "2025"[DP])'
+            )
+            citations = fetch_pubmed_citations(pm_query, max_results=5)
+
+            # If empty, bias toward lab-authored queries that are likely relevant
+            if not citations:
+                lab = detect_lab(core_q)
+                pm_query = build_lab_query(core_q, lab=lab)
+                citations = fetch_pubmed_citations(pm_query, max_results=5)
+
+        resp = _post_clean(explanation)
+        if citations:
+            resp += "\n\nCitations:\n" + "\n".join(citations)
+        else:
+            resp += f'\n\nSorry, no relevant citations found for “{core_q}.”'
+        return _ensure_min_answer(resp)
+
+    # ==========================================================================
+    # Non-citation path
+    # ==========================================================================
+    if det_text:
         dlog("ASK", "Deterministic definition/workflow fired")
-        return det
+        return det_text
 
+    # If the query doesn't clearly hit MSK/FE tokens, try retrieval then fallback
     if not (_MSK_MUST.search(q) or _is_fe_override(q)):
-        chunks = retrieve_context(q, top_k=3)
+        chunks = retrieve_context(q, top_k=5)
         if chunks:
             answer = _answer_from_chunks(chunks, q)
-            # tiny safety to release VRAM between turns
-            try:
-                torch.cuda.empty_cache()
-            except Exception:
-                pass
             return _ensure_min_answer(_post_clean(answer)) or direct_llm_fallback(q)
         return direct_llm_fallback(q)
 
-    chunks = retrieve_context(q, top_k=3)
+    # Clear MSK/FE intent → do the normal grounded path
+    chunks = retrieve_context(q, top_k=5)
     if not chunks:
         return direct_llm_fallback(q)
+
     answer = _answer_from_chunks(chunks, q)
-    try:
-        torch.cuda.empty_cache()
-    except Exception:
-        pass
     return _ensure_min_answer(_post_clean(answer)) or direct_llm_fallback(q)
 
+
+
 def deterministic_definitions_text(core_q: str) -> Optional[str]:
     q_lower = core_q.lower()
     if "define axial rigidity" in q_lower or "what is axial rigidity" in q_lower: