Update pages/02_Workflow_UI.py

pages/02_Workflow_UI.py

@@ -1,141 +1,134 @@
 # pages/02_Workflow_UI.py
-# -----------------------------------------------------------------------------
-# Why this change
-# -----------------------------------------------------------------------------
-# - Aligns the UI with the Phase 2 orchestration:
-#   * Removed obsolete arguments (e.g., rag_toggles that forced mixing)
-#   * Keeps generation natural-language for Assessment/Plan (no RAG injection)
-# - Adds a compact Explainability hook: if src.explainability.render_explainability_block
-#   is present, the page will render the post-hoc block under an expander.
-# - Keeps logging minimal and UI latency obvious.
-
 from __future__ import annotations
 
-import json
+import hashlib
+import os
 import time
 import uuid
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List
 
 import streamlit as st
 
-from src.ai_core import generate_soap_draft
 from src.prompt_builder import build_referral_summary
+from src.ai_core import generate_soap_draft
+from src.explainability import compute_referral_tokens_for_section, segment_claims
+from src.guideline_annotator import find_guideline_refs_for_section, assign_endnotes
+from src.log_utils import write_audit_log
 
-# Optional evidence mapping panel (if present)
-try:
-    from src.reasoning_panel import render_reasoning_panel  # type: ignore
-    _HAS_REASONING_PANEL = True
-except Exception:
-    _HAS_REASONING_PANEL = False
-
-# Explainability provider (render_explainability_block)
+# Canonical UI helpers from Phase 1A prototype (preferred)
 try:
-    from src.explainability import render_explainability_block  # type: ignore
-    _HAS_EXPLAINABILITY = True
+    from src.ui_helpers import render_referral_tokens, render_guideline_refs, ensure_dirs  # type: ignore
 except Exception:
-    _HAS_EXPLAINABILITY = False
-
-
-# --------------------------- Small utilities ---------------------------------
-
-def _parse_semver(v: str) -> tuple:
-    try:
-        parts = v.split(".")
-        major = int(parts[0]); minor = int(parts[1]); patch = int((parts[2] or "0").split("+")[0])
-        return (major, minor, patch)
-    except Exception:
-        return (0, 0, 0)
-
-
-def _render_timings(t0: float, t1: float, label: str = "Total"):
-    st.caption(f"{label} runtime: **{t1 - t0:.2f}s**")
-
-
-def _render_json(data: Dict[str, Any], label: str, key: str):
-    with st.expander(label, expanded=False):
-        st.json(data, expanded=False)
-
-
-def _render_explainability(soap: Dict[str, Any], mapping: Dict[str, Any], raw_text: str, summary: str):
-    if not _HAS_EXPLAINABILITY:
-        st.info("Explainability provider not available.")
-        return
-    try:
-        import inspect
-        sig = inspect.signature(render_explainability_block)
-        kwargs = {}
-        if "soap" in sig.parameters:
-            kwargs["soap"] = soap
-        if "mapping" in sig.parameters:
-            kwargs["mapping"] = mapping
-        if "raw_text" in sig.parameters:
-            kwargs["raw_text"] = raw_text
-        if "summary" in sig.parameters:
-            kwargs["summary"] = summary
-        if kwargs:
-            render_explainability_block(**kwargs)
-        else:
-            # Positional fallback: try soap first
-            render_explainability_block(soap)  # type: ignore
-    except Exception as e:
-        st.warning(f"Explainability failed: {e}")
-
-
-# ----------------------------- Page layout -----------------------------------
-
-st.set_page_config(page_title="AI E-Consult — Workflow", layout="wide")
+    def render_referral_tokens(tokens: List[Dict]) -> None:
+        st.write(tokens)
+    def render_guideline_refs(refs: List[Dict]) -> None:
+        for r in refs:
+            st.markdown(f"- [{r.get('n')}] **{r.get('doc')}** — p.{r.get('page')}: {r.get('excerpt')}")
+    def ensure_dirs(paths: List[str]) -> None:
+        import os
+        for p in paths:
+            os.makedirs(p, exist_ok=True)
+
+EXPLAINABILITY_ENABLED = (os.environ.get("EXPLAINABILITY_ENABLED","1").lower() in {"1","true","yes"})
+
+PAGE_KEY = "workflow_v3"
+
+# ------------------------- helpers -------------------------
+
+def _hash_text(s: str) -> str:
+    return hashlib.sha256((s or "").strip().encode("utf-8")).hexdigest()
+
+def _section_status(ai_text: str, current_text: str) -> Dict[str, Any]:
+    ai_h = _hash_text(ai_text)
+    cur_h = _hash_text(current_text)
+    modified = ai_h != cur_h
+    return {"ai_hash": ai_h, "cur_hash": cur_h, "modified": modified}
+
+def _banner_for_section(section_name: str, modified: bool):
+    if modified:
+        st.warning(f"⚠️ **Manual Modification — {section_name} referral tokens frozen.**")
+    else:
+        st.success(f"✅ **AI‑generated — {section_name}**")
+
+def _setup_session():
+    if PAGE_KEY not in st.session_state:
+        st.session_state[PAGE_KEY] = {
+            "case_id": None,
+            "endnote_registry": {},  # (doc::pX) -> n
+            "assessment": {"ai": "", "text": "", "tokens": [], "refs": []},
+            "plan": {"ai": "", "text": "", "tokens": [], "refs": []},
+            "summary": "",
+            "gen_params": {},
+            "metrics": {"coverage_percent": 0.0, "evidence_count": 0},
+        }
+
+def _calc_coverage(section_text: str, refs: List[Dict]) -> float:
+    claims = segment_claims(section_text)
+    if not claims:
+        return 0.0
+    # Approximate: any refs -> 100% if single claim, else 50% baseline if multi-claim
+    has_any = 1 if refs else 0
+    return 100.0 * (has_any / 2) if len(claims) >= 2 else 100.0 * has_any
+
+# ------------------------- UI -------------------------
+
+st.set_page_config(page_title="AI E‑Consult — Explainability Workflow", layout="wide")
+st.title("E‑Consult Workflow — Assessment/Plan with Explainability")
 
-st.title("Workflow — Intake → Draft → Assessment/Plan")
-
-# Environment/version hint
-try:
-    import transformers  # noqa
-    ver = getattr(transformers, "__version__", "unknown")
-    min_needed = (4, 50, 0)  # recommended for med-gemma paths + generate kwargs
-    ok = _parse_semver(ver) >= min_needed
-    st.caption(
-        f"transformers detected: {ver} — "
-        + ("✅ meets recommended minimum." if ok else f"⚠️ recommend ≥ {'.'.join(map(str, min_needed))}.")
-    )
-except Exception:
-    st.caption("transformers not detected.")
+with st.sidebar:
+    st.header("Inputs & Options")
+    sample = st.selectbox("Load sample case", ["—", "lipids", "ckd_dose"])
+    st.caption("Samples populate intake fields; you can edit them.")
+    st.divider()
+    max_new_tokens = st.slider("Max new tokens", 256, 1200, 700, step=32)
+    temperature = st.slider("Temperature", 0.0, 1.2, 0.2, step=0.05)
+    top_p = st.slider("Top‑p", 0.1, 1.0, 0.95, step=0.05)
+    st.caption("MedGemma generates natural‑language Assessment & Plan (no RAG in prompt).")
 
-# --------------------------- Intake form -------------------------------------
+    st.header("Explainability")
+    st.caption("Referral tokens are post‑hoc (E5 similarity) and **never exported** to EMR.")
+    st.caption("Guideline refs are post‑hoc via FAISS; you can re‑run them after edits.")
 
-with st.form("intake_form", clear_on_submit=False):
+# Intake form
+with st.form("intake_form_v3", clear_on_submit=False):
     c1, c2, c3 = st.columns([1, 1, 2])
     with c1:
-        age = st.number_input("Age", min_value=0, max_value=120, value=63, key="age")
-        sex = st.selectbox("Sex", ["M", "F", "Other"], index=0, key="sex")
+        age = st.number_input("Age", min_value=0, max_value=120, value=63, key="age_v3")
+        sex = st.selectbox("Sex", ["M", "F", "Other"], index=0, key="sex_v3")
     with c2:
-        chief_complaint = st.text_input("Chief complaint", value="Statin intolerance with high LDL", key="chief_complaint")
-        question = st.text_input("Consult question", value="How to manage LDL with statin intolerance?", key="question")
+        chief_complaint = st.text_input("Chief complaint", value="Statin intolerance with high LDL", key="cc_v3")
+        question = st.text_input("Consult question", value="How to manage LDL with statin intolerance?", key="q_v3")
     with c3:
