agent.py

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

# Import agent logic and message types
from agent import ClinicalAgent, AgentState, check_red_flags # Import necessary components
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage

# --- 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 is now primarily defined in agent.py, but can keep for display
    MODEL_NAME_DISPLAY = "Llama3-70b (via Groq)"


# --- 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
    # Summary state for future memory enhancement
    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... (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")

        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.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 - Uses messages from st.session_state
    for msg in st.session_state.messages: # Removed enumerate and key
        if isinstance(msg, HumanMessage):
            with st.chat_message("user"): st.markdown(msg.content)
        elif isinstance(msg, AIMessage):
            with st.chat_message("assistant"):
                # ... (Keep the detailed AI message display logic, including JSON parsing) ...
                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")
                    # ... (Keep detailed JSON display logic for assessment, ddx, plan, etc.) ...
                    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: [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."); 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); [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("")
                    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="🛠️"):
                # ... (Keep ToolMessage display logic) ...
                try:
                    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 status == "success" or status == "clear" or status == "flagged": st.success(f"{message}", icon="✅" if status != "flagged" else "🚨")
                    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 not st.session_state.agent: st.error("Agent not initialized. Check logs."); st.stop() # Add check for agent

        # 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, including existing messages and patient data
        current_state_dict = {
            "messages": st.session_state.messages,
            "patient_data": st.session_state.patient_data,
            "summary": st.session_state.get("summary"), # Include summary if implemented
            "interaction_warnings": None # Always start turn with no pending warnings
        }

        # Invoke the agent's graph for one turn
        with st.spinner("SynapseAI is thinking..."):
            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') # Update summary if implemented

            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="❌")
                # Optionally append an error message to the chat display
                # st.session_state.messages.append(AIMessage(content=f"Error processing request: {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()