Multi Agent Orchestration

"""
Framework-Specific Multi-Agent Implementations

Templates for CrewAI, AutoGen, LangGraph, and Swarm.
"""

from typing import Dict, List, Optional, Any
from dataclasses import dataclass


@dataclass
class CrewAITemplate:
    """Template for CrewAI implementation."""

    @staticmethod
    def create_financial_team() -> Dict[str, Any]:
        """Create CrewAI financial analysis team."""
        return {
            "agents": [
                {
                    "name": "Market Analyst",
                    "role": "Market Research Specialist",
                    "goal": "Analyze market conditions and trends",
                    "backstory": "Expert market analyst with 10+ years experience",
                    "tools": ["market_data_api", "financial_news"],
                },
                {
                    "name": "Financial Analyst",
                    "role": "Financial Analysis Specialist",
                    "goal": "Analyze company financial statements",
                    "backstory": "CFA with deep financial analysis expertise",
                    "tools": ["financial_statements", "ratio_calculator"],
                },
                {
                    "name": "Risk Manager",
                    "role": "Risk Assessment Specialist",
                    "goal": "Assess investment risks and scenarios",
                    "backstory": "Risk management expert",
                    "tools": ["risk_models", "scenario_analysis"],
                },
                {
                    "name": "Report Writer",
                    "role": "Report Generation Specialist",
                    "goal": "Synthesize findings into comprehensive report",
                    "backstory": "Professional technical writer",
                    "tools": ["document_generator"],
                },
            ],
            "tasks": [
                {
                    "agent": "Market Analyst",
                    "description": "Analyze market conditions for {company}",
                },
                {
                    "agent": "Financial Analyst",
                    "description": "Analyze Q3 financial results for {company}",
                },
                {
                    "agent": "Risk Manager",
                    "description": "Assess risks and create scenarios for {company}",
                },
                {
                    "agent": "Report Writer",
                    "description": "Write comprehensive investment analysis report",
                },
            ],
            "process": "sequential",
        }

    @staticmethod
    def create_legal_team() -> Dict[str, Any]:
        """Create CrewAI legal team."""
        return {
            "agents": [
                {
                    "name": "Contract Analyzer",
                    "role": "Legal Contract Specialist",
                    "goal": "Analyze and review contract terms",
                    "backstory": "Senior contract attorney",
                },
                {
                    "name": "Precedent Researcher",
                    "role": "Legal Research Specialist",
                    "goal": "Research relevant case law and precedents",
                    "backstory": "Legal researcher specializing in case law",
                },
                {
                    "name": "Risk Assessor",
                    "role": "Legal Risk Specialist",
                    "goal": "Identify and assess legal risks",
                    "backstory": "Risk management expert in legal domain",
                },
                {
                    "name": "Document Drafter",
                    "role": "Legal Document Specialist",
                    "goal": "Draft legal documents and recommendations",
                    "backstory": "Legal document drafting expert",
                },
            ],
            "process": "sequential",
        }


@dataclass
class AutoGenTemplate:
    """Template for AutoGen implementation."""

    @staticmethod
    def create_group_chat_config() -> Dict[str, Any]:
        """Create AutoGen group chat configuration."""
        return {
            "agents": [
                {
                    "name": "analyst",
                    "system_message": "You are a financial analyst. Provide analysis based on data.",
                    "llm_config": {"model": "gpt-4", "temperature": 0.7},
                },
                {
                    "name": "researcher",
                    "system_message": "You are a market researcher. Research trends and competition.",
                    "llm_config": {"model": "gpt-4", "temperature": 0.7},
                },
                {
                    "name": "critic",
                    "system_message": "You are a critical evaluator. Challenge assumptions and findings.",
                    "llm_config": {"model": "gpt-4", "temperature": 0.7},
                },
            ],
            "group_chat_config": {
                "agents": ["analyst", "researcher", "critic"],
                "max_round": 10,
                "speaker_selection_method": "auto",
            },
        }

    @staticmethod
    def create_hierarchical_structure() -> Dict[str, Any]:
        """Create hierarchical AutoGen structure."""
        return {
            "primary_agents": [
                {
                    "name": "senior_analyst",
                    "role": "Senior analyst coordinating teams",
                    "subordinates": ["junior_analyst_1", "junior_analyst_2"],
                }
            ],
            "secondary_agents": [
                {
                    "name": "junior_analyst_1",
                    "role": "Fundamental analysis specialist",
                },
                {
                    "name": "junior_analyst_2",
                    "role": "Technical analysis specialist",
                },
            ],
        }


@dataclass
class LangGraphTemplate:
    """Template for LangGraph workflow implementation."""

    @staticmethod
    def create_research_workflow() -> Dict[str, Any]:
        """Create LangGraph research workflow."""
        return {
            "name": "research_workflow",
            "state_schema": {
                "topic": str,
                "research_findings": str,
                "analysis": str,
                "recommendations": str,
            },
            "nodes": [
                {
                    "name": "researcher",
                    "agent": "research_agent",
                    "description": "Research the topic",
                },
                {
                    "name": "analyst",
                    "agent": "analyst_agent",
                    "description": "Analyze research findings",
                },
                {
                    "name": "critic",
                    "agent": "critic_agent",
                    "description": "Critique and improve",
                },
                {
                    "name": "writer",
                    "agent": "writer_agent",
                    "description": "Write final report",
                },
            ],
            "edges": [
                ("researcher", "analyst"),
                ("analyst", "critic"),
                ("critic", "writer"),
            ],
            "entry_point": "researcher",
            "exit_point": "writer",
        }

    @staticmethod
    def create_dynamic_workflow() -> Dict[str, Any]:
        """Create dynamic LangGraph workflow with conditions."""
        return {
            "name": "dynamic_workflow",
            "nodes": [
                {
                    "name": "start",
                    "type": "input",
                },
                {
                    "name": "analyze",
                    "type": "agent",
                    "agent": "analyzer",
                },
                {
                    "name": "quality_check",
                    "type": "decision",
                    "condition": "analyze_quality > threshold",
                },
                {
                    "name": "refine",
                    "type": "agent",
                    "agent": "refiner",
                },
                {
                    "name": "end",
                    "type": "output",
                },
            ],
            "conditional_edges": [
                {
                    "source": "quality_check",
                    "true_target": "end",
                    "false_target": "refine",
                }
            ],
        }


