diff --git a/ai_agent_tutorials/ai_mental_health_crisis_agent/ai_mental_health_crisis_agent.py b/ai_agent_tutorials/ai_mental_health_crisis_agent/ai_mental_health_crisis_agent.py index d586722..54d95ec 100644 --- a/ai_agent_tutorials/ai_mental_health_crisis_agent/ai_mental_health_crisis_agent.py +++ b/ai_agent_tutorials/ai_mental_health_crisis_agent/ai_mental_health_crisis_agent.py @@ -1,4 +1,3 @@ -import asyncio import streamlit as st from autogen import ( SwarmAgent, @@ -10,22 +9,18 @@ from autogen import ( ) import os -# Add this at the top of the file, before any other code os.environ["AUTOGEN_USE_DOCKER"] = "0" -# Initialize session state with 3 key components if 'output' not in st.session_state: st.session_state.output = { - 'assessment': '', # Combined psychology/analysis - 'action': '', # Immediate actions and resources - 'followup': '' # Long-term planning + 'assessment': '', + 'action': '', + 'followup': '' } -# Sidebar for API key input st.sidebar.title("OpenAI API Key") api_key = st.sidebar.text_input("Enter your OpenAI API Key", type="password") -# Add privacy notice in sidebar st.sidebar.warning(""" ## ⚠️ Important Notice @@ -36,10 +31,8 @@ This application is a supportive tool and does not replace professional mental h - Seek immediate professional help """) -# Main app UI st.title("🧠 Mental Health Crisis Navigator") -# Update UI description st.info(""" **Meet Your Mental Health Support Team:** @@ -48,7 +41,6 @@ st.info(""" 🔄 **Follow-up Agent** - Designs your long-term support strategy """) -# User inputs st.subheader("Personal Information") col1, col2 = st.columns(2) @@ -68,7 +60,6 @@ with col2: ["Family", "Friends", "Therapist", "Support Groups", "None"] ) -# Additional context recent_changes = st.text_area( "Any significant life changes or events recently?", placeholder="Job changes, relationships, losses, etc..." @@ -81,8 +72,6 @@ current_symptoms = st.multiselect( "Mood Swings", "Physical Discomfort"] ) - -# Button to start the agent collaboration if st.button("Get Support Plan"): if not api_key: st.error("Please enter your OpenAI API key.") @@ -100,7 +89,6 @@ if st.button("Get Support Plan"): Current Symptoms: {', '.join(current_symptoms) if current_symptoms else 'None reported'} """ - # Update system messages for 3 agents system_messages = { "assessment_agent": """ You are an experienced mental health professional speaking directly to the user. Your task is to: @@ -142,69 +130,55 @@ if st.button("Get Support Plan"): """ } - # Then modify the agent configurations llm_config = { "config_list": [{"model": "gpt-4o", "api_key": api_key}] } - # Context management for agent communication context_variables = { "assessment": None, "action": None, "followup": None, } - # Update functions for each agent def update_assessment_overview(assessment_summary: str, context_variables: dict) -> SwarmResult: - """Keep the summary as short as possible.""" 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: - """Keep the summary as short as possible.""" 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: - """Keep the summary as short as possible.""" context_variables["followup"] = followup_summary st.sidebar.success('Follow-up Strategy: ' + followup_summary) return SwarmResult(agent="assessment_agent", context_variables=context_variables) def update_system_message_func(agent: SwarmAgent, messages) -> str: - """""" system_prompt = system_messages[agent.name] - current_gen = agent.name.split("_")[0] + if agent._context_variables.get(current_gen) is None: 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." agent.llm_config['tool_choice'] = {"type": "function", "function": {"name": f"update_{current_gen}_overview"}} - agent.client = OpenAIWrapper(**agent.llm_config) else: - # remove the tools to avoid the agent from using it and reduce cost agent.llm_config["tools"] = None agent.llm_config['tool_choice'] = None - agent.client = OpenAIWrapper(**agent.llm_config) - # the agent has given a summary, now it should generate a detailed response 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'." - - # Remove all messages except the first one with less cost k = list(agent._oai_messages.keys())[-1] agent._oai_messages[k] = agent._oai_messages[k][:1] system_prompt += f"\n\n\nBelow are some context for you to refer to:" - # Add context variables to the prompt for k, v in agent._context_variables.items(): if v is not None: system_prompt += f"\n{k.capitalize()} Summary:\n{v}" + agent.client = OpenAIWrapper(**agent.llm_config) return system_prompt state_update = UPDATE_SYSTEM_MESSAGE(update_system_message_func) - # Initialize agents assessment_agent = SwarmAgent( "assessment_agent", llm_config=llm_config, @@ -226,12 +200,10 @@ if st.button("Get Support Plan"): update_agent_state_before_reply=[state_update] ) - # Update handoffs assessment_agent.register_hand_off(AFTER_WORK(action_agent)) action_agent.register_hand_off(AFTER_WORK(followup_agent)) followup_agent.register_hand_off(AFTER_WORK(assessment_agent)) - # Update result handling result, _, _ = initiate_swarm_chat( initial_agent=assessment_agent, agents=[assessment_agent, action_agent, followup_agent], @@ -240,14 +212,12 @@ if st.button("Get Support Plan"): max_rounds=13, ) - # Update session state with responses st.session_state.output = { 'assessment': result.chat_history[-3]['content'], 'action': result.chat_history[-2]['content'], 'followup': result.chat_history[-1]['content'] } - # Display outputs with st.expander("Situation Assessment"): st.markdown(st.session_state.output['assessment']) @@ -257,7 +227,6 @@ if st.button("Get Support Plan"): with st.expander("Long-term Support Strategy"): st.markdown(st.session_state.output['followup']) - # Display success message after completion st.success('✨ Mental health support plan generated successfully!') except Exception as e: