Update app.py

app.py

@@ -2,11 +2,14 @@ import os
 import json
 import pickle
 import re
+from urllib.parse import urlparse, quote_plus
 
 import faiss
 import gradio as gr
 import pandas as pd
+import requests
 import torch
+from bs4 import BeautifulSoup
 from peft import PeftModel
 from sentence_transformers import SentenceTransformer
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
@@ -14,23 +17,25 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 # ============================================================
 # SPACE PATHS / ENV
 # ============================================================
-# In Space Settings -> Variables:
+# Variables:
 # BASE_MODEL_PATH=meta-llama/Llama-3.1-8B-Instruct
 # ALLOW_REBUILD_FROM_CSV=false
 # USE_4BIT=true
 # TOP_K=4
 # SCORE_THRESHOLD=0.40
 # USE_SCORE_THRESHOLD=true
+# ENABLE_WEB_FALLBACK=true
 #
-# In Space Settings -> Secrets:
+# Secrets:
 # HF_TOKEN=hf_xxxxxxxxx
+# SERPAPI_API_KEY=xxxxxxxxx   (optional, for better web search)
 
 BASE_MODEL_PATH = os.getenv("BASE_MODEL_PATH", "meta-llama/Llama-3.1-8B-Instruct")
 HF_TOKEN = os.getenv("HF_TOKEN", None)
+SERPAPI_API_KEY = os.getenv("SERPAPI_API_KEY", None)
 
 ADAPTER_DIR = "adapter_refined_v10"
 
-# IMPORTANT: this must match your actual folder name in the Space repo
 RAG_DIR = "RAGmaterials"
 FAISS_INDEX_PATH = os.path.join(RAG_DIR, "ecg_rag.index")
 META_PATH = os.path.join(RAG_DIR, "ecg_rag_meta.pkl")
@@ -42,6 +47,7 @@ USE_4BIT = os.getenv("USE_4BIT", "true").lower() == "true"
 TOP_K = int(os.getenv("TOP_K", "4"))
 USE_SCORE_THRESHOLD = os.getenv("USE_SCORE_THRESHOLD", "true").lower() == "true"
 SCORE_THRESHOLD = float(os.getenv("SCORE_THRESHOLD", "0.40"))
+ENABLE_WEB_FALLBACK = os.getenv("ENABLE_WEB_FALLBACK", "true").lower() == "true"
 
 # ============================================================
 # EMBEDDING / RAG SETTINGS
@@ -49,36 +55,79 @@ SCORE_THRESHOLD = float(os.getenv("SCORE_THRESHOLD", "0.40"))
 EMBED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
 MAX_CONTEXT_CHARS = 3200
 
