Update app.py

app.py

@@ -1,189 +1,295 @@
-# app.py
 import streamlit as st
 import json
 import re
 import os
 import traceback
+import logging
 from dotenv import load_dotenv
 
+# Configure logging
+tlogging = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
 # Import agent logic and message types
 try:
-    from agent import ClinicalAgent, AgentState, check_red_flags # Import necessary components
+    from agent import ClinicalAgent, AgentState, check_red_flags
     from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
 except ImportError as e:
+    logging.exception("Failed to import from agent.py")
     st.error(f"Failed to import from agent.py: {e}. Make sure agent.py is in the same directory.")
     st.stop()
 
-
 # --- Environment Variable Loading & Validation ---
 load_dotenv()
-# Check keys required by agent.py are present before initializing the agent
-UMLS_API_KEY = os.environ.get("UMLS_API_KEY")
-GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
-TAVILY_API_KEY = os.environ.get("TAVILY_API_KEY")
-missing_keys = []
-if not UMLS_API_KEY: missing_keys.append("UMLS_API_KEY")
-if not GROQ_API_KEY: missing_keys.append("GROQ_API_KEY")
-if not TAVILY_API_KEY: missing_keys.append("TAVILY_API_KEY")
-if missing_keys:
-    st.error(f"Missing required API Key(s): {', '.join(missing_keys)}. Please set them in Hugging Face Space Secrets or environment variables.")
+required_keys = ["UMLS_API_KEY", "GROQ_API_KEY", "TAVILY_API_KEY"]
+missing = [key for key in required_keys if not os.getenv(key)]
+if missing:
+    st.error(f"Missing required API Key(s): {', '.join(missing)}. Please set them in environment variables.")
     st.stop()
 
 # --- App Configuration ---
 class ClinicalAppSettings:
     APP_TITLE = "SynapseAI (UMLS/FDA Integrated)"
     PAGE_LAYOUT = "wide"
