Update pages/02_Workflow_UI.py

pages/02_Workflow_UI.py

@@ -3,20 +3,28 @@
 # Phase 3 — Tabbed Workflow UI (PCP Referral ↔ Specialist Review)
 # Fully integrated with V2 backend: store, billing, modal_templates, config,
 # explainability, guideline_annotator, ai_core.
+# Consolidated fixes:
+# - Clinical Summary (consult.summary) + Meds + Labs
+# - Safe seeding, New Case, Reset Demo
+# - SOAP list flattening + reliable UI refresh
+# - Explainability/guidelines visibility + re-run
+# - Billing autosuggest + finalize exports + modals
 # -----------------------------------------------------------------------------
 
 from __future__ import annotations
 
 import os
+import json
+import shutil
 from pathlib import Path
-from typing import Any, Dict, List, Tuple, Optional
+from typing import Any, Dict, List, Optional
 
 import streamlit as st
 
 from src import store, billing, modal_templates, config
 from src.styles import inject_base_css
 from src.paths import faiss_index_dir, initialize_environment
-from src.prompt_builder import build_referral_summary, normalize_intake
+from src.prompt_builder import build_referral_summary
 from src.explainability import chips_from_text, ensure_chip_schema
 from src.guideline_annotator import annotate_guidelines
 from src.ai_core import generate_soap_draft
@@ -28,7 +36,7 @@ inject_base_css()
 initialize_environment()
 
 st.title("Step 2 — Workflow")
-st.caption("PCP Referral → Specialist Review. Demo only — de‑identified data; not for clinical use.")
+st.caption("PCP Referral → Specialist Review. Demo only — de-identified data; not for clinical use.")
 
 # --------------------------- Helpers ------------------------------------
 
@@ -70,9 +78,7 @@ def _as_text(x: Any) -> str:
     if isinstance(x, str):
         return x
     if isinstance(x, list):
-        # flatten list items to bullet lines
-        items = [str(s).strip() for s in x if str(s).strip()]
-        return "\n".join(f"- {s}" for s in items)
+        return "\n".join(str(s) for s in x)
     return str(x)
 
 def _soap_is_empty(soap: Dict[str, str]) -> bool:
@@ -135,7 +141,9 @@ def _note_markdown(case: Dict[str, Any], summary: str, soap: Dict[str, str], end
     )
     cites = ""
     if endnotes:
-        cites = "\n## Guideline Citations\n" + "\n".join([f"- [{e.get('n',i+1)}] {e.get('doc','Guideline')} p.{e.get('page','')}" for i, e in enumerate(endnotes)])
+        cites = "\n## Guideline Citations\n" + "\n".join(
+            [f"- [{e.get('n', i+1)}] {e.get('doc','Guideline')} p.{e.get('page','')}" for i, e in enumerate(endnotes)]
+        )
     return header + body + cites + "\n"
 
 def _seed_once() -> None:
@@ -151,8 +159,7 @@ def _case_options() -> List[Dict[str, Any]]:
     return items
 
 def _get_case(case_id: str) -> Dict[str, Any]:
-    case = store.read_case(case_id) or {}
-    return case
+    return store.read_case(case_id) or {}
 
 # --------------------------- Top banners --------------------------------
 
@@ -180,9 +187,8 @@ if not items:
     st.error("No cases available.")
     st.stop()
 
-# --- Optional: create a new case for ad hoc testing ---
+# --- New Case ---
 if st.button("➕ New Case", use_container_width=True):
-    # Generate a new unique case ID
     new_id = store.new_case_id()
     store.create_case({
         "case_id": new_id,
@@ -196,28 +202,18 @@ if st.button("➕ New Case", use_container_width=True):
             "labs": "",
             "consent_obtained": False,
         },
-        "soap_draft": {
-            "subjective": "",
-            "objective": "",
-            "assessment": "",
-            "plan": "",
-        },
+        "soap_draft": {"subjective": "", "objective": "", "assessment": "", "plan": ""},
         "billing": {"minutes": 5, "spoke": False, "cpt_code": None, "attested": False},
         "explainability": {"manually_modified": False},
     })
     st.success(f"Created new case {new_id}. Refreshing list…")
     st.rerun()
 
-# --- Demo reset control ---
-import shutil
-from src import config
-
+# --- Reset Demo ---
 if st.button("🗑️ Reset Demo (clear all cases)", use_container_width=True):
-    cases_dir = config.paths_dict()["cases_dir"]
     try:
-        shutil.rmtree(cases_dir)
+        shutil.rmtree(config.paths_dict()["cases_dir"])
         st.success("✅ All demo cases cleared. The 4 default draft cases will be recreated on reload.")
-        # Reset seeding flag so defaults are reloaded automatically
         st.session_state["_seeded"] = False
         st.rerun()
     except Exception as e:
@@ -238,14 +234,11 @@ st.session_state["_current_case_id"] = selected
 case = _get_case(selected)
 _case_summary_header(case)
 