+# ============================================================
+# WEB FALLBACK SETTINGS
+# ============================================================
+WEB_TIMEOUT = 12
+MAX_WEB_RESULTS = 3
+MAX_WEB_CHARS = 2200
+
+TRUSTED_MEDICAL_DOMAINS = [
+    # UK / general trusted clinical sources
+    "nice.org.uk",
+    "nhs.uk",
+    "gov.uk",
+    "bhf.org.uk",
+
+    # Major cardiology societies / journals
+    "escardio.org",
+    "heart.org",
+    "acc.org",
+    "ahajournals.org",
+    "hrsonline.org",
+    "nejm.org",
+    "bmj.com",
+    "jamanetwork.com",
+    "pubmed.ncbi.nlm.nih.gov",
+    "ncbi.nlm.nih.gov",
+
+    # ECG / cardiology education sites
+    "litfl.com",              # LITFL ECG Library
+    "ecgwaves.com",
+    "ecgpedia.org",
+    "en.ecgpedia.org",
+    "ecgacademy.com",
+    "ecglibrary.com",
+    "geekymedics.com",
+
+]
+
 # ============================================================
 # GENERATION SETTINGS
 # ============================================================
 SYSTEM_PROMPT = (
     "You are a highly precise medical assistant specializing strictly in cardiology and ECG.\n\n"
     "You must follow ALL rules exactly:\n"
-    "1) Provide only a direct medical definition or explanation of the asked term.\n"
-    "2) Do NOT add background information, causes, risk factors, complications, prognosis, or treatment unless explicitly requested.\n"
-    "3) Limit the response to a maximum of 2 sentences.\n"
-    "4) Use formal medical terminology.\n"
-    "5) Do NOT include citations, references, brackets, or external sources.\n"
-    "6) Do NOT include examples unless explicitly requested.\n"
-    "7) Do NOT ask questions.\n"
-    "8) Do NOT include conversational phrases.\n"
-    "9) Do NOT include notes, formatting labels, or navigation text.\n"
-    "10) Do NOT repeat information.\n"
-    "11) End with a complete declarative sentence.\n"
-    "12) If the answer is uncertain, respond exactly with: I’m not sure about that."
+    "1) Provide only a direct medical definition or explanation of the asked term unless the user asks for management, diagnosis, contraindications, indications, or differential diagnosis.\n"
+    "2) Use only the retrieved evidence provided in CONTEXT and WEB_CONTEXT.\n"
+    "3) Prefer local CONTEXT over WEB_CONTEXT whenever both are available.\n"
+    "4) Do NOT fabricate clinical recommendations or sources.\n"
+    "5) Use formal medical terminology.\n"
+    "6) Do NOT ask questions.\n"
+    "7) Do NOT include conversational phrases.\n"
+    "8) Do NOT include notes, formatting labels, or navigation text.\n"
+    "9) Do NOT repeat information.\n"
+    "10) End with complete declarative sentences.\n"
+    "11) If the answer is uncertain or unsupported, respond exactly with: I’m not sure about that."
 )
 
 RAG_INSTRUCTION = (
     "You MUST follow these rules:\n"
-    "1) Use ONLY the provided CONTEXT to answer.\n"
-    "2) If CONTEXT is empty OR does not contain the answer, say exactly: \"I couldn't find this in my sources.\" "
+    "1) Use ONLY the provided CONTEXT and WEB_CONTEXT to answer.\n"
+    "2) If CONTEXT and WEB_CONTEXT are both empty OR do not contain the answer, say exactly: \"I couldn't find this in my sources.\" "
     "Then give ONLY a very general 1–2 sentence reply.\n"
-    "3) If you use CONTEXT, EVERY sentence MUST end with citations like [S1].\n"
-    "4) Output ONLY the answer text. No headings, no 'Sources:', no 'References:', no URLs.\n"
-    "5) Do NOT paste or quote the CONTEXT.\n"
-    "6) Do NOT repeat yourself.\n"
-    "7) Stop immediately after the final sentence.\n"
+    "3) If you use CONTEXT, every sentence supported by CONTEXT must end with citations like [S1].\n"
+    "4) If you use WEB_CONTEXT, every sentence supported by WEB_CONTEXT must end with citations like [W1].\n"
+    "5) Output ONLY the answer text. No headings, no 'Sources:', no 'References:', no URLs.\n"
+    "6) Do NOT paste or quote the CONTEXT or WEB_CONTEXT.\n"
+    "7) Do NOT repeat yourself.\n"
+    "8) Stop immediately after the final sentence.\n"
+)
+
+MEMORY_INSTRUCTION = (
+    "SESSION_MEMORY and recent conversation are supplementary only.\n"
+    "If SESSION_MEMORY or conversation conflicts with CONTEXT or WEB_CONTEXT, prioritize CONTEXT, then WEB_CONTEXT.\n"
+    "Do not invent patient facts from memory.\n"
 )
 
 MAX_LEN = 2048
@@ -121,15 +170,15 @@ def every_sentence_has_citation(txt: str) -> bool:
     parts = [p.strip() for p in re.split(r"(?<=[.!?])\s+", txt.strip()) if p.strip()]
     if not parts:
         return False
