agents/core/orchestrator.py

"""
Multi-Agent Orchestrator
Coordinates multiple agents to handle complex queries
"""

from typing import List, Dict, Any
import json

class MultiAgentOrchestrator:
    """
    Orchestrates multiple agents to handle complex queries
    that require expertise from multiple domains
    """
    
    def __init__(self, agents: Dict[str, Any]):
        """
        Initialize orchestrator with available agents
        
        Args:
            agents: Dictionary of agent_name -> agent_instance
        """
        self.agents = agents
    
    def orchestrate(self, query: str, agent_names: List[str], chat_history: List = None) -> str:
        """
        Orchestrate multiple agents to answer a complex query
        
        Args:
            query: User query
            agent_names: List of agents to use
            chat_history: Conversation history
        
        Returns:
            Combined response from all agents
        """
        if len(agent_names) == 1:
            # Single agent - just call it
            return self._call_single_agent(agent_names[0], query, chat_history)
        
        # Multi-agent orchestration
        return self._orchestrate_multi_agent(query, agent_names, chat_history)
    
    def _call_single_agent(self, agent_name: str, query: str, chat_history: List) -> str:
        """Call a single agent"""
        agent = self.agents.get(agent_name)
        if not agent:
            return f"Agent {agent_name} not found"
        
        try:
            response = agent.handle(
                parameters={"user_query": query},
                chat_history=chat_history
            )
            return response
        except Exception as e:
            return f"Error calling {agent_name}: {str(e)}"
    
    def _orchestrate_multi_agent(self, query: str, agent_names: List[str], chat_history: List) -> str:
        """
        Orchestrate multiple agents
        
        Strategy:
        1. Analyze query to determine what each agent should focus on
        2. Call each agent with specific sub-query
        3. Combine responses intelligently
        """
        # Decompose query into sub-queries for each agent
        sub_queries = self._decompose_query(query, agent_names)
        
        # Call each agent with their sub-query
        responses = {}
        for agent_name, sub_query in sub_queries.items():
            if agent_name in self.agents:
                try:
                    response = self.agents[agent_name].handle(
                        parameters={"user_query": sub_query},
                        chat_history=chat_history
                    )
                    responses[agent_name] = response
                except Exception as e:
                    responses[agent_name] = f"Error: {str(e)}"
        
        # Combine responses
        return self._combine_responses(query, responses, agent_names)
    
    def _decompose_query(self, query: str, agent_names: List[str]) -> Dict[str, str]:
        """
        Decompose complex query into sub-queries for each agent
        
        Example:
        Query: "Tôi muốn giảm cân, nên ăn gì và tập gì?"
        →
        nutrition_agent: "Tư vấn chế độ ăn để giảm cân"
        exercise_agent: "Tư vấn lịch tập để giảm cân"
        """
        sub_queries = {}
        
        # Simple heuristic-based decomposition
        query_lower = query.lower()
        
        for agent_name in agent_names:
            if agent_name == "nutrition_agent":
                if "ăn" in query_lower or "dinh dưỡng" in query_lower or "calo" in query_lower:
                    sub_queries[agent_name] = f"Tư vấn dinh dưỡng cho: {query}"
                else:
                    sub_queries[agent_name] = query
            
            elif agent_name == "exercise_agent":
                if "tập" in query_lower or "gym" in query_lower or "luyện" in query_lower:
                    sub_queries[agent_name] = f"Tư vấn tập luyện cho: {query}"
                else:
                    sub_queries[agent_name] = query
            
            elif agent_name == "mental_health_agent":
                if "stress" in query_lower or "lo âu" in query_lower:
                    sub_queries[agent_name] = f"Tư vấn sức khỏe tinh thần cho: {query}"
                else:
                    sub_queries[agent_name] = query
            
            else:
                # Default: use original query
                sub_queries[agent_name] = query
        
        return sub_queries
    
    def _combine_responses(self, original_query: str, responses: Dict[str, str], agent_names: List[str]) -> str:
        """
        Combine responses from multiple agents into a coherent answer
        
        Strategy:
        1. Identify the main topic
        2. Structure response logically
        3. Avoid redundancy
        """
        if not responses:
            return "Xin lỗi, không thể xử lý câu hỏi này."
        
        # Build combined response
        combined = []
        
        # Add intro
        if len(responses) > 1:
            combined.append("Để giải đáp câu hỏi của bạn, tôi sẽ tư vấn từ nhiều góc độ:\n")
        
        # Add each agent's response with clear sections
        agent_labels = {
            "nutrition_agent": "📊 Dinh Dưỡng",
            "exercise_agent": "💪 Tập Luyện",
            "mental_health_agent": "🧠 Sức Khỏe Tinh Thần",
            "symptom_agent": "🩺 Đánh Giá Triệu Chứng",
            "general_health_agent": "🏥 Tổng Quan"
        }
        
        for agent_name in agent_names:
            if agent_name in responses:
                label = agent_labels.get(agent_name, agent_name)
                response = responses[agent_name]
                
                # Clean up response (remove redundant intro)
                response = self._clean_response(response)
                
                combined.append(f"\n{label}:\n{response}\n")
        
        # Add conclusion
        if len(responses) > 1:
            combined.append("\n---\n")
            combined.append("💡 Lưu ý: Để đạt kết quả tốt nhất, hãy kết hợp cả dinh dưỡng, tập luyện và nghỉ ngơi hợp lý.")
        
        return "".join(combined)
    
    def _clean_response(self, response: str) -> str:
        """Clean up response by removing redundant intros"""
        # Remove common intro phrases
        intros_to_remove = [
            "Chào bạn!",
            "Xin chào!",
            "Để giải đáp câu hỏi của bạn",
            "Mình sẽ giúp bạn",
        ]
        
        for intro in intros_to_remove:
            if response.startswith(intro):
                response = response[len(intro):].strip()
        
        return response


# Example usage
if __name__ == "__main__":
    # Mock agents for testing
    class MockAgent:
        def __init__(self, name):
            self.name = name
        
        def handle(self, parameters, chat_history=None):
            query = parameters.get("user_query", "")
            return f"[{self.name}] Response to: {query}"
    
    agents = {
        "nutrition_agent": MockAgent("Nutrition"),
        "exercise_agent": MockAgent("Exercise"),
        "mental_health_agent": MockAgent("Mental Health")
    }
    
    orchestrator = MultiAgentOrchestrator(agents)
    
    # Test multi-agent
    query = "Tôi muốn giảm cân, nên ăn gì và tập gì?"
    agent_names = ["nutrition_agent", "exercise_agent"]
    
    response = orchestrator.orchestrate(query, agent_names)
    print(response)