-
 # --------------------------- Tabs ----------------------------------------
 
 tab_pcp, tab_spec = st.tabs(["PCP Referral", "Specialist Review"])
 
 # ===== PCP Referral Tab =====
-# Inside the PCP Referral tab (replace your existing block):
-
 with tab_pcp:
     is_draft = case.get("status") == "draft"
     p = case.get("patient", {}) or {}
@@ -259,9 +252,8 @@ with tab_pcp:
             "Clinical Summary",
             value=c.get("summary", "") or c.get("history", ""),
             height=160,
-            placeholder=("Example: 72yo man with NYHA class III HFrEF (EF 30%). "
-                         "BP 110/70 HR 62. Echo: mild MR, no effusion. "
-                         "BNP 230; here for med optimization."),
+            placeholder=("E.g., NYHA III HFrEF (EF 30%). BP 110/70 HR 62, euvolemic. "
+                         "Echo: EF 30%, mild MR, no effusion. Stable; med optimization."),
             disabled=not is_draft,
             key=f"summary_{selected}",
         )
@@ -301,7 +293,6 @@ with tab_pcp:
         else:
             st.info("Referral is not editable (status is In Review or Completed).")
 
-
 # ===== Specialist Review Tab =====
 with tab_spec:
     status = case.get("status")
@@ -326,7 +317,6 @@ with tab_spec:
 
         st.markdown("**CPT Autosuggest**")
         elig_codes = [s.code for s in suggestions if s.eligible]
-        # Show suggestion cards
         for s in suggestions:
             with st.container():
                 status_icon = "✅" if s.eligible else "⚠️"
@@ -341,23 +331,18 @@ with tab_spec:
             chosen = None
 
         attested_default = bool((case.get("billing", {}) or {}).get("attested", False))
-        attested = st.checkbox("I attest that the consult was performed and documented per interprofessional e‑consult requirements.", value=attested_default if read_only else False, disabled=read_only, key=f"att_{selected}")
+        attested = st.checkbox("I attest that the consult was performed and documented per interprofessional e-consult requirements.", value=attested_default if read_only else False, disabled=read_only, key=f"att_{selected}")
 
         can_finalize = (not read_only) and (chosen is not None) and attested and (minutes >= 5)
-
-        # Mirror finalize CTA on right
         finalize_right = st.button("Finalize Consult Note", type="primary", disabled=not can_finalize, use_container_width=True, key=f"finalize_r_{selected}")
 
     # ---------- LEFT: Referral & SOAP ----------
     with left:
         st.subheader("Referral Summary")
         summary = build_referral_summary(case)
-
-        # --- Debug expander for referral summary (shows text sent to MedGemma) ---
         with st.expander("🔍 Referral Summary (debug view)", expanded=False):
             st.code(summary, language="markdown")
 
-
         st.subheader("SOAP Draft")
         soap = case.get("soap_draft", {}) or {"subjective": "", "objective": "", "assessment": "", "plan": ""}
         gen_needed = (status == "submitted") and _soap_is_empty(soap)
@@ -365,7 +350,6 @@ with tab_spec:
         if gen_needed and not read_only:
             if st.button("Generate SOAP Draft (LLM)", key=f"gen_{selected}"):
                 try:
-                    # Use normalized intake (patient + consult)
                     result = generate_soap_draft(
                         intake=case,
                         mode="mapping",
@@ -374,47 +358,28 @@ with tab_spec:
                         top_p=0.95,
                     )
                     s = result.get("soap", {}) or {}
-
-                    # --- Normalize and flatten list fields ---
                     subj = _as_text(s.get("subjective", ""))
                     obj = _as_text(s.get("objective", ""))
 
                     assess_raw = s.get("assessment", "")
                     plan_raw = s.get("plan", "")
-
-                    # Convert lists into newline-separated strings
                     assess = "\n".join(assess_raw) if isinstance(assess_raw, list) else str(assess_raw)
                     plan = "\n".join(plan_raw) if isinstance(plan_raw, list) else str(plan_raw)
 
-                except Exception as e:  # graceful fallback
-                    subj = summary
-                    obj = ""
-                    assess = "—"
-                    plan = "—"
+                except Exception as e:
+                    subj, obj, assess, plan = summary, "", "—", "—"
                     st.warning(f"LLM generation unavailable; seeded a minimal draft. ({type(e).__name__})")
                     result = {"error": str(e)}
 
-                # --- Debug expander for raw MedGemma output (full JSON response) ---
                 with st.expander("🧠 MedGemma raw response (debug view)", expanded=False):
-                    import json
                     st.code(json.dumps(result, indent=2), language="json")
 
