app.py

import streamlit as st
from doctor_agent import DoctorAgent

st.session_state.setdefault("conversation", [])
st.session_state.setdefault("doctor_agent", DoctorAgent())
st.session_state.setdefault("assessment_complete", False)
st.session_state.setdefault("use_rag", True)
st.session_state.setdefault("pending_input", "")
st.session_state.setdefault("waiting_for_doctor", False)

st.set_page_config(
    page_title="Doctor Assist - AI Medical Interview",
    page_icon="🩺",
    layout="wide"
)

st.title("🩺 Doctor Assist - AI Medical Interview")
st.markdown("""
This AI acts as a Senior Doctor conducting dynamic interviews with ASHA workers.
Start by entering initial symptoms and the doctor will ask follow-up questions.
""")

with st.sidebar:
    st.header("⚙️ Settings")
    
    st.session_state.use_rag = st.checkbox(
        "Use Medical Knowledge Base (RAG)", 
        value=st.session_state.use_rag,
        help="Enable to use retrieved medical knowledge for better accuracy"
    )
    
    st.divider()
    
    if st.button("🔄 Start New Consultation", type="primary"):
        st.session_state.conversation = []
        st.session_state.doctor_agent.reset()
        st.session_state.assessment_complete = False
        st.session_state.pending_input = ""
        st.session_state.waiting_for_doctor = False
        st.experimental_rerun()
    
    st.divider()
    
    st.markdown("""
    ### How it works:
    1. **Describe symptoms** - Enter patient's symptoms
    2. **Click "Send to Doctor"** - Get doctor's response
    3. **Answer questions** - Respond to follow-ups
    4. **Get assessment** - Receive diagnosis
    
    ### Tips:
    - Be specific about symptoms
    - Include duration and severity
    - Wait for doctor's questions before replying
    """)

st.header("💬 Consultation")

for msg in st.session_state.conversation:
    role = "user" if msg["role"] == "user" else "assistant"
    with st.chat_message(role):
        st.markdown(msg["content"])

if not st.session_state.assessment_complete:
    user_input = st.text_area(
        "Enter your message:",
        value=st.session_state.pending_input,
        height=100,
        placeholder="Describe the patient's symptoms or answer the doctor's questions..."
    )
    
    send_button = st.button("📤 Send to Doctor", type="primary", use_container_width=True)

    if send_button:
        if not user_input.strip():
            st.warning("⚠️ Please enter a message before sending.")
        else:
            st.session_state.waiting_for_doctor = False

            st.session_state.conversation.append({"role": "user", "content": user_input})
            
            with st.chat_message("user"):
                st.markdown(user_input)
            
            if not st.session_state.waiting_for_doctor:
                with st.chat_message("assistant"):
                    with st.spinner("Doctor is analyzing..."):
                        response = st.session_state.doctor_agent.analyze_and_respond(
                            user_input,
                            use_rag=st.session_state.use_rag
                        )
                        st.markdown(response)
                
                st.session_state.conversation.append({"role": "assistant", "content": response})
                
                if "ASSESSMENT:" in response.upper() or "DIAGNOSIS:" in response.upper():
                    st.session_state.assessment_complete = True
                    st.success("✅ Assessment complete! Start a new consultation to continue.")
                else:
                    
                    st.session_state.waiting_for_doctor = True
            
            st.session_state.pending_input = ""

else:
    st.info("✅ Assessment complete. Click 'Start New Consultation' in the sidebar to begin a new case.")