|
| 1 | +import streamlit as st |
| 2 | +from autogen import (SwarmAgent, SwarmResult, initiate_swarm_chat, OpenAIWrapper,AFTER_WORK,UPDATE_SYSTEM_MESSAGE) |
| 3 | +import os |
| 4 | + |
| 5 | +os.environ["AUTOGEN_USE_DOCKER"] = "0" |
| 6 | + |
| 7 | +if 'output' not in st.session_state: |
| 8 | + st.session_state.output = { |
| 9 | + 'assessment': '', |
| 10 | + 'action': '', |
| 11 | + 'followup': '' |
| 12 | + } |
| 13 | + |
| 14 | +st.sidebar.title("OpenAI API Key") |
| 15 | +api_key = st.sidebar.text_input("Enter your OpenAI API Key", type="password") |
| 16 | + |
| 17 | +st.sidebar.warning(""" |
| 18 | +## ⚠️ Important Notice |
| 19 | +
|
| 20 | +This application is a supportive tool and does not replace professional mental health care. If you're experiencing thoughts of self-harm or severe crisis: |
| 21 | +
|
| 22 | +- Call National Crisis Hotline: 988 |
| 23 | +- Call Emergency Services: 911 |
| 24 | +- Seek immediate professional help |
| 25 | +""") |
| 26 | + |
| 27 | +st.title("🧠 Mental Wellbeing Agent") |
| 28 | + |
| 29 | +st.info(""" |
| 30 | +**Meet Your Mental Wellbeing Agent Team:** |
| 31 | +
|
| 32 | +🧠 **Assessment Agent** - Analyzes your situation and emotional needs |
| 33 | +🎯 **Action Agent** - Creates immediate action plan and connects you with resources |
| 34 | +🔄 **Follow-up Agent** - Designs your long-term support strategy |
| 35 | +""") |
| 36 | + |
| 37 | +st.subheader("Personal Information") |
| 38 | +col1, col2 = st.columns(2) |
| 39 | + |
| 40 | +with col1: |
| 41 | + mental_state = st.text_area("How have you been feeling recently?", |
| 42 | + placeholder="Describe your emotional state, thoughts, or concerns...") |
| 43 | + sleep_pattern = st.select_slider( |
| 44 | + "Sleep Pattern (hours per night)", |
| 45 | + options=[f"{i}" for i in range(0, 13)], |
| 46 | + value="7" |
| 47 | + ) |
| 48 | + |
| 49 | +with col2: |
| 50 | + stress_level = st.slider("Current Stress Level (1-10)", 1, 10, 5) |
| 51 | + support_system = st.multiselect( |
| 52 | + "Current Support System", |
| 53 | + ["Family", "Friends", "Therapist", "Support Groups", "None"] |
| 54 | + ) |
| 55 | + |
| 56 | +recent_changes = st.text_area( |
| 57 | + "Any significant life changes or events recently?", |
| 58 | + placeholder="Job changes, relationships, losses, etc..." |
| 59 | +) |
| 60 | + |
| 61 | +current_symptoms = st.multiselect( |
| 62 | + "Current Symptoms", |
| 63 | + ["Anxiety", "Depression", "Insomnia", "Fatigue", "Loss of Interest", |
| 64 | + "Difficulty Concentrating", "Changes in Appetite", "Social Withdrawal", |
| 65 | + "Mood Swings", "Physical Discomfort"] |
| 66 | +) |
| 67 | + |
| 68 | +if st.button("Get Support Plan"): |
| 69 | + if not api_key: |
| 70 | + st.error("Please enter your OpenAI API key.") |
| 71 | + else: |
| 72 | + with st.spinner('🤖 AI Agents are analyzing your situation...'): |
| 73 | + try: |
| 74 | + task = f""" |
| 75 | + Create a comprehensive mental health support plan based on: |
| 76 | + |
| 77 | + Emotional State: {mental_state} |
| 78 | + Sleep: {sleep_pattern} hours per night |
| 79 | + Stress Level: {stress_level}/10 |
| 80 | + Support System: {', '.join(support_system) if support_system else 'None reported'} |
| 81 | + Recent Changes: {recent_changes} |
| 82 | + Current Symptoms: {', '.join(current_symptoms) if current_symptoms else 'None reported'} |
| 83 | + """ |
| 84 | + |
| 85 | + system_messages = { |
| 86 | + "assessment_agent": """ |
| 87 | + You are an experienced mental health professional speaking directly to the user. Your task is to: |
| 88 | + 1. Create a safe space by acknowledging their courage in seeking support |
| 89 | + 2. Analyze their emotional state with clinical precision and genuine empathy |
| 90 | + 3. Ask targeted follow-up questions to understand their full situation |
| 91 | + 4. Identify patterns in their thoughts, behaviors, and relationships |
| 92 | + 5. Assess risk levels with validated screening approaches |
| 93 | + 6. Help them understand their current mental health in accessible language |
| 94 | + 7. Validate their experiences without minimizing or catastrophizing |
| 95 | +
|
| 96 | + Always use "you" and "your" when addressing the user. Blend clinical expertise with genuine warmth and never rush to conclusions. |
| 97 | + """, |
| 98 | + |
| 99 | + "action_agent": """ |
| 100 | + You are a crisis intervention and resource specialist speaking directly to the user. Your task is to: |
| 101 | + 1. Provide immediate evidence-based coping strategies tailored to their specific situation |
| 102 | + 2. Prioritize interventions based on urgency and effectiveness |
| 103 | + 3. Connect them with appropriate mental health services while acknowledging barriers (cost, access, stigma) |
| 104 | + 4. Create a concrete daily wellness plan with specific times and activities |
| 105 | + 5. Suggest specific support communities with details on how to join |
| 106 | + 6. Balance crisis resources with empowerment techniques |
| 107 | + 7. Teach simple self-regulation techniques they can use immediately |
| 108 | +
|
| 109 | + Focus on practical, achievable steps that respect their current capacity and energy levels. Provide options ranging from minimal effort to more involved actions. |
| 110 | + """, |
| 111 | + |
| 112 | + "followup_agent": """ |
| 113 | + You are a mental health recovery planner speaking directly to the user. Your task is to: |
| 114 | + 1. Design a personalized long-term support strategy with milestone markers |
| 115 | + 2. Create a progress monitoring system that matches their preferences and habits |
| 116 | + 3. Develop specific relapse prevention strategies based on their unique triggers |
| 117 | + 4. Establish a support network mapping exercise to identify existing resources |
| 118 | + 5. Build a graduated self-care routine that evolves with their recovery |
| 119 | + 6. Plan for setbacks with self-compassion techniques |
| 120 | + 7. Set up a maintenance schedule with clear check-in mechanisms |
| 121 | +
|
| 122 | + Focus on building sustainable habits that integrate with their lifestyle and values. Emphasize progress over perfection and teach skills for self-directed care. |
| 123 | + """ |
| 124 | + } |
| 125 | + |
| 126 | + llm_config = { |
| 127 | + "config_list": [{"model": "gpt-4o", "api_key": api_key}] |
| 128 | + } |
| 129 | + |
| 130 | + context_variables = { |
| 131 | + "assessment": None, |
| 132 | + "action": None, |
| 133 | + "followup": None, |
| 134 | + } |
| 135 | + |
| 136 | + def update_assessment_overview(assessment_summary: str, context_variables: dict) -> SwarmResult: |
| 137 | + context_variables["assessment"] = assessment_summary |
| 138 | + st.sidebar.success('Assessment: ' + assessment_summary) |
| 139 | + return SwarmResult(agent="action_agent", context_variables=context_variables) |
| 140 | + |
| 141 | + def update_action_overview(action_summary: str, context_variables: dict) -> SwarmResult: |
| 142 | + context_variables["action"] = action_summary |
| 143 | + st.sidebar.success('Action Plan: ' + action_summary) |
| 144 | + return SwarmResult(agent="followup_agent", context_variables=context_variables) |
| 145 | + |
| 146 | + def update_followup_overview(followup_summary: str, context_variables: dict) -> SwarmResult: |
| 147 | + context_variables["followup"] = followup_summary |
| 148 | + st.sidebar.success('Follow-up Strategy: ' + followup_summary) |
| 149 | + return SwarmResult(agent="assessment_agent", context_variables=context_variables) |
| 150 | + |
| 151 | + def update_system_message_func(agent: SwarmAgent, messages) -> str: |
| 152 | + system_prompt = system_messages[agent.name] |
| 153 | + current_gen = agent.name.split("_")[0] |
| 154 | + |
| 155 | + if agent._context_variables.get(current_gen) is None: |
| 156 | + system_prompt += f"Call the update function provided to first provide a 2-3 sentence summary of your ideas on {current_gen.upper()} based on the context provided." |
| 157 | + agent.llm_config['tool_choice'] = {"type": "function", "function": {"name": f"update_{current_gen}_overview"}} |
| 158 | + else: |
| 159 | + agent.llm_config["tools"] = None |
| 160 | + agent.llm_config['tool_choice'] = None |
| 161 | + system_prompt += f"\n\nYour task\nYou task is write the {current_gen} part of the report. Do not include any other parts. Do not use XML tags.\nStart your reponse with: '## {current_gen.capitalize()} Design'." |
| 162 | + k = list(agent._oai_messages.keys())[-1] |
| 163 | + agent._oai_messages[k] = agent._oai_messages[k][:1] |
| 164 | + |
| 165 | + system_prompt += f"\n\n\nBelow are some context for you to refer to:" |
| 166 | + for k, v in agent._context_variables.items(): |
| 167 | + if v is not None: |
| 168 | + system_prompt += f"\n{k.capitalize()} Summary:\n{v}" |
| 169 | + |
| 170 | + agent.client = OpenAIWrapper(**agent.llm_config) |
| 171 | + return system_prompt |
| 172 | + |
| 173 | + state_update = UPDATE_SYSTEM_MESSAGE(update_system_message_func) |
| 174 | + |
| 175 | + assessment_agent = SwarmAgent( |
| 176 | + "assessment_agent", |
| 177 | + llm_config=llm_config, |
| 178 | + functions=update_assessment_overview, |
| 179 | + update_agent_state_before_reply=[state_update] |
| 180 | + ) |
| 181 | + |
| 182 | + action_agent = SwarmAgent( |
| 183 | + "action_agent", |
| 184 | + llm_config=llm_config, |
| 185 | + functions=update_action_overview, |
| 186 | + update_agent_state_before_reply=[state_update] |
| 187 | + ) |
| 188 | + |
| 189 | + followup_agent = SwarmAgent( |
| 190 | + "followup_agent", |
| 191 | + llm_config=llm_config, |
| 192 | + functions=update_followup_overview, |
| 193 | + update_agent_state_before_reply=[state_update] |
| 194 | + ) |
| 195 | + |
| 196 | + assessment_agent.register_hand_off(AFTER_WORK(action_agent)) |
| 197 | + action_agent.register_hand_off(AFTER_WORK(followup_agent)) |
| 198 | + followup_agent.register_hand_off(AFTER_WORK(assessment_agent)) |
| 199 | + |
| 200 | + result, _, _ = initiate_swarm_chat( |
| 201 | + initial_agent=assessment_agent, |
| 202 | + agents=[assessment_agent, action_agent, followup_agent], |
| 203 | + user_agent=None, |
| 204 | + messages=task, |
| 205 | + max_rounds=13, |
| 206 | + ) |
| 207 | + |
| 208 | + st.session_state.output = { |
| 209 | + 'assessment': result.chat_history[-3]['content'], |
| 210 | + 'action': result.chat_history[-2]['content'], |
| 211 | + 'followup': result.chat_history[-1]['content'] |
| 212 | + } |
| 213 | + |
| 214 | + with st.expander("Situation Assessment"): |
| 215 | + st.markdown(st.session_state.output['assessment']) |
| 216 | + |
| 217 | + with st.expander("Action Plan & Resources"): |
| 218 | + st.markdown(st.session_state.output['action']) |
| 219 | + |
| 220 | + with st.expander("Long-term Support Strategy"): |
| 221 | + st.markdown(st.session_state.output['followup']) |
| 222 | + |
| 223 | + st.success('✨ Mental health support plan generated successfully!') |
| 224 | + |
| 225 | + except Exception as e: |
| 226 | + st.error(f"An error occurred: {str(e)}") |
0 commit comments