Add CrewAI orchestrator for agent coordination

agents/orchestrator.py

@@ -1,128 +1,195 @@
-from typing import Dict, List, Optional
-from .base_agent import BaseWellnessAgent
-from .conversation_agent import ConversationAgent
-from .assessment_agent import AssessmentAgent
-from .mindfulness_agent import MindfulnessAgent
-from .crisis_agent import CrisisAgent
+from typing import Dict, List
+from crewai import Crew, Process, Task
+from agents.conversation_agent import ConversationAgent
+from agents.assessment_agent import AssessmentAgent
+from agents.mindfulness_agent import MindfulnessAgent
+from agents.crisis_agent import CrisisAgent
+import logging
+from utils.log_manager import LogManager
 
 class WellnessOrchestrator:
-    """Orchestrates multiple specialized agents for mental wellness support"""
+    """Orchestrates the coordination between different agents"""
     
-    def __init__(self, config: Dict):
-        self.config = config
-        self.agents = self._initialize_agents()
-        self.current_agent: Optional[BaseWellnessAgent] = None
-        self.session_history: List[Dict] = []
+    def __init__(self, model_config: Dict):
+        self.model_config = model_config
+        self.log_manager = LogManager()
+        self.logger = self.log_manager.get_agent_logger("orchestrator")
         
-    def _initialize_agents(self) -> Dict[str, BaseWellnessAgent]:
-        """Initialize all specialized agents"""
-        return {
-            "conversation": ConversationAgent(
-                model_config=self.config["MODEL_CONFIGS"]
-            ),
-            "assessment": AssessmentAgent(
-                model_config=self.config["MODEL_CONFIGS"]
-            ),
-            "mindfulness": MindfulnessAgent(
-                model_config=self.config["MODEL_CONFIGS"]
-            ),
-            "crisis": CrisisAgent(
-                model_config=self.config["MODEL_CONFIGS"]
-            )
-        }
-        
-    def process_message(self, message: str, context: Optional[Dict] = None) -> Dict:
-        """Process incoming message and route to appropriate agent"""
-        # Update context for all agents
-        if context:
-            for agent in self.agents.values():
-                agent.update_context(context)
-                
-        # Check for crisis keywords first
-        if self._is_crisis_situation(message):
-            self.current_agent = self.agents["crisis"]
-            return self.current_agent.process_message(message)
-            
-        # Route to appropriate agent based on message content and context
-        agent_key = self._determine_best_agent(message, context)
-        self.current_agent = self.agents[agent_key]
-        
-        # Process message with selected agent
-        response = self.current_agent.process_message(message)
+        # Initialize agents
+        self.initialize_agents()
         
-        # Record interaction in session history
-        self._record_interaction(message, response, agent_key)
+        # Initialize CrewAI
+        self.initialize_crew()
         
-        return response
+    def initialize_agents(self):
+        """Initialize all agents with their specific roles and tools"""
+        self.logger.info("Initializing agents")
         