-    return all(re.search(r"\[S\d+\]\s*$", p) for p in parts)
+    return all(re.search(r"\[(S|W)\d+\]\s*$", p) for p in parts)
 
 def keep_only_cited_sentences(txt: str) -> str:
     sents = [p.strip() for p in re.split(r"(?<=[.!?])\s+", txt.strip()) if p.strip()]
-    kept = [s for s in sents if re.search(r"\[S\d+\]\s*$", s)]
+    kept = [s for s in sents if re.search(r"\[(S|W)\d+\]\s*$", s)]
     return " ".join(kept).strip()
 
 def remove_citations(txt: str) -> str:
-    return re.sub(r"\s*\[S\d+\]", "", txt).strip()
+    return re.sub(r"\s*\[(S|W)\d+\]", "", txt).strip()
 
 def clean_text(x: str) -> str:
     x = str(x).replace("\x00", " ").strip()
@@ -254,6 +303,211 @@ def build_context(hits, max_chars=MAX_CONTEXT_CHARS):
 
     return "\n\n".join(blocks)
 
+# ============================================================
+# MEMORY HELPERS
+# ============================================================
+def build_memory_messages(chat_history, max_turns=4):
+    memory_messages = []
+    if not chat_history:
+        return memory_messages
+
+    trimmed = chat_history[-(max_turns * 2):]
+
+    for msg in trimmed:
+        role = msg.get("role")
+        content = (msg.get("content") or "").strip()
+        if role in {"user", "assistant"} and content:
+            memory_messages.append({"role": role, "content": content})
+
+    return memory_messages
+
+def extract_session_memory(state_history):
+    memory = {
+        "age": None,
+        "sex": None,
+        "symptoms": [],
+        "conditions": [],
+        "medications": [],
+        "ecg_findings": [],
+    }
+
+    if not state_history:
+        return memory
+
+    for turn in state_history:
+        user_text = (turn.get("user") or "").lower()
+
+        age_match = re.search(r"\b(\d{1,3})\s*[- ]?\s*year[- ]?old\b", user_text)
+        if age_match:
+            memory["age"] = age_match.group(1)
+
+        if re.search(r"\bmale\b", user_text):
+            memory["sex"] = "male"
+        elif re.search(r"\bfemale\b", user_text):
+            memory["sex"] = "female"
+
+        symptom_keywords = ["chest pain", "palpitations", "syncope", "dyspnoea", "shortness of breath", "dizziness"]
+        condition_keywords = ["atrial fibrillation", "heart failure", "hypertension", "myocardial infarction", "nstemi", "stemi"]
+        medication_keywords = ["aspirin", "clopidogrel", "apixaban", "rivaroxaban", "warfarin", "metoprolol", "bisoprolol", "atorvastatin"]
+        ecg_keywords = ["st elevation", "st depression", "t wave inversion", "qt prolongation", "atrial fibrillation", "bundle branch block"]
+
+        for kw in symptom_keywords:
+            if kw in user_text and kw not in memory["symptoms"]:
+                memory["symptoms"].append(kw)
+
+        for kw in condition_keywords:
+            if kw in user_text and kw not in memory["conditions"]:
+                memory["conditions"].append(kw)
+
+        for kw in medication_keywords:
+            if kw in user_text and kw not in memory["medications"]:
+                memory["medications"].append(kw)
+
+        for kw in ecg_keywords:
+            if kw in user_text and kw not in memory["ecg_findings"]:
+                memory["ecg_findings"].append(kw)
+
+    return memory
+
+def format_session_memory(memory):
+    lines = []
+    for k, v in memory.items():
+        if isinstance(v, list) and v:
+            lines.append(f"{k}: {', '.join(v)}")
+        elif v:
+            lines.append(f"{k}: {v}")
+    return "\n".join(lines) if lines else "No session memory."
+
+# ============================================================
+# TRUSTED WEB FALLBACK
+# ============================================================
+def needs_current_guidance(query: str) -> bool:
+    q = (query or "").lower()
+    trigger_terms = [
+        "guideline", "guidelines", "recommended", "recommendation",
+        "contraindication", "contraindications", "indication", "indications",
+        "management", "treatment", "treat", "anticoagulation", "drug",
+        "dose", "dosing", "safety", "trial", "evidence", "recent", "latest",
+        "current", "monitoring", "screening", "algorithm", "approach",
+        "differential", "diagnosis", "workup", "work-up"
+    ]
+    return any(t in q for t in trigger_terms)
+
+def is_trusted_domain(url: str) -> bool:
+    try:
+        netloc = urlparse(url).netloc.lower()
+        return any(netloc.endswith(domain) for domain in TRUSTED_MEDICAL_DOMAINS)
+    except Exception:
+        return False
+
+def fetch_page_text(url: str, max_chars: int = MAX_WEB_CHARS) -> str:
+    try:
+        resp = requests.get(
+            url,
+            timeout=WEB_TIMEOUT,
+            headers={"User-Agent": "Mozilla/5.0"}
+        )
+        resp.raise_for_status()
+        soup = BeautifulSoup(resp.text, "html.parser")
+
+        for tag in soup(["script", "style", "noscript", "svg"]):
+            tag.extract()
+
+        text = " ".join(soup.get_text(separator=" ").split())
+        return text[:max_chars].strip()
+    except Exception as e:
+        print(f"Failed to fetch page {url}: {e}")
+        return ""
+
+def serpapi_search(query: str, max_results: int = 10):
+    if not SERPAPI_API_KEY:
+        return []
+
+    try:
+        params = {
+            "engine": "google",
+            "q": query,
+            "api_key": SERPAPI_API_KEY,
+            "num": max_results,
+        }
+        resp = requests.get("https://serpapi.com/search.json", params=params, timeout=WEB_TIMEOUT)
+        resp.raise_for_status()
+        data = resp.json()
+
+        results = []
+        for item in data.get("organic_results", []):
+            title = item.get("title", "")
+            url = item.get("link", "")
+            if title and url:
+                results.append({"title": title, "url": url})
+        return results
+    except Exception as e:
+        print(f"SerpAPI search failed: {e}")
+        return []
+
+def duckduckgo_html_search(query: str, max_results: int = 10):
+    try:
+        url = f"https://html.duckduckgo.com/html/?q={quote_plus(query)}"
+        resp = requests.get(
+            url,
+            timeout=WEB_TIMEOUT,
+            headers={"User-Agent": "Mozilla/5.0"}
+        )
+        resp.raise_for_status()
+        soup = BeautifulSoup(resp.text, "html.parser")
+
+        results = []
+        for a in soup.select("a.result__a"):
+            title = " ".join(a.get_text(separator=" ").split()).strip()
+            href = a.get("href", "").strip()
+            if title and href:
+                results.append({"title": title, "url": href})
+            if len(results) >= max_results:
+                break
+
+        return results
+    except Exception as e:
+        print(f"DuckDuckGo HTML search failed: {e}")
+        return []
+
+def web_search_medical(query: str, max_results: int = MAX_WEB_RESULTS):
+    raw_results = serpapi_search(query, max_results=10) if SERPAPI_API_KEY else duckduckgo_html_search(query, max_results=10)
+
+    trusted = []
+    for item in raw_results:
+        url = item.get("url", "")
+        title = item.get("title", "")
+        if not is_trusted_domain(url):
+            continue
+
+        page_text = fetch_page_text(url)
+        if not page_text:
+            continue
+
+        trusted.append({
+            "title": title,
+            "url": url,
+            "text": page_text,
+        })
+
+        if len(trusted) >= max_results:
+            break
+
+    return trusted
+
+def build_web_context(web_hits, max_chars=MAX_WEB_CHARS):
+    blocks = []
+    total = 0
+
+    for i, hit in enumerate(web_hits, 1):
+        block = f"[W{i}]\nTitle: {hit['title']}\nText: {hit['text']}"
+        if total + len(block) > max_chars:
+            break
+        blocks.append(block)
+        total += len(block) + 2
+
+    return "\n\n".join(blocks)
+
 # ============================================================
 # LOAD BASE MODEL + LORA
 # ============================================================
