Update pages/02_Workflow_UI.py

pages/02_Workflow_UI.py

@@ -2,14 +2,15 @@
 # -----------------------------------------------------------------------------
 # Phase-3+ Simplified — 3-tab workflow
 # Tabs:
-#   1) PCP Referral (unchanged)
-#   2) Specialist Review (SOAP + Guideline Rationale + 3 specialist controls)
-#   3) Documentation & Billing (Completed cases only: EHR note expander, CPT, 837)
+#   1) PCP Referral
+#   2) Specialist Review (SOAP + Guideline Rationale + time/spoke/attestation)
+#   3) Documentation & Billing (completed cases only: EHR note, CPT + ICD, 837 JSON)
 #
-# Notes:
-# - No Referral Rationale. Single explainability layer: Guideline Rationale (bullets).
-# - Guideline Rationale persists into case JSON to survive page switches.
-# - SOAP draft generation triggers st.rerun() to populate text areas immediately.
+# Assumes:
+#   - No Referral Rationale
+#   - Single Guideline Rationale layer (bullets) via guideline_annotator
+#   - ai_core.generate_soap_draft returns 4-string SOAP keys
+#   - billing.autosuggest_cpt + billing.autosuggest_icd + billing.build_837_claim
 # -----------------------------------------------------------------------------
 
 from __future__ import annotations
@@ -19,6 +20,7 @@ import shutil
 import html
 from pathlib import Path
 from typing import Any, Dict, List, Optional
+from datetime import datetime
 
 import streamlit as st
 
@@ -30,7 +32,6 @@ from src.explainability import text_hash, normalize_text, is_stale
 from src.guideline_annotator import generate_guideline_rationale
 from src.ai_core import generate_soap_draft
 from src.model_loader import active_model_status
-from datetime import datetime
 
 
 # --------------------------- Page Setup ---------------------------------
@@ -133,10 +134,9 @@ def _latest_export(case_id: str, pattern_suffix: str) -> Optional[Path]:
     matches = sorted(exp.glob(f"EC-{case_id}_*{pattern_suffix}"), key=lambda p: p.stat().st_mtime)
     return matches[-1] if matches else None
 
-
 def _note_markdown(case: Dict[str, Any], summary: str, soap: Dict[str, str],
                    guideline_points: List[str], endnotes: List[Dict[str, Any]]) -> str:
-    """Generate formatted EHR-style consultation report (matches sample format)."""
+    """Generate formatted EHR-style consultation report (matches sample-like format)."""
     p = case.get("patient", {}) or {}
     c = case.get("consult", {}) or {}
     consult_id = case.get("case_id", "")