-    def _is_crisis_situation(self, message: str) -> bool:
-        """Check if message indicates a crisis situation"""
-        crisis_keywords = [
-            "suicide", "kill myself", "end it all", "self harm",
-            "hurt myself", "die", "death", "emergency"
-        ]
-        return any(keyword in message.lower() for keyword in crisis_keywords)
-        
-    def _determine_best_agent(self, message: str, context: Optional[Dict]) -> str:
-        """Determine the most appropriate agent for the message"""
-        # Check for explicit commands or keywords
-        if any(cmd in message.lower() for cmd in ["assess", "evaluation", "test"]):
-            return "assessment"
-        if any(cmd in message.lower() for cmd in ["meditate", "breathe", "relax"]):
-            return "mindfulness"
-            
-        # Default to conversation agent if no specific needs detected
-        return "conversation"
-        
-    def _record_interaction(self, message: str, response: Dict, agent_key: str):
-        """Record interaction in session history"""
-        self.session_history.append({
-            "timestamp": response.get("timestamp"),
-            "user_message": message,
-            "agent_response": response,
-            "agent_type": agent_key
-        })
+        try:
+            # Conversation Agent
+            self.conversation_agent = ConversationAgent(
+                name="Therapeutic Conversation Agent",
+                role="Lead conversation therapist",
+                goal="Guide therapeutic conversations and provide emotional support",
+                backstory="Expert in therapeutic dialogue and emotional support",
+                tools=["chat", "emotion_detection"],
+                model_config=self.model_config
+            )
+            
+            # Assessment Agent
+            self.assessment_agent = AssessmentAgent(
+                name="Mental Health Assessment Agent",
+                role="Mental health evaluator",
+                goal="Conduct mental health assessments and track progress",
+                backstory="Specialist in mental health evaluation and monitoring",
+                tools=["assessment_tools", "progress_tracking"],
+                model_config=self.model_config
+            )
+            
+            # Mindfulness Agent
+            self.mindfulness_agent = MindfulnessAgent(
+                name="Mindfulness Guide Agent",
+                role="Mindfulness and meditation instructor",
+                goal="Guide mindfulness exercises and meditation sessions",
+                backstory="Expert in mindfulness techniques and meditation",
+                tools=["meditation_guide", "breathing_exercises"],
+                model_config=self.model_config
+            )
+            
+            # Crisis Agent
+            self.crisis_agent = CrisisAgent(
+                name="Crisis Intervention Agent",
+                role="Emergency response specialist",
+                goal="Provide immediate support in crisis situations",
+                backstory="Trained in crisis intervention and emergency response",
+                tools=["crisis_protocol", "emergency_resources"],
+                model_config=self.model_config
+            )
+            
+            self.logger.info("All agents initialized successfully")
+            
+        except Exception as e:
+            self.logger.error(f"Error initializing agents: {str(e)}")
+            raise
+            
+    def initialize_crew(self):
+        """Initialize CrewAI with agents and tasks"""
+        self.logger.info("Initializing CrewAI")
         
-    def start_assessment(self, assessment_type: str) -> Dict:
-        """Start a new assessment"""
-        self.current_agent = self.agents["assessment"]
-        return self.agents["assessment"].start_assessment(assessment_type)
+        try:
+            # Create the crew
+            self.crew = Crew(
+                agents=[
+                    self.conversation_agent,
+                    self.assessment_agent,
+                    self.mindfulness_agent,
+                    self.crisis_agent
+                ],
+                tasks=[],  # Tasks will be added dynamically
+                process=Process.sequential  # Can be changed to parallel if needed
+            )
+            
+            self.logger.info("CrewAI initialized successfully")
+            
+        except Exception as e:
+            self.logger.error(f"Error initializing CrewAI: {str(e)}")
+            raise
+            
+    def create_task(self, task_type: str, description: str, agent) -> Task:
+        """Create a task for an agent"""
+        return Task(
+            description=description,
+            agent=agent,
+            expected_output="Detailed response with next steps"
+        )
         
-    def start_mindfulness_session(self, session_type: str) -> Dict:
-        """Start a new mindfulness session"""
-        self.current_agent = self.agents["mindfulness"]
-        return self.agents["mindfulness"].start_session(session_type)
+    def process_message(self, message: str, context: Dict = None) -> Dict:
+        """Process user message through appropriate agents"""
+        self.logger.info("Processing message through agents")
         