-    MODEL_NAME_DISPLAY = "Llama3-70b (via Groq)" # Defined in agent.py
+    MODEL_NAME_DISPLAY = "Llama3-70b (via Groq)"
+
+# Cache the agent to avoid re-initialization on each rerun
+@st.cache_resource
+def get_agent():
+    try:
+        return ClinicalAgent()
+    except Exception as e:
+        logging.exception("Failed to initialize ClinicalAgent")
+        st.error(f"Failed to initialize Clinical Agent: {e}. Check API keys and dependencies.")
+        st.stop()
+
+# Sidebar patient intake helper
+def load_patient_intake():
+    st.header("📄 Patient Intake Form")
+    # Demographics
+    age = st.number_input("Age", min_value=0, max_value=120, value=55, key="sb_age")
+    sex = st.selectbox("Sex", ["Male", "Female", "Other"], key="sb_sex")
+
+    # HPI
+    chief_complaint = st.text_input("Chief Complaint", "Chest pain", key="sb_cc")
+    hpi_details = st.text_area("HPI Details", "55 y/o male...", height=100, key="sb_hpi")
+    symptoms = st.multiselect(
+        "Symptoms",
+        ["Nausea", "Diaphoresis", "SOB", "Dizziness", "Severe Headache", "Syncope", "Hemoptysis"],
+        default=["Nausea", "Diaphoresis"],
+        key="sb_sym"
+    )
+
+    # History
+    pmh = st.text_area("PMH", "HTN, HLD, DM2, History of MI", key="sb_pmh")
+    psh = st.text_area("PSH", "Appendectomy", key="sb_psh")
+
+    # Meds & Allergies
+    current_meds_str = st.text_area(
+        "Current Meds",
+        "Lisinopril 10mg daily\nMetformin 1000mg BID\nWarfarin 5mg daily",
+        key="sb_meds"
+    )
+    allergies_str = st.text_area("Allergies", "Penicillin (rash), Aspirin", key="sb_allergies")
+
+    # Social/Family
+    social_history = st.text_area("SH", "Smoker", key="sb_sh")
+    family_history = st.text_area("FHx", "Father MI", key="sb_fhx")
+
+    # Vitals & Exam
+    col1, col2 = st.columns(2)
+    with col1:
+        temp_c = st.number_input("Temp C", min_value=35.0, max_value=42.0, value=36.8, format="%.1f", key="sb_temp")
+        hr_bpm = st.number_input("HR", min_value=30, max_value=250, value=95, key="sb_hr")
+        rr_rpm = st.number_input("RR", min_value=5, max_value=50, value=18, key="sb_rr")
+    with col2:
+        bp_mmhg = st.text_input("BP", "155/90", key="sb_bp")
+        spo2_percent = st.number_input("SpO2", min_value=70, max_value=100, value=96, key="sb_spo2")
+        pain_scale = st.slider("Pain", min_value=0, max_value=10, value=8, key="sb_pain")
+    exam_notes = st.text_area("Exam Notes", "Awake, alert...", height=50, key="sb_exam")
+
+    # Process meds and allergies with comprehensions
+    current_meds_list = [m.strip() for m in current_meds_str.splitlines() if m.strip()]
+    current_med_names_only = [
+        m.group(1).lower()
+        for med in current_meds_list
+        if (m := re.match(r"^\s*([A-Za-z-]+)", med))
+    ]
+    allergies_list = [
+        (m.group(1).strip().lower() if (m := re.match(r"^\s*([A-Za-z\s/-]+)", a.strip())) else a.strip().lower())
+        for a in allergies_str.split(",")
+        if a.strip()
+    ]
+
+    # Parse blood pressure
+    bp_sys, bp_dia = None, None
+    if "/" in bp_mmhg:
+        try:
+            bp_sys, bp_dia = map(int, bp_mmhg.split("/"))
+        except ValueError:
+            logging.warning(f"Unable to parse BP '{bp_mmhg}'")
+
+    return {
+        "demographics": {"age": age, "sex": sex},
+        "hpi": {"chief_complaint": chief_complaint, "details": hpi_details, "symptoms": symptoms},
+        "pmh": {"conditions": pmh},
+        "psh": {"procedures": psh},
+        "medications": {"current": current_meds_list, "names_only": current_med_names_only},
+        "allergies": allergies_list,
+        "social_history": {"details": social_history},
+        "family_history": {"details": family_history},
+        "vitals": {
+            "temp_c": temp_c,
+            "hr_bpm": hr_bpm,
+            "bp_mmhg": bp_mmhg,
+            "bp_sys": bp_sys,
+            "bp_dia": bp_dia,
+            "rr_rpm": rr_rpm,
+            "spo2_percent": spo2_percent,
+            "pain_scale": pain_scale
+        },
+        "exam_findings": {"notes": exam_notes},
+    }
 
+# Main application
 
-# --- Streamlit UI ---
 def main():
     st.set_page_config(page_title=ClinicalAppSettings.APP_TITLE, layout=ClinicalAppSettings.PAGE_LAYOUT)
     st.title(f"🩺 {ClinicalAppSettings.APP_TITLE}")
     st.caption(f"Interactive Assistant | LangGraph/Groq/Tavily/UMLS/OpenFDA | Model: {ClinicalAppSettings.MODEL_NAME_DISPLAY}")
 
     # Initialize session state
-    if "messages" not in st.session_state: st.session_state.messages = []
-    if "patient_data" not in st.session_state: st.session_state.patient_data = None
-    if "summary" not in st.session_state: st.session_state.summary = None
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+    if "patient_data" not in st.session_state:
+        st.session_state.patient_data = None
+    if "summary" not in st.session_state:
+        st.session_state.summary = None
     if "agent" not in st.session_state:
-        try:
-            st.session_state.agent = ClinicalAgent()
-            print("ClinicalAgent successfully initialized in Streamlit session state.")
-        except Exception as e:
-            st.error(f"Failed to initialize Clinical Agent: {e}. Check API keys and dependencies.")
-            print(f"ERROR Initializing ClinicalAgent: {e}"); traceback.print_exc(); st.stop()
+        st.session_state.agent = get_agent()
 
-    # --- Patient Data Input Sidebar ---
+    # Sidebar intake
     with st.sidebar:
-        st.header("📄 Patient Intake Form")
-        # Input fields... (Assume full fields as before)
-        st.subheader("Demographics"); age = st.number_input("Age", 0, 120, 55, key="sb_age"); sex = st.selectbox("Sex", ["Male", "Female", "Other"], key="sb_sex")
-        st.subheader("HPI"); chief_complaint = st.text_input("Chief Complaint", "Chest pain", key="sb_cc"); hpi_details = st.text_area("HPI Details", "55 y/o male...", height=100, key="sb_hpi"); symptoms = st.multiselect("Symptoms", ["Nausea", "Diaphoresis", "SOB", "Dizziness", "Severe Headache", "Syncope", "Hemoptysis"], default=["Nausea", "Diaphoresis"], key="sb_sym")
-        st.subheader("History"); pmh = st.text_area("PMH", "HTN, HLD, DM2, History of MI", key="sb_pmh"); psh = st.text_area("PSH", "Appendectomy", key="sb_psh")
-        st.subheader("Meds & Allergies"); current_meds_str = st.text_area("Current Meds", "Lisinopril 10mg daily\nMetformin 1000mg BID\nWarfarin 5mg daily", key="sb_meds"); allergies_str = st.text_area("Allergies", "Penicillin (rash), Aspirin", key="sb_allergies")
-        st.subheader("Social/Family"); social_history = st.text_area("SH", "Smoker", key="sb_sh"); family_history = st.text_area("FHx", "Father MI", key="sb_fhx")
-        st.subheader("Vitals & Exam"); col1, col2 = st.columns(2);
-        with col1: temp_c = st.number_input("Temp C", 35.0, 42.0, 36.8, format="%.1f", key="sb_temp"); hr_bpm = st.number_input("HR", 30, 250, 95, key="sb_hr"); rr_rpm = st.number_input("RR", 5, 50, 18, key="sb_rr")
-        with col2: bp_mmhg = st.text_input("BP", "155/90", key="sb_bp"); spo2_percent = st.number_input("SpO2", 70, 100, 96, key="sb_spo2"); pain_scale = st.slider("Pain", 0, 10, 8, key="sb_pain")
-        exam_notes = st.text_area("Exam Notes", "Awake, alert...", height=50, key="sb_exam")
-
+        patient_data = load_patient_intake()
         if st.button("Start/Update Consultation", key="sb_start"):
-            # Compile data...
-            current_meds_list = [med.strip() for med in current_meds_str.split('\n') if med.strip()]
-            current_med_names_only = [];
-            for med in current_meds_list: match = re.match(r"^\s*([a-zA-Z\-]+)", med);
-            if match: current_med_names_only.append(match.group(1).lower())
-            allergies_list = []
-            for a in allergies_str.split(','): cleaned_allergy = a.strip();
-            if cleaned_allergy: match = re.match(r"^\s*([a-zA-Z\-\s/]+)(?:\s*\(.*\))?", cleaned_allergy); name_part = match.group(1).strip().lower() if match else cleaned_allergy.lower(); allergies_list.append(name_part)
-            # Update patient data in session state
-            st.session_state.patient_data = { "demographics": {"age": age, "sex": sex}, "hpi": {"chief_complaint": chief_complaint, "details": hpi_details, "symptoms": symptoms}, "pmh": {"conditions": pmh}, "psh": {"procedures": psh}, "medications": {"current": current_meds_list, "names_only": current_med_names_only}, "allergies": allergies_list, "social_history": {"details": social_history}, "family_history": {"details": family_history}, "vitals": { "temp_c": temp_c, "hr_bpm": hr_bpm, "bp_mmhg": bp_mmhg, "rr_rpm": rr_rpm, "spo2_percent": spo2_percent, "pain_scale": pain_scale}, "exam_findings": {"notes": exam_notes} }
-            # Call check_red_flags from agent module
-            red_flags = check_red_flags(st.session_state.patient_data); st.sidebar.markdown("---");
-            if red_flags: st.sidebar.warning("**Initial Red Flags:**"); [st.sidebar.warning(f"- {flag.replace('Red Flag: ','')}") for flag in red_flags]
-            else: st.sidebar.success("No immediate red flags.")
-            # Reset conversation and summary on new intake
-            initial_prompt = "Initiate consultation. Review patient data and begin analysis."
-            st.session_state.messages = [HumanMessage(content=initial_prompt)]
-            st.session_state.summary = None # Reset summary
+            st.session_state.patient_data = patient_data
+            red_flags = check_red_flags(patient_data)
+            st.sidebar.markdown("---")
+            if red_flags:
+                st.sidebar.warning("**Initial Red Flags:**")
+                for flag in red_flags:
+                    st.sidebar.warning(f"- {flag.replace('Red Flag: ', '')}")
+            else:
+                st.sidebar.success("No immediate red flags.")
+            st.session_state.messages = [HumanMessage(content="Initiate consultation. Review patient data and begin analysis.")]
+            st.session_state.summary = None
             st.success("Patient data loaded/updated.")
             st.rerun()
 
