app.py

# app.py
import streamlit as st
import json
import re
import os
import traceback
from dotenv import load_dotenv

# Import agent logic and message types from agent.py
try:
    from agent import ClinicalAgent, AgentState, check_red_flags
    from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
except ImportError as e:
    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.")
    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

# --- 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

    # Initialize the agent instance only once
    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()

    # --- Patient Data Input Sidebar ---
    with st.sidebar:
        st.header("📄 Patient Intake Form")
        # Input fields... (Using shorter versions for brevity, assume full fields are here)
        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")  # Added Warfarin/Aspirin for testing

        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")

        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:**")
                for flag in red_flags:
                    st.sidebar.warning(f"- {flag.replace('Red Flag: ', '')}")
            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.success("Patient data loaded/updated.")
            # Rerun might be needed if the main area should clear or update based on new data
            st.rerun()

    # --- Main Chat Interface 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)
        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)  # Display non-JSON 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]:
                        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')
                                if likelihood and likelihood[0] in 'HML':
                                    medal = "🥇" if likelihood[0] == 'H' else "🥈" if likelihood[0] == 'M' else "🥉"
                                else:
                                    medal = "?"
                                expander_title = f"{medal} {item.get('diagnosis', 'Unknown')} ({likelihood})"
                                with st.expander(expander_title):
                                    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)}")
                        if not flags and not concerns:
                            st.success("No major 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 items and isinstance(items, list):
                                for item in items:
                                    st.markdown(f"- {item}")
                            elif items:
                                st.markdown(f"- {items}")
                            else:
                                st.markdown("_None_")
                            st.markdown("")
                        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}")
                        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")
                    # Display flagged risks immediately if the tool signals it
                    if tool_name_display == "flag_risk" and status == "flagged":
                        st.error(f"🚨 **RISK FLAGGED:** {message}", icon="🚨")
                    elif status in ["success", "clear"]:
                        st.success(f"{message}", icon="✅")
                    elif status == "warning":
                        st.warning(f"{message}", icon="⚠️")
                    else:
                        st.error(f"{message}", icon="❌")
                    if warnings and isinstance(warnings, list):
                        st.caption("Details:")
                        for warn in warnings:
                            st.caption(f"- {warn}")
                    if details:
                        st.caption(f"Details: {details}")
                except json.JSONDecodeError:
                    st.info(f"{msg.content}")
                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()

        # Append user message and display immediately
        user_message = HumanMessage(content=prompt)
        st.session_state.messages.append(user_message)
        with st.chat_message("user"):
            st.markdown(prompt)

        # Prepare state for the agent
        current_state_dict = {
            "messages": st.session_state.messages,
            "patient_data": st.session_state.patient_data,
            "summary": st.session_state.get("summary"),
            "interaction_warnings": None  # Start clean
        }

        # Invoke the agent's graph for one turn
        with st.spinner("SynapseAI is processing..."):
            try:
                # Call the agent instance's method
                final_state = st.session_state.agent.invoke_turn(current_state_dict)

                # Update Streamlit session state from the returned agent 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 during agent invocation: {type(e).__name__} - {e}")
                traceback.print_exc()
                st.error(f"An error occurred during processing: {e}", icon="❌")
                # Append error to messages for user visibility
                st.session_state.messages.append(AIMessage(content=f"Error during processing: {e}"))

        # Rerun Streamlit script to update the chat display
        st.rerun()

    # Disclaimer
    st.markdown("---")
    st.warning("**Disclaimer:** SynapseAI is for demonstration...")

if __name__ == "__main__":
    main()