-    def get_session_summary(self) -> Dict:
-        """Generate summary of current session"""
+        try:
+            # Clear previous tasks
+            self.crew.tasks = []
+            
+            # Initial assessment by conversation agent
+            initial_task = self.create_task(
+                "initial_assessment",
+                f"Analyze this message and determine required support: {message}",
+                self.conversation_agent
+            )
+            self.crew.tasks.append(initial_task)
+            
+            # Analyze for crisis indicators
+            crisis_check = self.create_task(
+                "crisis_check",
+                f"Check for crisis indicators in: {message}",
+                self.crisis_agent
+            )
+            self.crew.tasks.append(crisis_check)
+            
+            # Execute the crew tasks
+            result = self.crew.kickoff()
+            
+            # Process results and determine next steps
+            if "crisis" in result.lower():
+                # Add crisis intervention task
+                crisis_task = self.create_task(
+                    "crisis_intervention",
+                    f"Provide crisis intervention for: {message}",
+                    self.crisis_agent
+                )
+                self.crew.tasks = [crisis_task]
+                response = self.crew.kickoff()
+                
+            elif "assessment" in result.lower():
+                # Add assessment task
+                assessment_task = self.create_task(
+                    "mental_health_assessment",
+                    f"Conduct mental health assessment based on: {message}",
+                    self.assessment_agent
+                )
+                self.crew.tasks = [assessment_task]
+                response = self.crew.kickoff()
+                
+            elif "mindfulness" in result.lower():
+                # Add mindfulness task
+                mindfulness_task = self.create_task(
+                    "mindfulness_session",
+                    f"Guide mindfulness exercise based on: {message}",
+                    self.mindfulness_agent
+                )
+                self.crew.tasks = [mindfulness_task]
+                response = self.crew.kickoff()
+                
+            else:
+                # Continue therapeutic conversation
+                conversation_task = self.create_task(
+                    "therapeutic_conversation",
+                    f"Continue therapeutic conversation: {message}",
+                    self.conversation_agent
+                )
+                self.crew.tasks = [conversation_task]
+                response = self.crew.kickoff()
+            
+            return {
+                "message": response,
+                "agent_type": self.crew.tasks[-1].agent.name,
+                "task_type": self.crew.tasks[-1].task_type
+            }
+            
+        except Exception as e:
+            self.logger.error(f"Error processing message: {str(e)}")
+            return {
+                "message": "I apologize, but I encountered an error. Please try again.",
+                "agent_type": "error",
+                "task_type": "error_handling"
+            }
+            
+    def get_agent_status(self) -> Dict:
+        """Get status of all agents"""
         return {
-            "interaction_count": len(self.session_history),
-            "agent_usage": self._calculate_agent_usage(),
-            "key_insights": self._generate_insights(),
-            "recommendations": self._generate_recommendations()
-        }
-        
-    def _calculate_agent_usage(self) -> Dict[str, int]:
-        """Calculate how often each agent was used"""
-        usage = {}
-        for interaction in self.session_history:
-            agent_type = interaction["agent_type"]
-            usage[agent_type] = usage.get(agent_type, 0) + 1
-        return usage
-        
-    def _generate_insights(self) -> List[str]:
-        """Generate insights from session history"""
-        # Implement insight generation logic
-        return []
-        
-    def _generate_recommendations(self) -> List[str]:
-        """Generate recommendations based on session history"""
-        # Implement recommendation generation logic
-        return []
-        
-    def reset_session(self):
-        """Reset the current session"""
-        self.session_history = []
-        self.current_agent = None
-        for agent in self.agents.values():
-            agent.clear_state() 
+            "conversation": self.conversation_agent.get_status(),
+            "assessment": self.assessment_agent.get_status(),
+            "mindfulness": self.mindfulness_agent.get_status(),
+            "crisis": self.crisis_agent.get_status()
+        } 

interface/app.py

@@ -5,10 +5,7 @@ import logging
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 from utils.log_manager import LogManager
 from utils.analytics_logger import AnalyticsLogger
-from agents.conversation_agent import ConversationAgent
-from agents.assessment_agent import AssessmentAgent
-from agents.mindfulness_agent import MindfulnessAgent
-from agents.crisis_agent import CrisisAgent
+from agents.orchestrator import WellnessOrchestrator
 
 # Force CPU-only mode
 torch.cuda.is_available = lambda: False
@@ -29,8 +26,8 @@ class WellnessInterface:
         # Initialize models
         self.initialize_models()
         
-        # Initialize agents
-        self.initialize_agents()
+        # Initialize orchestrator
+        self.initialize_orchestrator()
         
         # Initialize interface
         self.setup_interface()