@@ -191,6 +191,23 @@ def _note_markdown(case: Dict[str, Any], summary: str, soap: Dict[str, str],
 
     return header + referral_section + soap_section + guideline_section + endnote_section + attestation
 
+def _soap_is_empty(soap: Dict[str, str]) -> bool:
+    return not any((soap or {}).get(k, "").strip() for k in ("subjective", "objective", "assessment", "plan"))
+
+def _seed_once() -> None:
+    if not st.session_state.get("_seeded", False):
+        store.seed_cases(reset=True)
+        st.session_state["_seeded"] = True
+
+def _case_options() -> List[Dict[str, Any]]:
+    items = store.list_cases()
+    if not items:
+        _seed_once()
+        items = store.list_cases()
+    return items
+
+def _get_case(case_id: str) -> Dict[str, Any]:
+    return store.read_case(case_id) or {}
 
 
 # --------------------------- Global callouts --------------------------------
@@ -243,15 +260,15 @@ if st.button("➕ New Case", use_container_width=True):
     st.success(f"Created new case {new_id}. Refreshing list…")
     st.rerun()
 
-# Reset demo
+# Reset demo (use seed reset; don't rely on internal paths)
 if st.button("🗑️ Reset Demo (clear all cases)", use_container_width=True):
     try:
-        shutil.rmtree(store._cases_dir(), ignore_errors=True)  # type: ignore
-        st.success("✅ All demo cases cleared. The default draft cases will be recreated on reload.")
-        st.session_state["_seeded"] = False
+        store.seed_cases(reset=True)
+        st.success("✅ All demo cases cleared and re-seeded.")
+        st.session_state["_seeded"] = True
         st.rerun()
     except Exception as e:
-        st.error(f"Failed to clear cases: {e}")
+        st.error(f"Failed to reset cases: {e}")
 
 # Selection
 default_case_id = st.session_state.get("_current_case_id") or items[0]["case_id"]
@@ -335,23 +352,13 @@ with tab_spec:
         st.info("Submit the referral on the PCP tab to activate Specialist Review.")
         st.stop()
 
-        # Prevent guideline rationale or re-render during immediate SOAP rerun
-        if st.session_state.pop("_suppress_guidelines", False):
-            st.stop()
-
-
     read_only = status == "completed"
 
     st.subheader("SOAP Draft")
+    soap = case.get("soap_draft", {}) or {"subjective": "", "objective": "", "assessment": "", "plan": ""}
 
-    # Cache SOAP in session to avoid stale disk reads
-    st.session_state[f"soap_{selected}"] = case.get("soap_draft", {})
-    soap = st.session_state[f"soap_{selected}"]
-
-
-    
     # Generate SOAP if needed
-    gen_needed = (status == "submitted") and (not any(soap.get(k, "").strip() for k in ("subjective", "objective", "assessment", "plan")))
+    gen_needed = (status == "submitted") and _soap_is_empty(soap)
     if gen_needed and not read_only:
         if st.button("Generate SOAP Draft (LLM)", key=f"gen_{selected}"):
             try:
@@ -377,14 +384,12 @@ with tab_spec:
                 st.warning(f"LLM generation unavailable; seeded a minimal draft. ({type(e).__name__})")
                 result = {"error": str(e)}
 
-            # Persist and force rerun so text areas populate instantly
             store.update_case(selected, {
                 "soap_draft": {"subjective": subj.strip(), "objective": obj.strip(), "assessment": assess.strip(), "plan": plan.strip()},
                 "explainability": explain_patch,
             })
-            # Set flag so guideline rationale doesn't run on immediate rerun
-            st.session_state["_suppress_guidelines"] = True
-            st.rerun()
+            # allow simple tab toggle to show new SOAP; no extra rerun complexity here
+            st.success("SOAP draft generated. If fields appear empty, toggle tabs once to refresh the view.")
 
     # Editable fields
     c1, c2 = st.columns(2)
@@ -395,7 +400,7 @@ with tab_spec:
     assess_new = st.text_area("Assessment", value=soap.get("assessment",""), height=140, disabled=read_only, key=f"assess_{selected}")
     plan_new   = st.text_area("Plan", value=soap.get("plan",""), height=180, disabled=read_only, key=f"plan_{selected}")
 
-    # Persist edits (whitespace-insensitive)
+    # Persist edits (whitespace-insensitive) while in submitted state
     if (not read_only) and (status == "submitted"):
         changed = any([
             normalize_text(subj_new)   != normalize_text(soap.get("subjective") or ""),
@@ -416,51 +421,44 @@ with tab_spec:
 
     stale = is_stale(assess_new, assess_hash0) or is_stale(plan_new, plan_hash0)
 
-    # Guard: skip guideline generation until SOAP fields exist
-    if not any((soap.get(k, "") or "").strip() for k in ("subjective", "objective", "assessment", "plan")):
-        st.info("Generate SOAP draft first to enable guideline rationale.")
-        st.stop()
-
-
-
-    if not stale and (normalize_text(assess_new + plan_new)):
-        # Use stored bullets if available; else compute and persist once
-        stored_points = (exp.get("guideline_rationale") or [])
-        if stored_points:
-            _render_guideline_bullets(stored_points)
-        else:
-            g = generate_guideline_rationale(assessment_text=assess_new, plan_text=plan_new)
-            warn = (g.get("warning") or "").strip()
-            if warn:
-                st.info(warn)
-                try:
-                    st.page_link("pages/01_RAG_Corpus_Prep.py", label="Open RAG Prep")
-                except Exception:
-                    pass
-            points = g.get("rationale", []) or []
-            _render_guideline_bullets(points)
-            if points:
-                # Persist once so it survives navigation
-                store.update_case(selected, {"explainability": {"guideline_rationale": points}})
-    else:
-        # Hide until re-run; allow refresh which sets new baselines and persists bullets
-        if (not read_only) and (status == "submitted"):
-            if st.button("↻ Re-run Guidelines", key=f"rerun_{selected}"):
+    # Guard: only show guideline rationale if we have SOAP content
+    if not _soap_is_empty(soap):
+        if not stale and (normalize_text(assess_new + plan_new)):
+            # Use stored bullets if available; else compute once and store
+            stored_points = (exp.get("guideline_rationale") or [])
+            if stored_points:
+                _render_guideline_bullets(stored_points)
+            else:
                 g = generate_guideline_rationale(assessment_text=assess_new, plan_text=plan_new)
+                warn = (g.get("warning") or "").strip()
+                if warn:
+                    st.info(warn)
+                    try:
+                        st.page_link("pages/01_RAG_Corpus_Prep.py", label="Open RAG Prep")
+                    except Exception:
+                        pass
                 points = g.get("rationale", []) or []
-                store.update_case(selected, {
-                    "explainability": {
-                        "baseline": {"assessment_hash": text_hash(assess_new), "plan_hash": text_hash(plan_new)},
-                        "guideline_rationale": points,
-                    }
-                })
-                st.rerun()
-            else:
-                st.caption("Guideline rationale hidden after edits. Click **Re-run Guidelines** to refresh.")
+                _render_guideline_bullets(points)
+                if points:
+                    store.update_case(selected, {"explainability": {"guideline_rationale": points}})
+        else:
+            if (not read_only) and (status == "submitted"):
+                if st.button("↻ Re-run Guidelines", key=f"rerun_{selected}"):
+                    g = generate_guideline_rationale(assessment_text=assess_new, plan_text=plan_new)
+                    points = g.get("rationale", []) or []
+                    store.update_case(selected, {
+                        "explainability": {
+                            "baseline": {"assessment_hash": text_hash(assess_new), "plan_hash": text_hash(plan_new)},
+                            "guideline_rationale": points,
+                        }
+                    })
+                    st.rerun()
+                else:
+                    st.caption("Guideline rationale hidden after edits. Click **Re-run Guidelines** to refresh.")
 
     st.divider()
 
-    # ---- Specialist Review Metadata (3 controls) ----
+    # Specialist Review metadata (3 controls)
     st.subheader("Specialist Review")
     b = case.get("billing", {}) or {}
     minutes_val = int(b.get("minutes", 5))
@@ -470,7 +468,7 @@ with tab_spec:
     attested = bool(b.get("attested", False)) if read_only else False
     attested = st.checkbox("I attest that the consult note is complete.", value=attested, disabled=read_only, key=f"att_{selected}")
 
-    # Persist time/spoke as they change (do not persist attested until finalize)
+    # Persist time/spoke as they change (not attested until finalize)
     if (not read_only) and (status == "submitted"):
         if (minutes != minutes_val) or (spoke != bool(b.get("spoke", False))):
             store.update_case(selected, {"billing": {"minutes": minutes, "spoke": bool(spoke)}})
@@ -480,20 +478,19 @@ with tab_spec:
     finalize_btn = st.button("Finalize Consult", type="primary", disabled=not can_finalize, key=f"finalize_{selected}")
 
     if finalize_btn:
-        # Compose note with current SOAP + persisted/stored bullets (fallback to latest generation)
-        g = generate_guideline_rationale(assessment_text=assess_new, plan_text=plan_new)
-        points = (exp.get("guideline_rationale") or []) or (g.get("rationale", []) or [])
-        endnotes = g.get("endnotes", []) or []
+        # Use persisted guideline bullets; do NOT regenerate at finalize
+        points = (case.get("explainability", {}) or {}).get("guideline_rationale") or []
+        endnotes: List[Dict[str, Any]] = []
 
         soap_now = case.get("soap_draft", {}) or {}
         note_md = _note_markdown(
             case,
             build_referral_summary(case),
             {
-                "subjective": locals().get("subj_new", soap_now.get("subjective","")),
-                "objective":  locals().get("obj_new",  soap_now.get("objective","")),
-                "assessment": locals().get("assess_new", soap_now.get("assessment","")),
-                "plan":       locals().get("plan_new",   soap_now.get("plan","")),
+                "subjective": subj_new or soap_now.get("subjective",""),
+                "objective":  obj_new or soap_now.get("objective",""),
+                "assessment": assess_new or soap_now.get("assessment",""),
+                "plan":       plan_new or soap_now.get("plan",""),
             },
             points, endnotes=endnotes
         )
@@ -502,14 +499,15 @@ with tab_spec:
         note_path = config.make_export_path(selected, "Sample Consultation Report.md")
         Path(note_path).write_text(note_md, encoding="utf-8")
 
-        # Persist attestation now and mark completed
-        store.update_case(selected, {"billing": {"minutes": minutes, "spoke": bool(spoke), "attested": True}})
-        store.set_status(selected, "completed")
+        # Persist billing info and mark completed
+        store.update_case(selected, {
+            "billing": {"minutes": minutes, "spoke": bool(spoke), "attested": True},
+            "status": "completed",
+        })
 
         st.session_state["_show_exports"] = True
-        st.success("Finalized. Consult note exported; proceed to the **Documentation & Billing** tab to submit a claim.")
+        st.success("✅ Finalized. Consult note exported; proceed to the **Documentation & Billing** tab to submit a claim.")
         st.caption(f"Note: {note_path}")
-        # Optional: st.rerun() here would reset tab focus; we leave it so user can jump tabs manually.
 
 
 # ===== DOCUMENTATION & BILLING TAB =====
@@ -525,7 +523,7 @@ with tab_bill:
     if note_path and note_path.exists():
         with st.expander("📄 View EHR Consultation Note", expanded=True):
             md = note_path.read_text(encoding="utf-8")
-            st.markdown(md)  # render formatted markdown, not raw text
+            st.markdown(md)
         st.caption(f"File: {note_path.name}")
     else:
         st.warning("No consultation note found for this case. Finalize the consult to generate one.")
@@ -548,11 +546,9 @@ with tab_bill:
             st.write(f"{icon} **{s.code}** — {s.descriptor} — ${s.rate:.2f}")
             st.caption(s.why)
 
-    # --- ICD-10 diagnosis suggestions (new) ---
-    from src.billing import autosuggest_icd  # top of file already imports billing; safe to import here too
-    
+    # ICD-10 suggestions block (from Assessment + Plan)
     soap_now = case.get("soap_draft", {}) or {}
-    icd_suggestions = autosuggest_icd(
+    icd_suggestions = billing.autosuggest_icd(
         assessment_text=soap_now.get("assessment", ""),
         plan_text=soap_now.get("plan", "")
     )
@@ -561,7 +557,6 @@ with tab_bill:
         for s in icd_suggestions:
             st.write(f"✅ **{s.code}** — {s.description}  (confidence {int(s.confidence*100)}%)")
             st.caption(s.why)
-        # Allow user to choose 1–3 codes
         icd_options = [f"{s.code} — {s.description}" for s in icd_suggestions]
         picked_icd = st.multiselect(
             "Select diagnosis code(s) to include on the claim",
@@ -572,12 +567,16 @@ with tab_bill:
         picked_icd_codes = [o.split(" — ")[0] for o in picked_icd]
     else:
         st.info("No ICD-10 suggestions available for this case.")
-        picked_icd_codes = []
+        picked_icd_codes: List[str] = []
 
-    
     chosen_default = b.get("cpt_code") or (elig_codes[0] if elig_codes else None)
     if elig_codes:
-        chosen = st.selectbox("Choose CPT (eligible)", options=elig_codes, index=elig_codes.index(chosen_default) if (chosen_default in elig_codes) else 0, key=f"cpt_{selected}")
+        chosen = st.selectbox(
+            "Choose CPT (eligible)",
+            options=elig_codes,
+            index=elig_codes.index(chosen_default) if (chosen_default in elig_codes) else 0,
+            key=f"cpt_{selected}"
+        )
     else:
         st.warning("No eligible CPT for the current minutes/spoke combination.", icon="⚠️")
         chosen = None
@@ -591,10 +590,10 @@ with tab_bill:
         store.update_case(selected, {
             "billing": {
                 "cpt_code": chosen,
-                "icd_codes": picked_icd_codes,   # NEW: persist ICDs in case JSON
+                "icd_codes": picked_icd_codes,
             }
         })
-        
+
         # Build and save 837 JSON (includes ICD codes)
         pick = next((s for s in suggestions if s.code == chosen), None)
         rate = float(getattr(pick, "rate", 0.0)) if pick else 0.0
@@ -605,11 +604,11 @@ with tab_bill:
             minutes=minutes,
             spoke=bool(spoke),
             attested=True,
-            icd_codes=picked_icd_codes,          # NEW: include ICD list in claim
+            icd_codes=picked_icd_codes,
         )
         claim_path = config.make_export_path(selected, "Sample 837 Json.json")
         Path(claim_path).write_text(json.dumps(claim, ensure_ascii=False, indent=2), encoding="utf-8")
-        
+
         st.success("Claim submitted. 837 JSON generated (see preview below).")
         st.caption(f"Claim: {claim_path.name}")
 
@@ -637,3 +636,4 @@ with tab_bill:
 
 
 
+