-        history = st.text_area("History / Notes", height=120, key="history", value="Stopped rosuvastatin due to myalgias; FHx premature ASCVD. Prior LDL 190.")
-        medications = st.text_area("Medications", height=80, key="medications", value="Ezetimibe 10 mg daily; rosuvastatin discontinued.")
-        allergies = st.text_area("Allergies", height=60, key="allergies", value="No known drug allergies.")
+        history = st.text_area("History / Notes", height=120, key="hx_v3", value="Stopped rosuvastatin due to myalgias; FHx premature ASCVD. Prior LDL 190.")
+        medications = st.text_area("Medications", height=80, key="meds_v3", value="Ezetimibe 10 mg daily; rosuvastatin discontinued.")
+        allergies = st.text_area("Allergies", height=60, key="allergies_v3", value="No known drug allergies.")
 
-    labs = st.text_area("Labs (optional)", height=80, key="labs", value="LDL 172 mg/dL; HDL 38 mg/dL; TG 210 mg/dL; HbA1c 5.8%")
-    imaging = st.text_area("Imaging (optional)", height=80, key="imaging", value="")
+    labs = st.text_area("Labs (optional)", height=80, key="labs_v3", value="LDL 172 mg/dL; HDL 38 mg/dL; TG 210 mg/dL; HbA1c 5.8%")
+    imaging = st.text_area("Imaging (optional)", height=80, key="img_v3", value="")
     submitted = st.form_submit_button("Generate Draft")
 
-# Toggles / params
-with st.sidebar:
-    st.header("Generation Options")
-    max_new_tokens = st.slider("Max new tokens", 256, 1200, 700, step=32)
-    temperature = st.slider("Temperature", 0.0, 1.2, 0.2, step=0.05)
-    top_p = st.slider("Top-p", 0.1, 1.0, 0.95, step=0.05)
-
-    st.header("Evidence / Explainability")
-    explain = st.checkbox("Enable Explainability (post‑hoc)", value=False, key="explainability_toggle")
-
-    st.caption("Explainability is post‑hoc only; no RAG is injected into prompts.")
-
-# ------------------------- Generate & render ---------------------------------
+# Seed samples
+if sample == "lipids":
+    st.session_state["cc_v3"] = "Statin intolerance with high LDL"
+    st.session_state["q_v3"] = "How to manage LDL with statin intolerance?"
+    st.session_state["hx_v3"] = "Prior LDL 190. Myalgias on multiple statins. No ASCVD event. Considering PCSK9."
+    st.session_state["meds_v3"] = "Ezetimibe 10 mg daily. CoQ10 OTC."
+    st.session_state["labs_v3"] = "LDL 172 mg/dL; HDL 38 mg/dL; TG 210 mg/dL"
+elif sample == "ckd_dose":
+    st.session_state["cc_v3"] = "CKD stage 4 — medication dosing question"
+    st.session_state["q_v3"] = "Adjust DOAC and antibiotics dosing?"
+    st.session_state["hx_v3"] = "eGFR 22 mL/min/1.73m2; HTN; T2DM; recent cellulitis."
+    st.session_state["meds_v3"] = "Apixaban 5 mg BID; metformin 1,000 mg BID; TMP-SMX DS; lisinopril 20 mg daily."
+    st.session_state["labs_v3"] = "Creatinine 2.9 mg/dL; K 4.9 mEq/L; HbA1c 7.8%"
+
+# state init
+_setup_session()
+state = st.session_state[PAGE_KEY]
 
 if submitted:
-    t0 = time.time()
-    # Build normalized referral summary for the generator
+    case_id = str(uuid.uuid4())
+    state["case_id"] = case_id
+
     intake = {
         "age": age,
         "sex": sex,
@@ -148,50 +141,159 @@ if submitted:
         "question": question,
     }
     summary = build_referral_summary(intake)
+    state["summary"] = summary
 