@@ -349,7 +603,7 @@ def _generate(messages):
     return answer
 
 @torch.no_grad()
-def generate_with_rag(user_query: str, show_citations: bool = True):
+def generate_with_rag(user_query: str, chat_history=None, state_history=None, show_citations: bool = True):
     hits = retrieve(user_query, k=TOP_K)
     best_score = hits[0][0] if hits else -1.0
 
@@ -360,11 +614,31 @@ def generate_with_rag(user_query: str, show_citations: bool = True):
         context = build_context(hits)
         used_context = bool(context.strip())
 
-    user_payload = f"QUESTION:\n{user_query}\n\nCONTEXT:\n{context}"
+    web_hits = []
+    web_context = ""
+    used_web = False
+
+    if ENABLE_WEB_FALLBACK and ((not used_context) or (best_score < SCORE_THRESHOLD)) and needs_current_guidance(user_query):
+        web_hits = web_search_medical(user_query, max_results=MAX_WEB_RESULTS)
+        web_context = build_web_context(web_hits)
+        used_web = bool(web_context.strip())
+
+    memory_messages = build_memory_messages(chat_history, max_turns=4) if chat_history else []
+    session_memory = format_session_memory(extract_session_memory(state_history or []))
+
+    user_payload = (
+        f"QUESTION:\n{user_query}\n\n"
+        f"SESSION_MEMORY:\n{session_memory}\n\n"
+        f"CONTEXT:\n{context}\n\n"
+        f"WEB_CONTEXT:\n{web_context}"
+    )
 
     messages = [
         {"role": "system", "content": SYSTEM_PROMPT},
         {"role": "system", "content": RAG_INSTRUCTION},
+        {"role": "system", "content": MEMORY_INSTRUCTION},
+        {"role": "system", "content": "If CONTEXT conflicts with WEB_CONTEXT, prioritize CONTEXT."},
+    ] + memory_messages + [
         {"role": "user", "content": user_payload},
     ]
 