@dataclass
class SwarmTemplate:
    """Template for OpenAI Swarm implementation."""

    @staticmethod
    def create_customer_support_swarm() -> Dict[str, Any]:
        """Create customer support Swarm."""
        return {
            "agents": [
                {
                    "name": "triage_agent",
                    "instructions": "Determine which specialist to route the customer to. "
                    "Ask clarifying questions if needed.",
                    "functions": [
                        "route_to_billing",
                        "route_to_technical",
                        "route_to_account",
                    ],
                },
                {
                    "name": "billing_specialist",
                    "instructions": "Handle all billing and payment related questions. "
                    "Can access billing records.",
                    "tools": ["billing_system", "payment_processor"],
                },
                {
                    "name": "technical_support",
                    "instructions": "Handle technical issues and troubleshooting. "
                    "Can access diagnostic tools.",
                    "tools": ["diagnostic_tools", "knowledge_base"],
                },
                {
                    "name": "account_specialist",
                    "instructions": "Handle account management and profile changes.",
                    "tools": ["account_system"],
                },
            ],
            "handoff_functions": {
                "route_to_billing": "Transfer to billing specialist",
                "route_to_technical": "Transfer to technical support",
                "route_to_account": "Transfer to account specialist",
            },
        }

    @staticmethod
    def create_sales_swarm() -> Dict[str, Any]:
        """Create sales team Swarm."""
        return {
            "agents": [
                {
                    "name": "sales_router",
                    "instructions": "Route customer inquiries to appropriate sales agent",
                    "functions": ["route_to_enterprise", "route_to_smb"],
                },
                {
                    "name": "enterprise_sales",
                    "instructions": "Handle enterprise customer inquiries",
                },
                {
                    "name": "smb_sales",
                    "instructions": "Handle small business inquiries",
                },
            ],
        }


class AgentCommunicationManager:
    """Manage communication patterns between agents."""

    def __init__(self, agents: Dict[str, Any]):
        """Initialize communication manager."""
        self.agents = agents
        self.message_queue = []

    def broadcast_message(self, sender: str, message: str, recipients: List[str]):
        """Broadcast message to multiple agents."""
        for recipient in recipients:
            self.message_queue.append({
                "from": sender,
                "to": recipient,
                "message": message,
                "type": "broadcast",
            })

    def direct_message(self, sender: str, recipient: str, message: str):
        """Send direct message between agents."""
        self.message_queue.append({
            "from": sender,
            "to": recipient,
            "message": message,
            "type": "direct",
        })

    def publish_shared_state(self, state_key: str, value: Any):
        """Publish to shared state (tool-mediated communication)."""
        self.message_queue.append({
            "type": "shared_state",
            "key": state_key,
            "value": value,
        })

    def get_pending_messages(self, agent: str) -> List[Dict]:
        """Get messages for specific agent."""
        return [msg for msg in self.message_queue if msg.get("to") == agent]

    def process_message_queue(self) -> Dict[str, Any]:
        """Process and summarize message queue."""
        direct_messages = [m for m in self.message_queue if m["type"] == "direct"]
        broadcasts = [m for m in self.message_queue if m["type"] == "broadcast"]
        shared_states = [m for m in self.message_queue if m["type"] == "shared_state"]

        return {
            "direct_messages": len(direct_messages),
            "broadcasts": len(broadcasts),
            "shared_states": len(shared_states),
            "total_messages": len(self.message_queue),
        }

"""
Multi-Agent Orchestration Patterns

Implements sequential, parallel, hierarchical, and consensus orchestration.
"""

from typing import List, Dict, Any, Optional
import asyncio
from abc import ABC, abstractmethod


class Agent(ABC):
    """Base agent class."""

    def __init__(self, name: str, role: str, goal: str):
        """Initialize agent."""
        self.name = name
        self.role = role
        self.goal = goal

    @abstractmethod
    def work(self, task: str) -> str:
        """Execute task."""
        pass

    async def work_async(self, task: str) -> str:
        """Execute task asynchronously."""
        return self.work(task)


class SimpleAgent(Agent):
    """Simple agent implementation."""

    def __init__(self, name: str, role: str, goal: str):
        """Initialize simple agent."""
        super().__init__(name, role, goal)

    def work(self, task: str) -> str:
        """Simulate agent work."""
        return f"{self.name}: Completed task - {task[:50]}..."


class SequentialOrchestrator:
    """Orchestrate agents to work sequentially."""

    def __init__(self, agents: List[Agent]):
        """
        Initialize orchestrator with agents.

        Args:
            agents: List of agents to orchestrate
        """
        self.agents = agents
        self.execution_log = []

    def execute(self, initial_task: str) -> Dict[str, Any]:
        """
        Execute agents sequentially.

        Args:
            initial_task: Initial task description

        Returns:
            Execution results dictionary
        """
        current_input = initial_task
        results = {}

        for agent in self.agents:
            result = agent.work(current_input)
            results[agent.name] = result
            self.execution_log.append({
                "agent": agent.name,
                "role": agent.role,
                "task": current_input,
                "result": result,
            })
            # Next agent uses this agent's output
            current_input = result

        return {
            "type": "sequential",
            "results": results,
            "final_output": current_input,
            "execution_log": self.execution_log,
        }


class ParallelOrchestrator:
    """Orchestrate agents to work in parallel."""

    def __init__(self, agents: List[Agent]):
        """Initialize parallel orchestrator."""
        self.agents = agents
        self.execution_log = []

    async def execute_async(self, task: str) -> Dict[str, Any]:
        """
        Execute agents in parallel.

        Args:
            task: Task description

        Returns:
            Execution results dictionary
        """
        # Create async tasks for all agents
        tasks = [agent.work_async(task) for agent in self.agents]

        # Execute all in parallel
        results_list = await asyncio.gather(*tasks)

        results = {
            agent.name: result
            for agent, result in zip(self.agents, results_list)
        }

        self.execution_log.append({
            "type": "parallel",
            "task": task,
            "agents": [a.name for a in self.agents],
            "results": results,
        })

        return {
            "type": "parallel",
            "results": results,
            "execution_log": self.execution_log,
        }

    def execute(self, task: str) -> Dict[str, Any]:
        """Execute agents synchronously (sequential fallback)."""
        return asyncio.run(self.execute_async(task))


class HierarchicalOrchestrator:
    """Orchestrate agents in hierarchical structure."""

    def __init__(self, manager_agent: Agent, specialist_teams: Dict[str, List[Agent]]):
        """
        Initialize hierarchical orchestrator.

        Args:
            manager_agent: Manager agent
            specialist_teams: Dict of team names to lists of agents
        """
        self.manager = manager_agent
        self.specialist_teams = specialist_teams
        self.execution_log = []

    def execute(self, main_task: str) -> Dict[str, Any]:
        """
        Execute with hierarchical structure.

        Args:
            main_task: Main task for manager

        Returns:
            Execution results
        """
        team_results = {}

        # Manager assigns tasks to teams
        for team_name, agents in self.specialist_teams.items():
            # Each team works on their aspect
            team_task = f"{main_task} - Team: {team_name}"
            team_result = {}

            for agent in agents:
                result = agent.work(team_task)
                team_result[agent.name] = result

            team_results[team_name] = team_result

            self.execution_log.append({
                "team": team_name,
                "agents": [a.name for a in agents],
                "results": team_result,
            })

        # Manager synthesizes results
        manager_result = self.manager.work(
            f"Synthesize findings from: {list(team_results.keys())}"
        )

        return {
            "type": "hierarchical",
            "team_results": team_results,
            "manager_synthesis": manager_result,
            "execution_log": self.execution_log,
        }


class ConsensusOrchestrator:
    """Orchestrate agent debate and consensus."""

    def __init__(self, agents: List[Agent], mediator_agent: Agent):
        """Initialize consensus orchestrator."""
        self.agents = agents
        self.mediator = mediator_agent
        self.debate_history = []

    def execute(self, question: str, rounds: int = 2) -> Dict[str, Any]:
        """
        Execute debate and reach consensus.

        Args:
            question: Question for debate
            rounds: Number of debate rounds

        Returns:
            Consensus results
        """
        # Round 1: Initial positions
        positions = {}
        for agent in self.agents:
            position = agent.work(f"Argue your position on: {question}")
            positions[agent.name] = position
            self.debate_history.append({
                "round": 1,
                "agent": agent.name,
                "position": position,
            })

        # Additional rounds: Response to others
        for round_num in range(2, rounds + 1):
            for agent in self.agents:
                other_positions = {
                    name: pos
                    for name, pos in positions.items()
                    if name != agent.name
                }
                response = agent.work(
                    f"Respond to these positions: {str(other_positions)}"
                )
                self.debate_history.append({
                    "round": round_num,
                    "agent": agent.name,
                    "response": response,
                })

        # Mediator reaches consensus
        consensus = self.mediator.work(
            f"Synthesize consensus from debate on: {question}"
        )

        return {
            "type": "consensus",
            "question": question,
            "initial_positions": positions,
            "debate_rounds": rounds,
            "consensus": consensus,
            "debate_history": self.debate_history,
        }


class AdaptiveOrchestrator:
    """Adapt orchestration based on progress."""

    def __init__(self, agents: List[Agent]):
        """Initialize adaptive orchestrator."""
        self.agents = agents
        self.execution_log = []

    def execute_with_adaptation(
        self,
        initial_task: str,
        progress_threshold: float = 0.7,
        quality_threshold: float = 0.6,
    ) -> Dict[str, Any]:
        """
        Execute with adaptive workflow changes.

        Args:
            initial_task: Initial task
            progress_threshold: Threshold for adding resources
            quality_threshold: Threshold for adding validation

        Returns:
            Adaptive execution results
        """
        results = {}
        current_task = initial_task
        active_agents = self.agents.copy()

        for iteration in range(3):
            # Execute with current agents
            iteration_results = {}
            for agent in active_agents:
                result = agent.work(current_task)
                iteration_results[agent.name] = result

            results[f"iteration_{iteration}"] = iteration_results

            # Assess progress
            progress = self._calculate_progress(iteration_results)
            quality = self._calculate_quality(iteration_results)

            log_entry = {
                "iteration": iteration,
                "agents": [a.name for a in active_agents],
                "progress": progress,
                "quality": quality,
            }

            # Adapt based on progress
            if progress < progress_threshold and iteration < 2:
                # Add more agents
                log_entry["adaptation"] = "Added specialist agent"
                self.execution_log.append(log_entry)
            elif quality < quality_threshold and iteration < 2:
                # Add validation step
                log_entry["adaptation"] = "Added validation agent"
                self.execution_log.append(log_entry)
            else:
                self.execution_log.append(log_entry)

        return {
            "type": "adaptive",
            "results": results,
            "adaptations": self.execution_log,
        }

    @staticmethod
    def _calculate_progress(results: Dict) -> float:
        """Calculate progress score."""
        # Simple heuristic based on results
        return min(len(results) / 3.0, 1.0)

    @staticmethod
    def _calculate_quality(results: Dict) -> float:
        """Calculate quality score."""
        # Simple heuristic
        return 0.7  # Placeholder