-    # Generate SOAP-like JSON (Assessment/Plan natural-language; no RAG in prompt)
-    soap, raw_text, mapping = generate_soap_draft(
+    t0 = time.time()
+    soap, raw_text = generate_soap_draft(
         referral_summary=summary,
         max_new_tokens=max_new_tokens,
         temperature=temperature,
         top_p=top_p,
-        mode="mapping",  # keep mapping mode to populate registry; still no RAG in prompt
     )
     t1 = time.time()
 
-    # Subjective/Objective editors (simple text)
-    st.subheader("Subjective / Objective")
-    col1, col2 = st.columns(2)
-    with col1:
-        subj = st.text_area("Subjective", value="\n".join(soap.get("subjective", [])), height=180, key="subj_text")
-    with col2:
-        obj = st.text_area("Objective", value="\n".join(soap.get("objective", [])), height=180, key="obj_text")
+    # Prepare AI outputs
+    assessment_ai = "\n".join(soap.get("assessment", []))
+    plan_ai = "\n".join(soap.get("plan", []))
+
+    # Compute referral tokens (frozen thereafter)
+    assess_tokens = compute_referral_tokens_for_section(assessment_ai, summary, top_n=4, min_weight=0.03)
+    plan_tokens = compute_referral_tokens_for_section(plan_ai, summary, top_n=4, min_weight=0.03)
+
+    # Seed guideline refs (initially for AI text)
+    endreg = {}
+    assess_refs_raw = find_guideline_refs_for_section(assessment_ai, top_k_per_claim=3, max_refs_per_claim=2)
+    endreg, assess_refs = assign_endnotes(endreg, assess_refs_raw)
+    plan_refs_raw = find_guideline_refs_for_section(plan_ai, top_k_per_claim=3, max_refs_per_claim=2)
+    endreg, plan_refs = assign_endnotes(endreg, plan_refs_raw)
+
+    # Store in state
+    state["endnote_registry"] = endreg
+    state["assessment"] = {"ai": assessment_ai, "text": assessment_ai, "tokens": assess_tokens, "refs": assess_refs}
+    state["plan"] = {"ai": plan_ai, "text": plan_ai, "tokens": plan_tokens, "refs": plan_refs}
+    state["gen_params"] = {"max_new_tokens": max_new_tokens, "temperature": temperature, "top_p": top_p}
+
+    # Metrics
+    coverage = _calc_coverage(assessment_ai, assess_refs)  # simple proxy
+    evidence_count = len(assess_refs) + len(plan_refs)
+    state["metrics"]["coverage_percent"] = coverage
+    state["metrics"]["evidence_count"] = evidence_count
+
+    # Auto-write audit snapshot
+    try:
+        from transformers import __version__ as tx_ver  # type: ignore
+    except Exception:
+        tx_ver = "unknown"
+
+    payload = {
+        "runtime_sec": (t1 - t0),
+        "sections": {
+            "assessment": {"tokens": state["assessment"]["tokens"], "refs": state["assessment"]["refs"]},
+            "plan": {"tokens": state["plan"]["tokens"], "refs": state["plan"]["refs"]}
+        },
+        "coverage_percent": coverage,
+        "evidence_count": evidence_count,
+        "section_edit_flags": {"assessment": False, "plan": False},
+        "explainability_mode": "posthoc_sim",
+        "model_id": "env:MODEL_ID",
+        "transformers_version": tx_ver,
+        "embedding_model_id": "intfloat/e5-base-v2",
+        "index_build_meta": {"created_at": None, "num_chunks": None},
+        "gpu_cuda": False,
+        "gen_params": state["gen_params"],
+    }
+    try:
+        fp = write_audit_log(case_id, payload)
+        st.caption(f"Audit snapshot saved: `{fp}`")
+    except Exception as e:
+        st.caption(f"Audit log skipped: {e}")
 
-    # Assessment & Plan (AI output, editable)
+# Render editors and explainability if we have state
+if state["assessment"]["ai"]:
     st.subheader("Assessment / Plan")
-    col3, col4 = st.columns(2)
-    with col3:
-        assess_text = st.text_area("Assessment (editable)", value="\n".join(soap.get("assessment", [])), height=220, key="assess_text")
-    with col4:
-        plan_text = st.text_area("Plan (editable)", value="\n".join(soap.get("plan", [])), height=220, key="plan_text")
-
-    # Evidence mapping (if available)
-    if _HAS_REASONING_PANEL:
-        try:
-            render_reasoning_panel(mapping)
-        except Exception as e:
-            st.warning(f"Evidence mapping failed: {e}")
-
-    # Explainability (optional)
-    if explain:
-        st.subheader("Explainability")
-        _render_explainability(soap, mapping, raw_text, summary)
-
-    # Timings
-    st.caption(f"Generation runtime: **{t1 - t0:.2f}s**")
-
-    # Raw JSON (for debugging/export)
-    _render_json(soap, "Generated SOAP JSON", key="soap_json")
+    colA, colB = st.columns(2)
+
+    # ---------------- Assessment ----------------
+    with colA:
+        st.markdown("### Assessment")
+        assessment_text = st.text_area("Edit assessment", value=state["assessment"]["text"], height=220, key="assess_text_v3")
+        state["assessment"]["text"] = assessment_text
+        status = _section_status(state["assessment"]["ai"], assessment_text)
+        _banner_for_section("Assessment", status["modified"])
+
+        if EXPLAINABILITY_ENABLED:
+            with st.expander("Explainability — Assessment", expanded=False):
+                st.caption("Referral cues (AI provenance)")
+                render_referral_tokens(state["assessment"]["tokens"])  # frozen tokens
+                st.caption("Guideline references")
+                if status["modified"]:
+                    if st.button("Re-run Guidelines (Assessment)", key="rerun_assess"):
+                        raw_refs = find_guideline_refs_for_section(assessment_text, top_k_per_claim=3, max_refs_per_claim=2)
+                        new_reg, numbered = assign_endnotes(state["endnote_registry"], raw_refs)
+                        state["endnote_registry"] = new_reg
+                        state["assessment"]["refs"] = numbered
+                render_guideline_refs(state["assessment"]["refs"])
+
+    # ---------------- Plan ----------------
+    with colB:
+        st.markdown("### Plan")
+        plan_text = st.text_area("Edit plan", value=state["plan"]["text"], height=220, key="plan_text_v3")
+        state["plan"]["text"] = plan_text
+        status_p = _section_status(state["plan"]["ai"], plan_text)
+        _banner_for_section("Plan", status_p["modified"])
+
+        if EXPLAINABILITY_ENABLED:
+            with st.expander("Explainability — Plan", expanded=False):
+                st.caption("Referral cues (AI provenance)")
+                render_referral_tokens(state["plan"]["tokens"])  # frozen tokens
+                st.caption("Guideline references")
+                if status_p["modified"]:
+                    if st.button("Re-run Guidelines (Plan)", key="rerun_plan"):
+                        raw_refs = find_guideline_refs_for_section(plan_text, top_k_per_claim=3, max_refs_per_claim=2)
+                        new_reg, numbered = assign_endnotes(state["endnote_registry"], raw_refs)
+                        state["endnote_registry"] = new_reg
+                        state["plan"]["refs"] = numbered
+                render_guideline_refs(state["plan"]["refs"])
+
+    # References block
+    if state["endnote_registry"]:
+        st.divider()
+        st.markdown("### References")
+        # Render combined unique refs by number
+        combined = { (r["n"], r["doc"], r["page"]): r for r in state["assessment"]["refs"] + state["plan"]["refs"] }
+        for key in sorted(combined.keys()):
+            r = combined[key]
+            st.markdown(f"[{r['n']}] **{r['doc']}** — p.{r['page']}")
+
+    # Save snapshot
+    st.divider()
+    if st.button("Save Audit Snapshot", key="save_snapshot_v3"):
+        flags = {
+            "assessment": _section_status(state['assessment']['ai'], state['assessment']['text'])["modified"],
+            "plan": _section_status(state['plan']['ai'], state['plan']['text'])["modified"],
+        }
+        cov = state["metrics"]["coverage_percent"]
+        payload = {
+            "runtime_sec": None,
+            "sections": {
+                "assessment": {"tokens": state["assessment"]["tokens"], "refs": state["assessment"]["refs"]},
+                "plan": {"tokens": state["plan"]["tokens"], "refs": state["plan"]["refs"]}
+            },
+            "coverage_percent": cov,
+            "evidence_count": len(state["assessment"]["refs"]) + len(state["plan"]["refs"]),
+            "section_edit_flags": flags,
+            "explainability_mode": "posthoc_sim",
+            "model_id": "env:MODEL_ID",
+            "transformers_version": "auto",
+            "embedding_model_id": "intfloat/e5-base-v2",
+            "index_build_meta": {"created_at": None, "num_chunks": None},
+            "gpu_cuda": False,
+            "gen_params": state["gen_params"],
+        }
+        case_id = state["case_id"] or str(uuid.uuid4())
+        fp = write_audit_log(case_id, payload)
+        st.success(f"Snapshot saved to `{fp}`")
+