-                # Persist generated draft and mark unmodified so chips/guidelines show
                 store.update_case(selected, {
-                    "soap_draft": {
-                        "subjective": subj.strip(),
-                        "objective": obj.strip(),
-                        "assessment": assess.strip(),
-                        "plan": plan.strip(),
-                    },
+                    "soap_draft": {"subjective": subj.strip(), "objective": obj.strip(), "assessment": assess.strip(), "plan": plan.strip()},
                     "explainability": {"manually_modified": False},
                 })
-
-                # Refresh case so UI shows the new SOAP fields immediately
-                case = _get_case(selected)
-                soap = case.get("soap_draft", {})
-                st.success("SOAP draft generated and saved. You can review it below.")
-                
+                # Force clean reload so text areas show new content
+                st.rerun()
 
         c1, c2 = st.columns(2)
         with c1:
@@ -437,21 +402,18 @@ with tab_spec:
                     "soap_draft": {"subjective": subj_new, "objective": obj_new, "assessment": assess_new, "plan": plan_new},
                     "explainability": {"manually_modified": True},
                 })
-                # Re-read for consistency
-                case = _get_case(selected)
-                soap = case.get("soap_draft", soap)
+                # Keep on same run; next interactions will show re-computed state
 
-        # Post‑hoc Explainability & Guidelines
+        # Post-hoc Explainability & Guidelines
         exp = (case.get("explainability", {}) or {})
         manually_modified = bool(exp.get("manually_modified", False))
-
         if read_only:
             manually_modified = False  # allow showing in completed
 
         if not manually_modified and (soap.get("plan") or "").strip():
-            with st.expander("Explainability tokens (post‑hoc)", expanded=False):
-                _render_token_chips(_explainability_chips(soap.get("plan","")))
-            with st.expander("Guideline citations (post‑hoc)", expanded=True):
+            with st.expander("Explainability tokens (post-hoc)", expanded=False):
+                _render_token_chips(ensure_chip_schema(chips_from_text(soap.get("plan",""))))
+            with st.expander("Guideline citations (post-hoc)", expanded=True):
                 g = _run_guidelines_cached(selected, soap.get("plan",""))
                 warn = (g.get("warning") or "").strip()
                 if warn:
@@ -468,13 +430,12 @@ with tab_spec:
                     st.caption("_No references found._")
         else:
             if not read_only and (status == "submitted"):
-                if st.button("Re‑run Guidelines", key=f"rerun_{selected}"):
-                    # Re-run and mark unmodified so blocks show
+                if st.button("Re-run Guidelines", key=f"rerun_{selected}"):
                     _run_guidelines_cached(selected, plan_new)
                     store.update_case(selected, {"explainability": {"manually_modified": False}})
                     st.rerun()
                 else:
-                    st.caption("Guideline citations hidden after edits. Click **Re‑run Guidelines** to refresh.")
+                    st.caption("Guideline citations hidden after edits. Click **Re-run Guidelines** to refresh.")
 
         # Mirror finalize CTA on left
         finalize_left = st.button("Finalize Consult Note", type="primary", disabled=not can_finalize, key=f"finalize_l_{selected}")
@@ -506,8 +467,7 @@ with tab_spec:
             rate = float(getattr(pick, "rate", 0.0)) if pick else 0.0
             claim = billing.build_837_claim(case, code=str(chosen), rate=rate, minutes=int(minutes), spoke=bool(spoke), attested=True)
             claim_path = config.make_export_path(selected, "Sample 837 Json.json")
-            import json as _json
-            Path(claim_path).write_text(_json.dumps(claim, ensure_ascii=False, indent=2), encoding="utf-8")
+            Path(claim_path).write_text(json.dumps(claim, ensure_ascii=False, indent=2), encoding="utf-8")
 
             # Status → completed
             store.set_status(selected, "completed")
@@ -521,21 +481,6 @@ with tab_spec:
             st.caption(f"Claim: {claim_path}")
             st.rerun()
 
-    # If already completed, surface actions to preview artifacts again
-    if status == "completed":
-        st.markdown("---")
-        st.subheader("Completed Actions")
-        if st.button("Send Consult Note to EHR (Preview)", key=f"preview_note_{selected}"):
-            # Recompose or read the most recent note (rebuild for simplicity)
-            g = _run_guidelines_cached(selected, (soap.get("plan") or ""))
-            endnotes = g.get("endnotes", [])
-            note_md = _note_markdown(case, summary, soap, endnotes=endnotes)
-            modal_templates.show_consult_note_preview(note_md)
-        if st.button("Submit 837 Claim (Preview)", key=f"preview_claim_{selected}"):
-            pick = next((s for s in billing.autosuggest_cpt(minutes=minutes, spoke=spoke) if s.code == (case.get('billing',{}).get('cpt_code'))), None)
-            rate = float(getattr(pick, "rate", 0.0)) if pick else 0.0
-            claim = billing.build_837_claim(case, code=str(case.get('billing',{}).get('cpt_code')), rate=rate, minutes=int(minutes), spoke=bool(spoke), attested=True)
-            modal_templates.show_837_claim_preview(claim)