removed unnecessary imports and code

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
+        'assessment': '',
-        'action': '',     # Immediate actions and resources
+        'action': '',
-        'followup': ''    # Long-term planning
+        '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: