Update agent initialization for CrewAI compatibility

mentalwellness_space/agents/conversation_agent.py

@@ -0,0 +1,96 @@
+from typing import Dict, List
+from .base_agent import BaseWellnessAgent
+from transformers import pipeline
+
+class ConversationAgent(BaseWellnessAgent):
+    """Agent specialized in therapeutic conversations"""
+    
+    def __init__(self, model_config: Dict, verbose: bool = False):
+        super().__init__(
+            name="Therapeutic Conversation Agent",
+            role="Mental Health Conversationalist",
+            goal="Engage in supportive therapeutic conversations while maintaining empathy and professionalism",
+            backstory="""I am a specialized AI agent trained in therapeutic conversation techniques. 
+            I use evidence-based approaches to provide emotional support and guidance while maintaining 
+            appropriate boundaries and recognizing when to escalate to crisis intervention.""",
+            tools=["emotion_detection", "conversation"],
+            model_config=model_config,
+            verbose=verbose
+        )
+        self._initialize_conversation_model()
+        
+    def _initialize_conversation_model(self):
+        """Initialize the conversation model"""
+        self.models["conversation"] = pipeline(
+            "text-generation",
+            model=self.model_config["conversation"]["model_id"]
+        )
+        
+    def _generate_response(self, prompt: str) -> str:
+        """Generate response using the conversation model"""
+        response = self.models["conversation"](
+            prompt,
+            max_length=self.model_config["conversation"]["max_length"],
+            temperature=self.model_config["conversation"]["temperature"]
+        )
+        return response[0]["generated_text"] if response else ""
+        
+    def process_message(self, message: str) -> Dict:
+        """Process user message and generate therapeutic response"""
+        # Analyze emotion
+        emotion = self.analyze_emotion(message)
+        
+        # Update context with emotion
+        self.update_context({"last_emotion": emotion})
+        
+        # Generate appropriate response based on emotion and context
+        prompt = self._create_prompt(message, emotion)
+        response = self._generate_response(prompt)
+        
+        # Add to conversation history
+        interaction = {
+            "user_message": message,
+            "emotion": emotion,
+            "response": response
+        }
+        self.add_to_history(interaction)
+        
+        return self.format_response(response)
+        
+    def _create_prompt(self, message: str, emotion: Dict) -> str:
+        """Create context-aware prompt for response generation"""
+        history = self.get_history()
+        context = self.get_context()
+        
+        # Build prompt with context
+        prompt = f"""As a therapeutic conversation agent, respond to the following message with empathy and support.
+        User's message: {message}
+        Detected emotion: {emotion}
+        Conversation history: {history[-3:] if history else 'No history'}
+        Additional context: {context}
+        
+        Response:"""
+        
+        return prompt
+        
+    def get_conversation_summary(self) -> Dict:
+        """Generate summary of the conversation"""
+        history = self.get_history()
+        emotions = [interaction["emotion"] for interaction in history]
+        
+        return {
+            "interaction_count": len(history),
+            "emotional_trajectory": emotions,
+            "key_topics": self._extract_key_topics(history),
+            "recommendations": self._generate_recommendations(history)
+        }
+        
+    def _extract_key_topics(self, history: List[Dict]) -> List[str]:
+        """Extract main topics from conversation history"""
+        # Implement topic extraction logic
+        return []
+        
+    def _generate_recommendations(self, history: List[Dict]) -> List[str]:
+        """Generate recommendations based on conversation"""
+        # Implement recommendation logic
+        return [] 

mentalwellness_space/agents/orchestrator.py

@@ -0,0 +1,195 @@
+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 the coordination between different agents"""
+    
+    def __init__(self, model_config: Dict):
+        self.model_config = model_config
+        self.log_manager = LogManager()
+        self.logger = self.log_manager.get_agent_logger("orchestrator")
+        
+        # Initialize agents
+        self.initialize_agents()
+        
+        # Initialize CrewAI
+        self.initialize_crew()
+        
+    def initialize_agents(self):
+        """Initialize all agents with their specific roles and tools"""
+        self.logger.info("Initializing agents")
+        
+        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")
+        
+        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 process_message(self, message: str, context: Dict = None) -> Dict:
+        """Process user message through appropriate agents"""
+        self.logger.info("Processing message through agents")
+        
+        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 {
+            "conversation": self.conversation_agent.get_status(),
+            "assessment": self.assessment_agent.get_status(),
+            "mindfulness": self.mindfulness_agent.get_status(),
+            "crisis": self.crisis_agent.get_status()
+        }