@@ -61,28 +58,16 @@ class WellnessInterface:
             self.logger.error(f"Error initializing models: {str(e)}")
             raise
             
-    def initialize_agents(self):
-        """Initialize AI agents"""
-        self.logger.info("Initializing AI agents")
+    def initialize_orchestrator(self):
+        """Initialize CrewAI orchestrator"""
+        self.logger.info("Initializing CrewAI orchestrator")
         try:
-            # Initialize all agents
-            self.conversation_agent = ConversationAgent(
+            self.orchestrator = WellnessOrchestrator(
                 model_config=self.config["MODEL_CONFIGS"]
             )
-            self.assessment_agent = AssessmentAgent(
-                model_config=self.config["MODEL_CONFIGS"]
-            )
-            self.mindfulness_agent = MindfulnessAgent(
-                model_config=self.config["MODEL_CONFIGS"]
-            )
-            self.crisis_agent = CrisisAgent(
-                model_config=self.config["MODEL_CONFIGS"]
-            )
-            
-            self.logger.info("AI agents initialized successfully")
-            
+            self.logger.info("Orchestrator initialized successfully")
         except Exception as e:
-            self.logger.error(f"Error initializing agents: {str(e)}")
+            self.logger.error(f"Error initializing orchestrator: {str(e)}")
             raise
             
     def setup_interface(self):
@@ -232,27 +217,30 @@ class WellnessInterface:
                 }}
             )
             
-            # Analyze emotion if text is present
-            emotion = None
-            if text:
-                emotion_result = self.emotion_model(text)
-                emotion = emotion_result[0] if emotion_result else None
-                self.logger.info(f"Detected emotion: {emotion}")
+            # Process through orchestrator
+            context = {
+                "history": history,
+                "emotion": self.emotion_model(text)[0] if text else None,
+                "has_audio": bool(audio),
+                "has_image": bool(image)
+            }
             
-            # Route to appropriate agent based on content and emotion
-            if emotion and emotion.get("label") in ["anxiety", "fear", "panic"]:
-                response = self.crisis_agent.process_message(text)
-            elif "meditate" in text.lower() or "mindfulness" in text.lower():
-                response = self.mindfulness_agent.process_message(text)
-            elif "assess" in text.lower() or "check" in text.lower():
-                response = self.assessment_agent.process_message(text)
-            else:
-                response = self.conversation_agent.process_message(text)
+            response = self.orchestrator.process_message(
+                message=text if text else "Sent media",
+                context=context
+            )
             
             # Add to chat history using message format
             history = history or []
             history.append({"role": "user", "content": text if text else "Sent media"})
-            history.append({"role": "assistant", "content": response["message"]})
+            history.append({
+                "role": "assistant",
+                "content": response["message"],
+                "metadata": {
+                    "agent": response["agent_type"],
+                    "task": response["task_type"]
+                }
+            })
             
             return history, ""  # Return empty string to clear text input
             
@@ -274,20 +262,19 @@ class WellnessInterface:
         """Provide emergency help information"""
         self.logger.info("Emergency help requested")
         
+        # Use crisis agent through orchestrator
+        response = self.orchestrator.process_message(
+            message="EMERGENCY_HELP_REQUESTED",
+            context={"is_emergency": True}
+        )
+        
         return [{
             "role": "assistant",
-            "content": """🚨 EMERGENCY RESOURCES 🚨
-
-If you're experiencing a mental health emergency:
-
-📞 Emergency Services: 911 (US)
-🆘 National Crisis Hotline: 988
-💭 Crisis Text Line: Text HOME to 741741
-
-These services are available 24/7 and are staffed by trained professionals.
-Your life matters, and help is available immediately.
-
-Please don't hesitate to reach out - caring people are ready to help."""
+            "content": response["message"],
+            "metadata": {
+                "agent": response["agent_type"],
+                "task": response["task_type"]
+            }
         }]
         
     def launch(self, **kwargs):