class WorkflowGraph:
    """Define workflow as directed acyclic graph (DAG)."""

    def __init__(self):
        """Initialize workflow graph."""
        self.nodes = {}
        self.edges = []

    def add_node(self, node_id: str, agent: Agent) -> None:
        """Add agent node."""
        self.nodes[node_id] = agent

    def add_edge(self, from_node: str, to_node: str) -> None:
        """Add dependency edge."""
        self.edges.append((from_node, to_node))

    def execute_dag(self, initial_task: str) -> Dict[str, Any]:
        """
        Execute workflow as DAG.

        Args:
            initial_task: Initial task

        Returns:
            Execution results
        """
        results = {}
        ready_nodes = self._find_ready_nodes()

        while ready_nodes:
            for node_id in ready_nodes:
                agent = self.nodes[node_id]
                # Get input from dependencies
                task_input = self._get_node_input(node_id, results, initial_task)
                result = agent.work(task_input)
                results[node_id] = result

            ready_nodes = self._find_ready_nodes(completed=set(results.keys()))

        return {
            "type": "dag",
            "results": results,
            "nodes": list(self.nodes.keys()),
            "edges": self.edges,
        }

    def _find_ready_nodes(self, completed: Optional[set] = None) -> List[str]:
        """Find nodes with all dependencies completed."""
        if completed is None:
            completed = set()

        ready = []
        for node_id in self.nodes:
            if node_id not in completed:
                dependencies = [
                    from_node
                    for from_node, to_node in self.edges
                    if to_node == node_id
                ]
                if all(dep in completed for dep in dependencies):
                    ready.append(node_id)

        return ready

    def _get_node_input(
        self, node_id: str, results: Dict, initial_task: str
    ) -> str:
        """Get input for node from dependencies."""
        dependencies = [
            from_node for from_node, to_node in self.edges if to_node == node_id
        ]
        if not dependencies:
            return initial_task
        # Combine outputs from dependencies
        return " + ".join(results.get(dep, "") for dep in dependencies)

"""
Agent Communication Management

Handle agent-to-agent communication, message passing, and shared state.
"""

from typing import Dict, List, Optional, Any, Callable
from dataclasses import dataclass, field
from enum import Enum
import json


class MessageType(Enum):
    """Types of messages between agents."""

    DIRECT = "direct"
    BROADCAST = "broadcast"
    REQUEST = "request"
    RESPONSE = "response"
    FEEDBACK = "feedback"
    ERROR = "error"


@dataclass
class Message:
    """Message between agents."""

    sender: str
    recipient: str
    content: str
    message_type: MessageType
    timestamp: float = field(default_factory=lambda: 0)
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict:
        """Convert to dictionary."""
        return {
            "sender": self.sender,
            "recipient": self.recipient,
            "content": self.content,
            "type": self.message_type.value,
            "timestamp": self.timestamp,
            "metadata": self.metadata,
        }


class MessageBroker:
    """Central message broker for agent communication."""

    def __init__(self):
        """Initialize message broker."""
        self.message_queue: List[Message] = []
        self.agent_inboxes: Dict[str, List[Message]] = {}
        self.message_handlers: Dict[MessageType, List[Callable]] = {}

    def send_message(self, message: Message) -> bool:
        """
        Send message from one agent to another.

        Args:
            message: Message object

        Returns:
            Whether message was delivered
        """
        self.message_queue.append(message)

        # Add to recipient's inbox
        if message.recipient not in self.agent_inboxes:
            self.agent_inboxes[message.recipient] = []
        self.agent_inboxes[message.recipient].append(message)

        # Trigger handlers
        self._trigger_handlers(message)

        return True

    def broadcast_message(self, sender: str, content: str, recipients: List[str]):
        """
        Broadcast message to multiple agents.

        Args:
            sender: Sending agent
            content: Message content
            recipients: List of recipient agents
        """
        for recipient in recipients:
            message = Message(
                sender=sender,
                recipient=recipient,
                content=content,
                message_type=MessageType.BROADCAST,
            )
            self.send_message(message)

    def request_response(
        self, sender: str, recipient: str, content: str, timeout: float = 5.0
    ) -> Optional[Message]:
        """
        Send request and wait for response.

        Args:
            sender: Requesting agent
            recipient: Agent to respond
            content: Request content
            timeout: Response timeout in seconds

        Returns:
            Response message or None
        """
        message = Message(
            sender=sender,
            recipient=recipient,
            content=content,
            message_type=MessageType.REQUEST,
            metadata={"timeout": timeout},
        )
        self.send_message(message)

        # Placeholder - wait for response
        # In production, implement actual timeout/wait mechanism
        return None

    def get_inbox(self, agent: str) -> List[Message]:
        """Get messages for specific agent."""
        return self.agent_inboxes.get(agent, [])

    def clear_inbox(self, agent: str) -> None:
        """Clear agent's inbox."""
        self.agent_inboxes[agent] = []

    def register_handler(
        self, message_type: MessageType, handler: Callable
    ) -> None:
        """Register handler for message type."""
        if message_type not in self.message_handlers:
            self.message_handlers[message_type] = []
        self.message_handlers[message_type].append(handler)

    def _trigger_handlers(self, message: Message) -> None:
        """Trigger registered handlers."""
        handlers = self.message_handlers.get(message.message_type, [])
        for handler in handlers:
            try:
                handler(message)
            except Exception:
                pass

    def get_statistics(self) -> Dict[str, Any]:
        """Get communication statistics."""
        type_counts = {}
        for message in self.message_queue:
            msg_type = message.message_type.value
            type_counts[msg_type] = type_counts.get(msg_type, 0) + 1

        return {
            "total_messages": len(self.message_queue),
            "by_type": type_counts,
            "total_agents": len(self.agent_inboxes),
            "agents": list(self.agent_inboxes.keys()),
        }