@@ -375,24 +649,24 @@ def generate_with_rag(user_query: str, show_citations: bool = True):
             return True
         if BAD_OUTPUT_PATTERN.search(a):
             return True
-        if used_context and show_citations and (not every_sentence_has_citation(a)):
+        if (used_context or used_web) and show_citations and (not every_sentence_has_citation(a)):
             return True
         return False
 
     if is_invalid(answer):
         strict_format = (
-            "Rewrite using ONLY CONTEXT.\n"
+            "Rewrite using ONLY CONTEXT and WEB_CONTEXT.\n"
             "Output ONLY the answer text.\n"
             "No headings, no 'Sources', no 'References', no URLs.\n"
-            "If not found in CONTEXT: say exactly \"I couldn't find this in my sources.\" "
+            "If nothing useful is found: say exactly \"I couldn't find this in my sources.\" "
             "then give only a very general 1–2 sentence reply.\n"
-            "If you use CONTEXT: EVERY sentence must end with [S#].\n"
+            "If you use CONTEXT or WEB_CONTEXT: EVERY sentence must end with [S#] or [W#].\n"
             "Stop immediately after the final sentence."
         )
         messages2 = messages + [{"role": "user", "content": strict_format}]
         answer = _generate(messages2)
 
-    if used_context and show_citations:
+    if (used_context or used_web) and show_citations:
         cited = keep_only_cited_sentences(answer)
         if cited.strip():
             answer = cited