-    # --- Main Chat Interface Area ---
+    # Chat area
     st.header("💬 Clinical Consultation")
-    # Display loop
     for msg in st.session_state.messages:
         if isinstance(msg, HumanMessage):
-            with st.chat_message("user"): st.markdown(msg.content)
+            with st.chat_message("user"):
+                st.markdown(msg.content)
         elif isinstance(msg, AIMessage):
             with st.chat_message("assistant"):
-                ai_content = msg.content; structured_output = None
-                try: # JSON Parsing logic...
-                    json_match = re.search(r"```json\s*(\{.*?\})\s*```", ai_content, re.DOTALL | re.IGNORECASE)
-                    if json_match: json_str = json_match.group(1); prefix = ai_content[:json_match.start()].strip(); suffix = ai_content[json_match.end():].strip();
-                    if prefix: st.markdown(prefix); structured_output = json.loads(json_str);
-                    if suffix: st.markdown(suffix)
-                    elif ai_content.strip().startswith("{") and ai_content.strip().endswith("}"): structured_output = json.loads(ai_content); ai_content = ""
-                    else: st.markdown(ai_content)
-                except Exception as e: st.markdown(ai_content); print(f"Error parsing/displaying AI JSON: {e}")
-
-                if structured_output and isinstance(structured_output, dict): # Structured JSON display logic...
-                    st.divider(); st.subheader("📊 AI Analysis & Recommendations")
-                    cols = st.columns(2);
-                    with cols[0]: # Assessment, DDx, Risk
-                        st.markdown("**Assessment:**"); st.markdown(f"> {structured_output.get('assessment', 'N/A')}")
-                        st.markdown("**Differential Diagnosis:**"); ddx = structured_output.get('differential_diagnosis', []);
-                        if ddx: [st.expander(f"{'🥇🥈🥉'[('High','Medium','Low').index(item.get('likelihood','Low')[0])] if item.get('likelihood','?')[0] in 'HML' else '?'} {item.get('diagnosis', 'Unknown')} ({item.get('likelihood','?')})").write(f"**Rationale:** {item.get('rationale', 'N/A')}") for item in ddx]
-                        else: st.info("No DDx provided.")
-
-                        # Risk Assessment Display (CORRECTED - Separate lines)
-                        st.markdown(f"**Risk Assessment:**")
+                ai_content = msg.content
+                structured_output = None
+                try:
+                    match = re.search(r"```json\s*(\{.*?\})\s*```", ai_content, re.DOTALL | re.IGNORECASE)
+                    if match:
+                        payload = match.group(1)
+                        structured_output = json.loads(payload)
+                        prefix = ai_content[:match.start()].strip()
+                        suffix = ai_content[match.end():].strip()
+                        if prefix:
+                            st.markdown(prefix)
+                        if suffix:
+                            st.markdown(suffix)
+                    else:
+                        st.markdown(ai_content)
+                except (AttributeError, json.JSONDecodeError) as e:
+                    logging.warning(f"JSON parse error: {e}")
+                    st.markdown(ai_content)
+
+                if structured_output and isinstance(structured_output, dict):
+                    st.divider()
+                    # Display structured JSON sections
+                    cols = st.columns(2)
+                    with cols[0]:
+                        st.markdown("**Assessment:**")
+                        st.markdown(f"> {structured_output.get('assessment', 'N/A')}" )
+                        st.markdown("**Differential Diagnosis:**")
+                        ddx = structured_output.get('differential_diagnosis', [])
+                        if ddx:
+                            for item in ddx:
+                                likelihood = item.get('likelihood', 'Low')
+                                icon = '🥇' if likelihood == 'High' else ('🥈' if likelihood == 'Medium' else '🥉')
+                                with st.expander(f"{icon} {item.get('diagnosis', 'Unknown')} ({likelihood})"):
+                                    st.write(f"**Rationale:** {item.get('rationale', 'N/A')}")
+                        else:
+                            st.info("No DDx provided.")
+
+                        st.markdown("**Risk Assessment:**")
                         risk = structured_output.get('risk_assessment', {})