class SharedMemory:
    """Shared memory for tool-mediated agent communication."""

    def __init__(self):
        """Initialize shared memory."""
        self.memory: Dict[str, Any] = {}
        self.access_log: List[Dict] = []

    def write(self, key: str, value: Any, agent: str = "system") -> None:
        """
        Write to shared memory.

        Args:
            key: Memory key
            value: Value to store
            agent: Agent writing
        """
        self.memory[key] = {
            "value": value,
            "writer": agent,
            "access_count": 0,
        }
        self.access_log.append({
            "action": "write",
            "key": key,
            "agent": agent,
        })

    def read(self, key: str, agent: str = "system") -> Optional[Any]:
        """
        Read from shared memory.

        Args:
            key: Memory key
            agent: Agent reading

        Returns:
            Stored value or None
        """
        if key in self.memory:
            entry = self.memory[key]
            entry["access_count"] += 1
            self.access_log.append({
                "action": "read",
                "key": key,
                "agent": agent,
            })
            return entry["value"]
        return None

    def append(self, key: str, value: Any, agent: str = "system") -> None:
        """Append to list in shared memory."""
        if key not in self.memory:
            self.memory[key] = {
                "value": [],
                "writer": agent,
                "access_count": 0,
            }
        if isinstance(self.memory[key]["value"], list):
            self.memory[key]["value"].append(value)

    def get_all(self) -> Dict[str, Any]:
        """Get all shared memory contents."""
        return {
            key: entry["value"] for key, entry in self.memory.items()
        }

    def get_statistics(self) -> Dict[str, Any]:
        """Get memory access statistics."""
        total_accesses = sum(
            entry["access_count"] for entry in self.memory.values()
        )

        return {
            "total_keys": len(self.memory),
            "total_accesses": total_accesses,
            "access_log_size": len(self.access_log),
        }


class ContextManager:
    """Manage context sharing between agents."""

    def __init__(self):
        """Initialize context manager."""
        self.contexts: Dict[str, Dict[str, Any]] = {}
        self.global_context: Dict[str, Any] = {}

    def create_context(self, context_id: str, initial_data: Optional[Dict] = None) -> None:
        """Create new context."""
        self.contexts[context_id] = initial_data or {}

    def update_context(self, context_id: str, data: Dict) -> None:
        """Update context data."""
        if context_id in self.contexts:
            self.contexts[context_id].update(data)

    def get_context(self, context_id: str) -> Dict[str, Any]:
        """Get context data."""
        return self.contexts.get(context_id, {})

    def set_global_context(self, key: str, value: Any) -> None:
        """Set global context variable."""
        self.global_context[key] = value

    def get_global_context(self, key: str) -> Optional[Any]:
        """Get global context variable."""
        return self.global_context.get(key)

    def context_to_string(self, context_id: str) -> str:
        """Convert context to formatted string."""
        context = self.get_context(context_id)
        return json.dumps(context, indent=2)


class CommunicationProtocol:
    """Define communication protocol between agents."""

    def __init__(self, broker: MessageBroker, shared_memory: SharedMemory):
        """Initialize protocol."""
        self.broker = broker
        self.shared_memory = shared_memory

    def request_analysis(
        self, requester: str, analyzer: str, subject: str
    ) -> None:
        """Request analysis from another agent."""
        message = Message(
            sender=requester,
            recipient=analyzer,
            content=f"Analyze: {subject}",
            message_type=MessageType.REQUEST,
            metadata={"action": "analyze"},
        )
        self.broker.send_message(message)

    def share_findings(self, source_agent: str, key: str, findings: Dict) -> None:
        """Share findings through shared memory."""
        self.shared_memory.write(key, findings, source_agent)

    def aggregate_results(
        self, agents: List[str], result_key: str
    ) -> Dict[str, Any]:
        """Aggregate results from multiple agents."""
        results = {}
        for agent in agents:
            agent_results = self.shared_memory.read(f"{agent}_{result_key}")
            if agent_results:
                results[agent] = agent_results

        return results

    def notify_status(self, agent: str, status: str) -> None:
        """Notify status update."""
        message = Message(
            sender=agent,
            recipient="orchestrator",
            content=f"Status: {status}",
            message_type=MessageType.FEEDBACK,
        )
        self.broker.send_message(message)

    def error_propagation(self, source_agent: str, error: str) -> None:
        """Propagate error to relevant agents."""
        message = Message(
            sender=source_agent,
            recipient="orchestrator",
            content=f"Error: {error}",
            message_type=MessageType.ERROR,
        )
        self.broker.send_message(message)

"""
Multi-Agent System Benchmarking

Evaluate team performance and agent effectiveness.
"""

from typing import Dict, List, Any, Callable, Optional
from dataclasses import dataclass
import time
import statistics


@dataclass
class BenchmarkResult:
    """Result from benchmark run."""

    test_name: str
    total_time: float
    task_count: int
    success_count: int
    failure_count: int
    quality_score: float
    cost_tokens: int


