Spaces:
Running
Running
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})
|