-                        flags = risk.get('identified_red_flags', [])
-                        concerns = risk.get("immediate_concerns", [])
-                        comps = risk.get("potential_complications", [])
-
-                        if flags:
-                            st.warning(f"**Flags:** {', '.join(flags)}")
-                        if concerns:
-                            st.warning(f"**Concerns:** {', '.join(concerns)}")
-                        if comps:
-                            st.info(f"**Potential Complications:** {', '.join(comps)}")
-                        # Add a message if no risks were highlighted by the AI assessment
-                        if not flags and not concerns and not comps:
-                            st.success("No specific risks highlighted in this AI assessment.")
-
-                    with cols[1]: # Plan
-                         st.markdown("**Recommended Plan:**"); plan = structured_output.get('recommended_plan', {});
-                         for section in ["investigations","therapeutics","consultations","patient_education"]: st.markdown(f"_{section.replace('_',' ').capitalize()}:_"); items = plan.get(section); [st.markdown(f"- {item}") for item in items] if items and isinstance(items, list) else (st.markdown(f"- {items}") if items else st.markdown("_None_")); st.markdown("")
-
-                    # Rationale & Interaction Summary
-                    st.markdown("**Rationale & Guideline Check:**"); st.markdown(f"> {structured_output.get('rationale_summary', 'N/A')}")
-                    interaction_summary = structured_output.get("interaction_check_summary", "");
-                    if interaction_summary: st.markdown("**Interaction Check Summary:**"); st.markdown(f"> {interaction_summary}");
+                        for key, style in [('identified_red_flags', st.warning), ('immediate_concerns', st.warning), ('potential_complications', st.info)]:
+                            items = risk.get(key, [])
+                            if items:
+                                style(f"**{key.replace('_', ' ').capitalize()}:** {', '.join(items)}")
+                        if not any(risk.get(k) for k in ['identified_red_flags', 'immediate_concerns', 'potential_complications']):
+                            st.success("No specific risks highlighted.")
+
+                    with cols[1]:
+                        st.markdown("**Recommended Plan:**")
+                        plan = structured_output.get('recommended_plan', {})
+                        for section in ["investigations","therapeutics","consultations","patient_education"]:
+                            st.markdown(f"_{section.replace('_',' ').capitalize()}:_")
+                            items = plan.get(section)
+                            if isinstance(items, list):
+                                for it in items:
+                                    st.markdown(f"- {it}")
+                            elif items:
+                                st.markdown(f"- {items}")
+                            else:
+                                st.markdown("_None_")
+
+                    st.markdown("**Rationale & Guideline Check:**")
+                    st.markdown(f"> {structured_output.get('rationale_summary', 'N/A')}" )
+                    if interaction := structured_output.get('interaction_check_summary'):
+                        st.markdown("**Interaction Check Summary:**")
+                        st.markdown(f"> {interaction}")
                     st.divider()
 
-                # Tool Call Display
-                if getattr(msg, 'tool_calls', None):
-                     with st.expander("🛠️ AI requested actions", expanded=False):
-                         if msg.tool_calls:
-                            for tc in msg.tool_calls:
-                                try: st.code(f"Action: {tc.get('name', 'Unknown Tool')}\nArgs: {json.dumps(tc.get('args', {}), indent=2)}", language="json")
-                                except Exception as display_e: st.error(f"Could not display tool call args: {display_e}", icon="⚠️"); st.code(f"Action: {tc.get('name', 'Unknown Tool')}\nRaw Args: {tc.get('args')}")
-                         else: st.caption("_No actions requested._")
         elif isinstance(msg, ToolMessage):