class TeamBenchmark:
    """Benchmark multi-agent team performance."""

    def __init__(self):
        """Initialize benchmarking suite."""
        self.results: List[BenchmarkResult] = []
        self.agent_metrics: Dict[str, Dict] = {}

    def run_benchmark(
        self,
        test_name: str,
        orchestrator,
        test_data: List[Dict],
        quality_metric: Optional[Callable] = None,
    ) -> BenchmarkResult:
        """
        Run benchmark test.

        Args:
            test_name: Name of benchmark
            orchestrator: Orchestrator instance
            test_data: Test cases
            quality_metric: Function to evaluate quality

        Returns:
            Benchmark result
        """
        start_time = time.time()
        results = []
        costs = []

        for test_case in test_data:
            try:
                result = orchestrator.execute(test_case)
                results.append(result)
                # Assume cost in test_case
                costs.append(test_case.get("cost", 0))
            except Exception:
                pass

        end_time = time.time()

        success_count = len(results)
        failure_count = len(test_data) - success_count
        total_time = end_time - start_time

        # Calculate quality score
        quality_score = (
            quality_metric(results) if quality_metric else success_count / len(test_data)
        )

        benchmark_result = BenchmarkResult(
            test_name=test_name,
            total_time=total_time,
            task_count=len(test_data),
            success_count=success_count,
            failure_count=failure_count,
            quality_score=quality_score,
            cost_tokens=sum(costs),
        )

        self.results.append(benchmark_result)
        return benchmark_result

    def compare_orchestration_methods(
        self,
        methods: Dict[str, Callable],
        test_data: List[Dict],
    ) -> Dict[str, Any]:
        """
        Compare different orchestration methods.

        Args:
            methods: Dict of method name to orchestrator
            test_data: Test cases

        Returns:
            Comparison results
        """
        comparison = {}

        for method_name, orchestrator in methods.items():
            result = self.run_benchmark(method_name, orchestrator, test_data)
            comparison[method_name] = {
                "total_time": result.total_time,
                "success_rate": result.success_count / result.task_count,
                "quality_score": result.quality_score,
                "efficiency": result.success_count / (result.total_time + 0.001),
                "cost_per_success": (
                    result.cost_tokens / result.success_count
                    if result.success_count > 0
                    else float("inf")
                ),
            }

        return comparison

    def get_summary(self) -> Dict[str, Any]:
        """Get benchmark summary."""
        if not self.results:
            return {"status": "no_results"}

        times = [r.total_time for r in self.results]
        success_rates = [
            r.success_count / r.task_count for r in self.results
        ]
        quality_scores = [r.quality_score for r in self.results]

        return {
            "total_benchmarks": len(self.results),
            "avg_time": statistics.mean(times),
            "median_time": statistics.median(times),
            "avg_success_rate": statistics.mean(success_rates),
            "avg_quality": statistics.mean(quality_scores),
            "benchmarks": [r.__dict__ for r in self.results],
        }


class AgentEffectiveness:
    """Measure individual agent effectiveness."""

    def __init__(self):
        """Initialize effectiveness tracker."""
        self.agent_stats: Dict[str, Dict] = {}

    def record_agent_task(
        self,
        agent: str,
        task_id: str,
        success: bool,
        quality_score: float,
        duration: float,
    ) -> None:
        """Record agent task execution."""
        if agent not in self.agent_stats:
            self.agent_stats[agent] = {
                "tasks": [],
                "successes": 0,
                "failures": 0,
            }

        stats = self.agent_stats[agent]
        stats["tasks"].append({
            "task_id": task_id,
            "success": success,
            "quality": quality_score,
            "duration": duration,
        })

        if success:
            stats["successes"] += 1
        else:
            stats["failures"] += 1

    def get_agent_metrics(self, agent: str) -> Dict[str, Any]:
        """Get metrics for specific agent."""
        if agent not in self.agent_stats:
            return {"status": "no_data"}

        stats = self.agent_stats[agent]
        tasks = stats["tasks"]

        if not tasks:
            return {"status": "no_tasks"}

        success_rate = stats["successes"] / (stats["successes"] + stats["failures"])
        quality_scores = [t["quality"] for t in tasks]
        durations = [t["duration"] for t in tasks]

        return {
            "agent": agent,
            "total_tasks": len(tasks),
            "success_rate": success_rate,
            "avg_quality": statistics.mean(quality_scores),
            "avg_duration": statistics.mean(durations),
            "reliability": success_rate,  # Alias for clarity
        }

    def rank_agents(self) -> List[tuple]:
        """Rank agents by effectiveness."""
        rankings = []

        for agent in self.agent_stats.keys():
            metrics = self.get_agent_metrics(agent)
            if "success_rate" in metrics:
                score = (
                    metrics["success_rate"] * 0.4 +
                    metrics["avg_quality"] * 0.4 +
                    (1 - min(metrics["avg_duration"] / 10.0, 1.0)) * 0.2
                )
                rankings.append((agent, score, metrics))

        rankings.sort(key=lambda x: x[1], reverse=True)
        return rankings

    def get_team_report(self) -> Dict[str, Any]:
        """Get team effectiveness report."""
        rankings = self.rank_agents()

        return {
            "total_agents": len(self.agent_stats),
            "rankings": [
                {
                    "rank": i + 1,
                    "agent": agent,
                    "score": score,
                    "metrics": metrics,
                }
                for i, (agent, score, metrics) in enumerate(rankings)
            ],
        }


