agents/orchestrator.py

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

class WellnessOrchestrator:
    """Orchestrates multiple specialized agents for mental wellness support"""
    
    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 _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)
        
        # Record interaction in session history
        self._record_interaction(message, response, agent_key)
        
        return response
        
    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
        })
        
    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)
        
    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 get_session_summary(self) -> Dict:
        """Generate summary of current session"""
        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()