AI Mental Wellbeing Agent

system_messages = {
    "assessment_agent": """
    You are an experienced mental health professional speaking 
    directly to the user. Your task is to:
    
    1. Create a safe space by acknowledging their courage 
       in seeking support
    2. Analyze their emotional state with clinical precision 
       and genuine empathy
    3. Ask targeted follow-up questions to understand 
       their full situation
    4. Identify patterns in their thoughts, behaviors, 
       and relationships
    5. Assess risk levels with validated screening approaches
    6. Help them understand their current mental health 
       in accessible language
    7. Validate their experiences without minimizing 
       or catastrophizing
    
    Always use "you" and "your" when addressing the user. 
    Blend clinical expertise with genuine warmth and never 
    rush to conclusions.
    """,
    
    "action_agent": """
    You are a crisis intervention and resource specialist 
    speaking directly to the user. Your task is to:
    
    1. Provide immediate evidence-based coping strategies 
       tailored to their specific situation
    2. Prioritize interventions based on urgency and effectiveness
    3. Connect them with appropriate mental health services 
       while acknowledging barriers (cost, access, stigma)
    4. Create a concrete daily wellness plan with specific 
       times and activities
    5. Suggest specific support communities with details 
       on how to join
    6. Balance crisis resources with empowerment techniques
    7. Teach simple self-regulation techniques they can 
       use immediately
    
    Focus on practical, achievable steps that respect their 
    current capacity and energy levels. Provide options ranging 
    from minimal effort to more involved actions.
    """,
    
    "followup_agent": """
    You are a mental health recovery planner speaking directly 
    to the user. Your task is to:
    
    1. Design a personalized long-term support strategy 
       with milestone markers
    2. Create a progress monitoring system that matches 
       their preferences and habits
    3. Develop specific relapse prevention strategies based 
       on their unique triggers
    4. Establish a support network mapping exercise to 
       identify existing resources
    5. Build a graduated self-care routine that evolves 
       with their recovery
    6. Plan for setbacks with self-compassion techniques
    7. Set up a maintenance schedule with clear check-in 
       mechanisms
    
    Focus on building sustainable habits that integrate with 
    their lifestyle and values. Emphasize progress over 
    perfection and teach skills for self-directed care.
    """
}

# Functions for agent coordination

def update_assessment_overview(
    assessment_summary: str,
    context_variables: dict
) -> SwarmResult:
    """Store assessment and hand off to action agent."""
    context_variables["assessment"] = assessment_summary
    st.sidebar.success('Assessment: ' + assessment_summary)
    return SwarmResult(
        agent="action_agent",
        context_variables=context_variables
    )

def update_action_overview(
    action_summary: str,
    context_variables: dict
) -> SwarmResult:
    """Store action plan and hand off to follow-up agent."""
    context_variables["action"] = action_summary
    st.sidebar.success('Action Plan: ' + action_summary)
    return SwarmResult(
        agent="followup_agent",
        context_variables=context_variables
    )

def update_followup_overview(
    followup_summary: str,
    context_variables: dict
) -> SwarmResult:
    """Store follow-up strategy and complete cycle."""
    context_variables["followup"] = followup_summary
    st.sidebar.success('Follow-up Strategy: ' + followup_summary)
    return SwarmResult(
        agent="assessment_agent",
        context_variables=context_variables
    )

# Prepare user task
task = f"""
Create a comprehensive mental health support plan based on:

Emotional State: {mental_state}
Sleep: {sleep_pattern} hours per night
Stress Level: {stress_level}/10
Support System: {', '.join(support_system)}
Recent Changes: {recent_changes}
Current Symptoms: {', '.join(current_symptoms)}
"""

# Execute swarm chat
result, _, _ = initiate_swarm_chat(
    initial_agent=assessment_agent,
    agents=[assessment_agent, action_agent, followup_agent],
    user_agent=None,
    messages=task,
    max_rounds=13,  # 2 phases per agent + initial
)

# Extract outputs
output = {
    'assessment': result.chat_history[-3]['content'],
    'action': result.chat_history[-2]['content'],
    'followup': result.chat_history[-1]['content']
}

# Display in expandable sections
with st.expander("Situation Assessment"):
    st.markdown(output['assessment'])

with st.expander("Action Plan & Resources"):
    st.markdown(output['action'])

with st.expander("Long-term Support Strategy"):
    st.markdown(output['followup'])

import streamlit as st

st.title("🧠 Mental Wellbeing Agent")

st.info("""
**Meet Your Mental Wellbeing Agent Team:**

🧠 **Assessment Agent** - Analyzes your situation and needs
🎯 **Action Agent** - Creates immediate action plan
🔄 **Follow-up Agent** - Designs long-term support strategy
""")

# User inputs
st.subheader("Personal Information")
col1, col2 = st.columns(2)

with col1:
    mental_state = st.text_area(
        "How have you been feeling recently?",
        placeholder="Describe your emotional state..."
    )
    sleep_pattern = st.select_slider(
        "Sleep Pattern (hours per night)",
        options=[f"{i}" for i in range(0, 13)],
        value="7"
    )

with col2:
    stress_level = st.slider(
        "Current Stress Level (1-10)",
        1, 10, 5
    )
    support_system = st.multiselect(
        "Current Support System",
        ["Family", "Friends", "Therapist", 
         "Support Groups", "None"]
    )

recent_changes = st.text_area(
    "Any significant life changes recently?",
    placeholder="Job changes, relationships, losses..."
)

current_symptoms = st.multiselect(
    "Current Symptoms",
    ["Anxiety", "Depression", "Insomnia", "Fatigue",
     "Loss of Interest", "Difficulty Concentrating",
     "Changes in Appetite", "Social Withdrawal",
     "Mood Swings", "Physical Discomfort"]
)

if st.button("Get Support Plan"):
    # Execute swarm and display results
    pass

# Sleep patterns
sleep_pattern = st.select_slider(
    "Sleep Pattern (hours per night)",
    options=[f"{i}" for i in range(0, 13)],
    value="7"
)

# Current symptoms
current_symptoms = st.multiselect(
    "Current Symptoms",
    [
        "Anxiety",
        "Depression",
        "Insomnia",
        "Fatigue",
        "Loss of Interest",
        "Difficulty Concentrating",
        "Changes in Appetite",
        "Social Withdrawal",
        "Mood Swings",
        "Physical Discomfort"
    ]
)

# Stress assessment
stress_level = st.slider(
    "Current Stress Level (1-10)",
    1, 10, 5
)

# Available support
support_system = st.multiselect(
    "Current Support System",
    [
        "Family",
        "Friends",
        "Therapist",
        "Support Groups",
        "None"
    ]
)

# Recent changes
recent_changes = st.text_area(
    "Any significant life changes recently?",
    placeholder="Job changes, relationships, "
               "losses, etc..."
)

# Helps identify:
# - Potential triggers
# - Context for current state
# - Environmental stressors
# - Life transitions
# - Major events