@@ -400,30 +674,52 @@ def generate_with_rag(user_query: str, show_citations: bool = True):
             if hits:
                 fallback = extract_fallback_answer_from_doc(hits[0][2])
                 answer = f"{fallback} [S1]"
+            elif web_hits:
+                fallback = web_hits[0]["text"][:220].rsplit(" ", 1)[0].strip()
+                if not fallback.endswith("."):
+                    fallback += "."
+                answer = f"{fallback} [W1]"
             else:
                 answer = "I couldn't find this in my sources."
 
     if not show_citations:
         answer = remove_citations(answer)
 
-    return answer, hits, best_score
+    return answer, hits, web_hits, best_score
 
-def format_sources_markdown(hits, max_chars=600):
-    if not hits:
-        return "No retrieved sources for the last answer."
+# ============================================================
+# SOURCE FORMATTING
+# ============================================================
+def format_sources_markdown(hits, web_hits=None, max_chars=600):
+    sections = []
+
+    if hits:
+        lines = []
+        for i, (score, idx, text, m) in enumerate(hits, 1):
+            snippet = text[:max_chars].strip().replace("\n", "\n> ")
+            lines.append(
+                f"### S{i}\n"
+                f"- score: `{score:.3f}`\n"
+                f"- id: `{m.get('id')}`\n"
+                f"- question: {m.get('question', '')}\n\n"
+                f"> {snippet}"
+            )
+        sections.append("\n\n".join(lines))
+    else:
+        sections.append("No retrieved local sources for the last answer.")
 
-    lines = []
-    for i, (score, idx, text, m) in enumerate(hits, 1):
-        snippet = text[:max_chars].strip().replace("\n", "\n> ")
-        lines.append(
-            f"### S{i}\n"
-            f"- score: `{score:.3f}`\n"
-            f"- id: `{m.get('id')}`\n"
-            f"- question: {m.get('question', '')}\n\n"
-            f"> {snippet}"
-        )
+    if web_hits:
+        web_lines = ["## Web sources"]
+        for i, hit in enumerate(web_hits, 1):
+            web_lines.append(
+                f"### W{i}\n"
+                f"- title: {hit['title']}\n"
+                f"- url: {hit['url']}\n\n"
+                f"> {hit['text'][:max_chars].strip()}"
+            )
+        sections.append("\n\n".join(web_lines))
 
-    return "\n\n".join(lines)
+    return "\n\n".join(sections).strip()
 
 # ============================================================
 # GRADIO UI
@@ -438,8 +734,10 @@ def respond(user_text, chat_history, state_history, show_citations):
     if not user_text:
         return "", chat_history, state_history, "No sources yet."
 
-    answer, hits, best_score = generate_with_rag(
+    answer, hits, web_hits, best_score = generate_with_rag(
         user_text,
+        chat_history=chat_history,
+        state_history=state_history,
         show_citations=show_citations
     )
 
@@ -452,10 +750,11 @@ def respond(user_text, chat_history, state_history, show_citations):
         "user": user_text,
         "assistant": answer,
         "hits": hits,
+        "web_hits": web_hits,
         "best_score": best_score,
     }]
 
-    sources_md = format_sources_markdown(hits)
+    sources_md = format_sources_markdown(hits, web_hits=web_hits)
     return "", chat_history, state_history, sources_md
 
 def reset_chat():
@@ -463,7 +762,7 @@ def reset_chat():
 
 with gr.Blocks(title="MrCardio") as demo:
     gr.Markdown("# MrCardio")
-    gr.Markdown("Fine-tuned cardiology / ECG LoRA model with FAISS-based RAG.")
+    gr.Markdown("Fine-tuned cardiology / ECG LoRA model with FAISS-based RAG, short-term memory, and trusted medical web fallback.")
 
     show_citations = gr.Checkbox(value=True, label="Show citations in answers")
     chatbot = gr.Chatbot(type="messages", height=520)