import streamlit as st import json import re import os import traceback import logging from dotenv import load_dotenv # Configure logging logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO) # Import agent logic and message types try: from agent import ClinicalAgent, AgentState, check_red_flags from langchain_core.messages import HumanMessage, AIMessage, ToolMessage except ImportError as e: logger.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() 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" PAGE_LAYOUT = "wide" 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: logger.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") # Updated minimum height to 68px to satisfy Streamlit requirement exam_notes = st.text_area("Exam Notes", "Awake, alert...", height=68, 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: logger.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 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 "agent" not in st.session_state: st.session_state.agent = get_agent() # Sidebar intake with st.sidebar: patient_data = load_patient_intake() if st.button("Start/Update Consultation", key="sb_start"): 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() # Chat area st.header("💬 Clinical Consultation") for msg in st.session_state.messages: if isinstance(msg, HumanMessage): 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: 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: logger.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', {}) 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() elif isinstance(msg, ToolMessage): 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) # --- 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) st.session_state.messages = final_state.get('messages', []) st.session_state.summary = final_state.get('summary') except Exception as e: logger.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 only and not for clinical use.") if __name__ == "__main__": main()