File size: 8,550 Bytes
fca6e17
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
import streamlit as st
import uuid
from langchain_core.messages import HumanMessage, AIMessage

# --- Backend Imports ---
# Ensure these files are in the same directory as your Streamlit app
from main_back import app
from quick_help import quick_help_app

# --- Page Configuration ---
st.set_page_config(page_title="Gen AI Mental Health Assistant", layout="wide")

st.title("Mental Health AI Assistant")
st.markdown("A safe space to find support and track your well-being.")

# --- Questionnaire Data ---
# Based on the DSM-5 PDF provided earlier
questionnaire_data = [
    {"id": "1", "text": "Little interest or pleasure in doing things?"},
    {"id": "2", "text": "Feeling down, depressed, or hopeless?"},
    {"id": "3", "text": "Feeling more irritated, grouchy, or angry than usual?"},
    {"id": "4", "text": "Sleeping less than usual, but still have a lot of energy?"},
    {"id": "5", "text": "Starting lots more projects than usual or doing more risky things than usual?"},
    {"id": "6", "text": "Feeling nervous, anxious, frightened, worried, or on edge?"},
    {"id": "7", "text": "Feeling panic or being frightened?"},
    {"id": "8", "text": "Avoiding situations that make you anxious?"},
    {"id": "9", "text": "Unexplained aches and pains (e.g., head, back, joints, abdomen, legs)?"},
    {"id": "10", "text": "Feeling that your illnesses are not being taken seriously enough?"},
    {"id": "11", "text": "Thoughts of actually hurting yourself?"},
    {"id": "12", "text": "Hearing things other people couldn't hear, such as voices even when no one was around?"},
    {"id": "13", "text": "Feeling that someone could hear your thoughts, or that you could hear what another person was thinking?"},
    {"id": "14", "text": "Problems with sleep that affected your sleep quality over all?"},
    {"id": "15", "text": "Problems with memory (e.g., learning new information) or with location (e.g., finding your way home)?"},
    {"id": "16", "text": "Unpleasant thoughts, urges, or images that repeatedly enter your mind?"},
    {"id": "17", "text": "Feeling driven to perform certain behaviors or mental acts over and over again?"},
    {"id": "18", "text": "Feeling detached or distant from yourself, your body, your physical surroundings, or your memories?"},
    {"id": "19", "text": "Not knowing who you really are or what you want out of life?"},
    {"id": "20", "text": "Not feeling close to other people or enjoying your relationships with them?"},
    {"id": "21", "text": "Drinking at least 4 drinks of any kind of alcohol in a single day?"},
    {"id": "22", "text": "Smoking any cigarettes, a cigar, or pipe, or using snuff or chewing tobacco?"},
    {"id": "23", "text": "Using any medicines ON YOUR OWN, that is, without a doctor's prescription, in greater amounts or longer than prescribed?"},
]
response_options = ["None (0)", "Slight (1)", "Mild (2)", "Moderate (3)", "Severe (4)"]


# --- App State Initialization ---
def initialize_session_state():
    # General app mode
    if 'app_mode' not in st.session_state:
        st.session_state.app_mode = "Quick Help"

    # Quick Help state
    if 'quick_help_history' not in st.session_state:
        st.session_state.quick_help_history = []
    
    # Tracking Health state
    if 'tracking_stage' not in st.session_state:
        st.session_state.tracking_stage = "questionnaire" # questionnaire -> chat
    if 'tracking_history' not in st.session_state:
        st.session_state.tracking_history = []
    if 'tracking_thread_id' not in st.session_state:
        st.session_state.tracking_thread_id = None

initialize_session_state()

# --- Sidebar for Navigation ---
with st.sidebar:
    st.header("Navigation")
    st.session_state.app_mode = st.radio(
        "Choose a feature:",
        ("Quick Help", "Track Your Health"),
        key="app_mode_selector"
    )
    st.info("Your conversations are private. We do not store personally identifiable information.")

# --- Main App Logic ---

# --- Quick Help Feature ---
if st.session_state.app_mode == "Quick Help":
    st.header("Quick Help Chat")
    st.markdown("Get immediate, supportive advice. How are you feeling right now?")

    # Display chat history
    for message in st.session_state.quick_help_history:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # Handle user input
    if user_input := st.chat_input("Share your thoughts..."):
        st.session_state.quick_help_history.append({"role": "user", "content": user_input})
        with st.chat_message("user"):
            st.markdown(user_input)

        with st.chat_message("assistant"):
            # The input for the non-persistent app is the full history each time
            history_for_input = [
                HumanMessage(content=msg["content"]) if msg["role"] == "user" else AIMessage(content=msg["content"])
                for msg in st.session_state.quick_help_history
            ]
            
            ai_response = quick_help_app.invoke({'messages': history_for_input}).get("messages",["Could not generate response"])[-1].content
            st.write(ai_response)

        st.session_state.quick_help_history.append({"role": "assistant", "content": ai_response})

# --- Track Your Health Feature ---
elif st.session_state.app_mode == "Track Your Health":
    
    # --- Stage 1: Questionnaire ---
    if st.session_state.tracking_stage == "questionnaire":
        st.header("Health & Well-being Questionnaire")
        st.markdown("Please answer the following questions based on your feelings over the **last two weeks**.")

        with st.form("health_questionnaire"):
            responses = {}
            for q in questionnaire_data:
                # Get the integer value from the selection
                response_str = st.radio(q["text"], options=response_options, key=q["id"], horizontal=True)
                responses[q["id"]] = response_options.index(response_str)

            submitted = st.form_submit_button("Analyze & Create My Plan")
            
            if submitted:
                # Generate a unique thread ID for this user's persistent session
                st.session_state.tracking_thread_id = f"user_{uuid.uuid4()}"
                config = {"configurable": {"thread_id": st.session_state.tracking_thread_id}}
                
                # Call the main backend app with the questionnaire responses
                initial_input = {"questionnaire_responses": responses}
                
                with st.spinner("Analyzing your responses and generating a personalized plan..."):
                    # Use .invoke() for the first call as we want the full plan at once
                    result = app.invoke(initial_input, config=config)
                    initial_plan = result.get('messages', ["Could not generate a plan."])[-1].content
                    
                # Store the initial plan and switch to chat mode
                st.session_state.tracking_history = [{"role": "assistant", "content": initial_plan}]
                st.session_state.tracking_stage = "chat"
                st.rerun()

    # --- Stage 2: Chat with the Plan ---
    elif st.session_state.tracking_stage == "chat":
        st.header("Your Personalized Plan & Chat")
        st.markdown("Here is an initial plan based on your responses. You can ask questions about it or request different exercises.")
        
        config = {"configurable": {"thread_id": st.session_state.tracking_thread_id}}

        # Display chat history
        for message in st.session_state.tracking_history:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])

        # Handle user input for follow-up questions
        if user_input := st.chat_input("Ask a question about your plan..."):
            st.session_state.tracking_history.append({"role": "user", "content": user_input})
            with st.chat_message("user"):
                st.markdown(user_input)

            with st.chat_message("assistant"):
                # For follow-ups, we only need to send the new message.
                # The checkpointer on the backend handles loading the history.
                follow_up_input = {"messages": [HumanMessage(content=user_input)]}
                result = app.invoke(follow_up_input, config=config)
                ai_response = result.get('messages', ["Could not generate a plan."])[-1].content

            st.session_state.tracking_history.append({"role": "assistant", "content": ai_response})