class CollaborationMetrics:
    """Measure collaboration effectiveness between agents."""

    def __init__(self):
        """Initialize collaboration metrics."""
        self.interactions: List[Dict] = []

    def record_interaction(
        self,
        agent_a: str,
        agent_b: str,
        message: str,
        response_time: float,
        successful: bool,
    ) -> None:
        """Record agent-to-agent interaction."""
        self.interactions.append({
            "agent_a": agent_a,
            "agent_b": agent_b,
            "message": message,
            "response_time": response_time,
            "successful": successful,
        })

    def get_collaboration_graph(self) -> Dict[str, List[str]]:
        """Get collaboration graph between agents."""
        graph = {}

        for interaction in self.interactions:
            agent_a = interaction["agent_a"]
            agent_b = interaction["agent_b"]

            if agent_a not in graph:
                graph[agent_a] = []
            if agent_b not in graph[agent_a]:
                graph[agent_a].append(agent_b)

        return graph

    def get_interaction_metrics(self) -> Dict[str, Any]:
        """Get interaction metrics."""
        if not self.interactions:
            return {"total_interactions": 0}

        response_times = [i["response_time"] for i in self.interactions]
        success_rate = sum(1 for i in self.interactions if i["successful"]) / len(
            self.interactions
        )

        return {
            "total_interactions": len(self.interactions),
            "success_rate": success_rate,
            "avg_response_time": statistics.mean(response_times),
            "median_response_time": statistics.median(response_times),
            "max_response_time": max(response_times),
            "collaboration_score": success_rate * (1 - min(statistics.mean(response_times) / 10.0, 1.0)),
        }

    def get_agent_collaboration_analysis(self, agent: str) -> Dict[str, Any]:
        """Get collaboration analysis for specific agent."""
        agent_interactions = [
            i for i in self.interactions
            if i["agent_a"] == agent or i["agent_b"] == agent
        ]

        if not agent_interactions:
            return {"agent": agent, "status": "no_interactions"}

        partners = set()
        for interaction in agent_interactions:
            if interaction["agent_a"] == agent:
                partners.add(interaction["agent_b"])
            else:
                partners.add(interaction["agent_a"])

        success_rate = sum(
            1 for i in agent_interactions if i["successful"]
        ) / len(agent_interactions)

        return {
            "agent": agent,
            "collaboration_partners": list(partners),
            "total_interactions": len(agent_interactions),
            "collaboration_success_rate": success_rate,
        }


def create_benchmark_suite() -> Dict[str, Callable]:
    """Create standard benchmark suite."""
    return {
        "sequential_tasks": lambda orchestrator: orchestrator.execute(
            {"type": "sequential", "task_count": 5}
        ),
        "parallel_tasks": lambda orchestrator: orchestrator.execute(
            {"type": "parallel", "task_count": 10}
        ),
        "complex_workflow": lambda orchestrator: orchestrator.execute(
            {"type": "complex", "task_count": 20}
        ),
        "error_handling": lambda orchestrator: orchestrator.execute(
            {"type": "with_errors", "task_count": 10}
        ),
    }

"""
Workflow Management for Multi-Agent Systems

Handle workflow execution, monitoring, and optimization.
"""

from typing import Dict, List, Any, Optional, Callable
from enum import Enum
from dataclasses import dataclass, field
import time


class WorkflowStatus(Enum):
    """Workflow execution status."""

    PENDING = "pending"
    RUNNING = "running"
    PAUSED = "paused"
    COMPLETED = "completed"
    FAILED = "failed"


class TaskStatus(Enum):
    """Task execution status."""

    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    SKIPPED = "skipped"


@dataclass
class Task:
    """Task to be executed by an agent."""

    task_id: str
    agent: str
    description: str
    status: TaskStatus = TaskStatus.PENDING
    result: Optional[str] = None
    error: Optional[str] = None
    start_time: float = field(default_factory=time.time)
    end_time: Optional[float] = None
    dependencies: List[str] = field(default_factory=list)


@dataclass
class Workflow:
    """Workflow orchestration."""

    workflow_id: str
    name: str
    tasks: Dict[str, Task] = field(default_factory=dict)
    status: WorkflowStatus = WorkflowStatus.PENDING
    start_time: Optional[float] = None
    end_time: Optional[float] = None


class WorkflowExecutor:
    """Execute workflows with multiple agents."""

    def __init__(self):
        """Initialize workflow executor."""
        self.workflows: Dict[str, Workflow] = {}
        self.execution_history: List[Dict] = []

    def create_workflow(self, workflow_id: str, name: str) -> Workflow:
        """Create new workflow."""
        workflow = Workflow(workflow_id=workflow_id, name=name)
        self.workflows[workflow_id] = workflow
        return workflow

    def add_task(
        self,
        workflow_id: str,
        task_id: str,
        agent: str,
        description: str,
        dependencies: Optional[List[str]] = None,
    ) -> Task:
        """Add task to workflow."""
        task = Task(
            task_id=task_id,
            agent=agent,
            description=description,
            dependencies=dependencies or [],
        )
        self.workflows[workflow_id].tasks[task_id] = task
        return task

    def execute_workflow(
        self, workflow_id: str, executor_func: Callable
    ) -> Dict[str, Any]:
        """
        Execute workflow.

        Args:
            workflow_id: Workflow ID
            executor_func: Function to execute tasks

        Returns:
            Execution results
        """
        workflow = self.workflows[workflow_id]
        workflow.status = WorkflowStatus.RUNNING
        workflow.start_time = time.time()

        executed_tasks = set()
        results = {}

        while len(executed_tasks) < len(workflow.tasks):
            # Find ready tasks
            ready_tasks = self._get_ready_tasks(workflow, executed_tasks)

            if not ready_tasks:
                # Check if workflow is stuck
                if len(executed_tasks) > 0:
                    break

            for task_id in ready_tasks:
                task = workflow.tasks[task_id]
                task.status = TaskStatus.RUNNING

                try:
                    # Execute task
                    result = executor_func(task)
                    task.result = result
                    task.status = TaskStatus.COMPLETED
                    results[task_id] = result
                except Exception as e:
                    task.error = str(e)
                    task.status = TaskStatus.FAILED
                    results[task_id] = None

                task.end_time = time.time()
                executed_tasks.add(task_id)

        # Finalize workflow
        workflow.end_time = time.time()
        if all(task.status == TaskStatus.COMPLETED for task in workflow.tasks.values()):
            workflow.status = WorkflowStatus.COMPLETED
        else:
            workflow.status = WorkflowStatus.FAILED

        execution_record = {
            "workflow_id": workflow_id,
            "status": workflow.status.value,
            "duration": workflow.end_time - workflow.start_time,
            "results": results,
        }
        self.execution_history.append(execution_record)

        return results

    def _get_ready_tasks(
        self, workflow: Workflow, executed: set
    ) -> List[str]:
        """Get tasks ready to execute."""
        ready = []
        for task_id, task in workflow.tasks.items():
            if task_id not in executed and task.status == TaskStatus.PENDING:
                # Check if all dependencies are complete
                deps_met = all(dep in executed for dep in task.dependencies)
                if deps_met:
                    ready.append(task_id)

        return ready

    def get_workflow_status(self, workflow_id: str) -> Dict[str, Any]:
        """Get workflow status."""
        workflow = self.workflows[workflow_id]
        return {
            "id": workflow.workflow_id,
            "name": workflow.name,
            "status": workflow.status.value,
            "tasks": {
                task_id: task.status.value
                for task_id, task in workflow.tasks.items()
            },
        }