-            tool_name_display = getattr(msg, 'name', 'tool_execution')
-            with st.chat_message(tool_name_display, avatar="🛠️"):
-                try: # Tool message display logic...
-                    tool_data = json.loads(msg.content); status = tool_data.get("status", "info"); message = tool_data.get("message", msg.content); details = tool_data.get("details"); warnings = tool_data.get("warnings");
-                    if tool_name_display == "flag_risk" and status == "flagged": st.error(f"🚨 **RISK FLAGGED:** {message}", icon="🚨") # Show flag in UI too
-                    elif status == "success" or status == "clear": st.success(f"{message}", icon="✅")
-                    elif status == "warning": st.warning(f"{message}", icon="⚠️");
-                    if warnings and isinstance(warnings, list): st.caption("Details:"); [st.caption(f"- {warn}") for warn in warnings]
-                    else: st.error(f"{message}", icon="❌") # Assume error if not known status
-                    if details: st.caption(f"Details: {details}")
-                except json.JSONDecodeError: st.info(f"{msg.content}") # Display raw if not JSON
-                except Exception as e: st.error(f"Error displaying tool message: {e}", icon="❌"); st.caption(f"Raw content: {msg.content}")
-
-    # --- Chat Input Logic ---
-    if prompt := st.chat_input("Your message or follow-up query..."):
-        if not st.session_state.patient_data: st.warning("Please load patient data first."); st.stop()
-        if 'agent' not in st.session_state or not st.session_state.agent: st.error("Agent not initialized. Check logs."); st.stop()
+            tool_name = getattr(msg, 'name', 'tool_execution')
+            with st.chat_message(tool_name, avatar="🛠️"):
+                try:
+                    data = json.loads(msg.content)
+                    status = data.get('status', 'info')
+                    message = data.get('message', msg.content)
+                    if tool_name == "flag_risk" and status == "flagged":
+                        st.error(f"🚨 **RISK FLAGGED:** {message}")
+                    elif status in ("success", "clear"):
+                        st.success(message)
+                    elif status == "warning":
+                        st.warning(message)
+                    else:
+                        st.error(message)
+                    if details := data.get('details'):
+                        st.caption(f"Details: {details}")
+                except json.JSONDecodeError:
+                    st.info(msg.content)
 
-        user_message = HumanMessage(content=prompt); st.session_state.messages.append(user_message)
-        with st.chat_message("user"): st.markdown(prompt)
-        current_state_dict = {"messages": st.session_state.messages, "patient_data": st.session_state.patient_data, "summary": st.session_state.get("summary"), "interaction_warnings": None}
+    # --- Chat Input ---
+    if prompt := st.chat_input("Your message or follow-up query..."):
+        if not st.session_state.patient_data:
+            st.warning("Please load patient data first.")
+            st.stop()
+        user_msg = HumanMessage(content=prompt)
+        st.session_state.messages.append(user_msg)
+        with st.chat_message("user"):
+            st.markdown(prompt)
+        current_state = {
+            "messages": st.session_state.messages,
+            "patient_data": st.session_state.patient_data,
+            "summary": st.session_state.summary,
+            "interaction_warnings": None
+        }
         with st.spinner("SynapseAI is processing..."):
             try:
-                final_state = st.session_state.agent.invoke_turn(current_state_dict)
+                final_state = st.session_state.agent.invoke_turn(current_state)
                 st.session_state.messages = final_state.get('messages', [])
                 st.session_state.summary = final_state.get('summary')
-            except Exception as e: print(f"CRITICAL ERROR: {e}"); traceback.print_exc(); st.error(f"Error: {e}"); st.session_state.messages.append(AIMessage(content=f"Error processing request: {e}"))
+            except Exception as e:
+                logging.exception("Error during agent.invoke_turn")
+                st.error(f"Error: {e}")
+                st.session_state.messages.append(AIMessage(content=f"Error processing request: {e}"))
         st.rerun()
 
     # Disclaimer
-    st.markdown("---"); st.warning("**Disclaimer:** SynapseAI is for demonstration...")
+    st.markdown("---")
+    st.warning("**Disclaimer:** SynapseAI is for demonstration only and not for clinical use.")
 
 if __name__ == "__main__":
     main()