class WorkflowOptimizer:
    """Optimize workflow execution."""

    @staticmethod
    def analyze_dependencies(workflow: Workflow) -> Dict[str, Any]:
        """Analyze task dependencies."""
        dep_graph = {}
        for task_id, task in workflow.tasks.items():
            dep_graph[task_id] = task.dependencies

        # Find critical path
        critical_path = WorkflowOptimizer._find_critical_path(
            dep_graph, workflow.tasks
        )

        # Find parallelizable tasks
        parallel_groups = WorkflowOptimizer._find_parallel_groups(dep_graph)

        return {
            "dependency_graph": dep_graph,
            "critical_path": critical_path,
            "parallelizable_groups": parallel_groups,
        }

    @staticmethod
    def _find_critical_path(dep_graph: Dict, tasks: Dict) -> List[str]:
        """Find tasks on critical path."""
        # Simple implementation - actual critical path is more complex
        return list(dep_graph.keys())

    @staticmethod
    def _find_parallel_groups(dep_graph: Dict) -> List[List[str]]:
        """Find groups of tasks that can execute in parallel."""
        groups = []
        remaining = set(dep_graph.keys())

        while remaining:
            # Find tasks with no dependencies in remaining set
            independent = [
                task
                for task in remaining
                if not any(dep in remaining for dep in dep_graph.get(task, []))
            ]
            if independent:
                groups.append(independent)
                remaining -= set(independent)
            else:
                break

        return groups

    @staticmethod
    def estimate_execution_time(
        workflow: Workflow, task_times: Dict[str, float]
    ) -> float:
        """Estimate total execution time."""
        parallel_groups = WorkflowOptimizer._find_parallel_groups({
            task_id: task.dependencies
            for task_id, task in workflow.tasks.items()
        })

        total_time = 0
        for group in parallel_groups:
            group_time = max(
                task_times.get(task_id, 1.0) for task_id in group
            )
            total_time += group_time

        return total_time

    @staticmethod
    def suggest_parallelization(workflow: Workflow) -> List[Dict[str, Any]]:
        """Suggest ways to parallelize workflow."""
        suggestions = []

        dep_graph = {
            task_id: task.dependencies
            for task_id, task in workflow.tasks.items()
        }

        # Find sequential chains that could be parallelized
        for task_id, dependencies in dep_graph.items():
            if len(dependencies) == 0:
                suggestions.append({
                    "task": task_id,
                    "suggestion": "Can be parallelized with other independent tasks",
                })

        return suggestions


class WorkflowMonitor:
    """Monitor workflow execution."""

    def __init__(self):
        """Initialize monitor."""
        self.events: List[Dict] = []

    def record_event(
        self, workflow_id: str, task_id: str, event_type: str, data: Dict
    ) -> None:
        """Record workflow event."""
        event = {
            "workflow_id": workflow_id,
            "task_id": task_id,
            "event_type": event_type,
            "timestamp": time.time(),
            "data": data,
        }
        self.events.append(event)

    def get_workflow_metrics(self, workflow_id: str) -> Dict[str, Any]:
        """Get workflow metrics."""
        workflow_events = [e for e in self.events if e["workflow_id"] == workflow_id]

        task_times = {}
        for event in workflow_events:
            task_id = event["task_id"]
            if event["event_type"] == "start":
                task_times[task_id] = {"start": event["timestamp"]}
            elif event["event_type"] == "complete":
                if task_id in task_times:
                    task_times[task_id]["end"] = event["timestamp"]

        # Calculate durations
        durations = {
            task_id: times.get("end", time.time()) - times["start"]
            for task_id, times in task_times.items()
        }

        return {
            "total_tasks": len(task_times),
            "task_durations": durations,
            "avg_duration": sum(durations.values()) / len(durations)
            if durations
            else 0,
            "max_duration": max(durations.values()) if durations else 0,
        }

    def get_performance_report(self) -> Dict[str, Any]:
        """Get overall performance report."""
        if not self.events:
            return {"status": "no_events"}

        workflow_ids = set(e["workflow_id"] for e in self.events)

        report = {}
        for workflow_id in workflow_ids:
            report[workflow_id] = self.get_workflow_metrics(workflow_id)

        return report