Multi-Agent Systems: Collaborative AI Agent Architecture

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

class Agent(ABC):
    """
    Base class for an AI agent in a multi-agent system.
    """
    
    def __init__(self, agent_id: str, role: str):
        self.agent_id = agent_id
        self.role = role
        self.beliefs = {}  # Agent's knowledge
        self.intentions = []  # Goals to achieve
        self.plans = []  # Planned actions
        
    @abstractmethod
    def perceive(self, observation: Dict) -> None:
        """Process incoming observation."""
        pass
    
    @abstractmethod
    def reason(self) -> None:
        """Update beliefs and select intentions."""
        pass
    
    @abstractmethod
    def act(self) -> Any:
        """Execute planned actions."""
        pass
    
    @abstractmethod
    def communicate(self, message: Dict, recipients: List['Agent']) -> None:
        """Send message to other agents."""
        pass


class Environment:
    """
    Shared environment for multi-agent interaction.
    """
    
    def __init__(self):
        self.state = {}
        self.agents = {}
        
    def add_agent(self, agent: Agent):
        """Register agent in environment."""
        self.agents[agent.agent_id] = agent
        
    def step(self) -> Dict:
        """Execute one time step."""
        # All agents perceive
        for agent in self.agents.values():
            agent.perceive(self.state)
            
        # All agents reason
        for agent in self.agents.values():
            agent.reason()
            
        # All agents act
        for agent in self.agents.values():
            action = agent.act()
            self._apply_action(action)
            
        return self.state
    
    def _apply_action(self, action: Dict):
        """Apply agent action to environment."""
        # Update state based on action
        pass

class ReactiveAgent(Agent):
    """
    Simple reactive agent: stimulus -> action mapping.
    """
    
    def __init__(self, agent_id: str, role: str, rules: Dict):
        super().__init__(agent_id, role)
        self.rules = rules  # condition -> action mappings
        
    def perceive(self, observation: Dict) -> None:
        self.beliefs = observation
        
    def reason(self) -> None:
        # Match conditions to actions
        for condition, action in self.rules.items():
            if self._check_condition(condition):
                self.plans = [action]
                return
                
    def _check_condition(self, condition: str) -> bool:
        """Check if condition is satisfied."""
        return eval(condition, {}, self.beliefs)
    
    def act(self) -> Any:
        if self.plans:
            return self.plans[0]
        return None

class DeliberativeAgent(Agent):
    """
    Deliberative agent with planning capabilities.
    """
    
    def __init__(self, agent_id: str, role: str, planner):
        super().__init__(agent_id, role)
        self.planner = planner
        
    def perceive(self, observation: Dict) -> None:
        self.beliefs.update(observation)
        
    def reason(self) -> None:
        # Update goals based on beliefs
        self._update_intentions()
        
        # If we have intentions but no plans, create plans
        if self.intentions and not self.plans:
            goal = self.intentions[0]
            self.plans = self.planner.create_plan(goal, self.beliefs)
            
    def _update_intentions(self):
        """Update goals based on current state."""
        # Simplified: maintain goals until achieved
        pass
        
    def act(self) -> Any:
        if self.plans:
            action = self.plans[0]
            self.plans = self.plans[1:]  # Remove executed action
            return action
        return None

class HybridAgent(Agent):
    """
    Hybrid agent combining reactive and deliberative layers.
    """
    
    def __init__(self, agent_id: str, role: str, reactive_rules: Dict, planner):
        super().__init__(agent_id, role)
        self.reactive_rules = reactive_rules
        self.planner = planner
        self.deliberative_layer = DeliberativeAgent(agent_id, role, planner)
        
    def perceive(self, observation: Dict) -> None:
        self.beliefs.update(observation)
        
    def reason(self) -> None:
        # First check reactive rules
        for condition, action in self.reactive_rules.items():
            if self._check_condition(condition):
                self.plans = [action]
                return
                
        # If no reactive rule applies, use deliberative reasoning
        self.deliberative_layer.beliefs = self.beliefs
        self.deliberative_layer.reason()
        self.plans = self.deliberative_layer.plans
        self.intentions = self.deliberative_layer.intentions
        
    def _check_condition(self, condition: str) -> bool:
        return eval(condition, {}, self.beliefs)
        
    def act(self) -> Any:
        if self.plans:
            action = self.plans[0]
            self.plans = self.plans[1:]
            return action
        return None

class Message:
    """
    Agent communication message.
    """
    
    def __init__(self, sender: str, receiver: str, 
                 performative: str, content: Dict):
        self.sender = sender
        self.receiver = receiver
        self.performative = performative  # request, inform, query, etc.
        self.content = content
        self.timestamp = asyncio.get_event_loop().time()


class AgentCommunication:
    """
    Communication infrastructure for agents.
    """
    
    def __init__(self):
        self.message_queue = asyncio.Queue()
        self.agents = {}
        
    def register_agent(self, agent_id: str, agent: Agent):
        self.agents[agent_id] = agent
        
    async def send(self, message: Message):
        """Send message to agent."""
        await self.message_queue.put(message)
        
    async def broadcast(self, sender: str, message: Message):
        """Broadcast to all agents except sender."""
        for agent_id, agent in self.agents.items():
            if agent_id != sender:
                await self.send(message)
                
    async def receive_loop(self):
        """Process incoming messages."""
        while True:
            message = await self.message_queue.get()
            recipient = self.agents.get(message.receiver)
            if recipient:
                recipient.communicate(message.content, [])

class FIPAMessage:
    """
    FIPA Agent Communication Language messages.
    """
    
    # Message performatives
    INFORM = "inform"
    REQUEST = "request"
    QUERY = "query"
    PROPAGATE = "propagate"
    SUBSCRIBE = "subscribe"
    
    @staticmethod
    def create_inform(sender: str, receiver: str, content: Dict) -> Message:
        return Message(sender, receiver, FIPAMessage.INFORM, content)
    
    @staticmethod
    def create_request(sender: str, receiver: str, action: str, 
                      params: Dict) -> Message:
        content = {"action": action, "parameters": params}
        return Message(sender, receiver, FIPAMessage.REQUEST, content)
    
    @staticmethod
    def create_query(sender: str, receiver: str, query: str) -> Message:
        content = {"query": query}
        return Message(sender, receiver, FIPAMessage.QUERY, content)

class TaskAllocator:
    """
    Distribute tasks among agents based on capabilities.
    """
    
    def __init__(self, agents: Dict[str, Agent]):
        self.agents = agents
        self.capabilities = {}  # agent_id -> set of capabilities
        
    def register_capability(self, agent_id: str, capability: str):
        if agent_id not in self.capabilities:
            self.capabilities[agent_id] = set()
        self.capabilities[agent_id].add(capability)
        
    def allocate_task(self, task: Dict) -> List[Agent]:
        """
        Allocate task to suitable agents.
        
        Returns list of agents that can handle the task.
        """
        required = task.get("required_capabilities", [])
        
        suitable_agents = []
        for agent_id, caps in self.capabilities.items():
            if all(req in caps for req in required):
                suitable_agents.append(self.agents[agent_id])
                
        return suitable_agents
    
    def greedy_allocate(self, tasks: List[Dict]) -> Dict[Agent, List[Dict]]:
        """
        Greedy task allocation.
        """
        allocation = {agent: [] for agent in self.agents.values()}
        
        for task in tasks:
            suitable = self.allocate_task(task)
            if suitable:
                # Pick least loaded agent
                agent = min(suitable, key=lambda a: len(allocation[a]))
                allocation[agent].append(task)
                
        return allocation

class NegotiationProtocol:
    """
    Contract net protocol for agent negotiation.
    """
    
    def __init__(self):
        self.auctions = {}  # task_id -> auction info
        
    def initiate_auction(self, task: Dict, initiator: Agent):
        """Start auction for task."""
        auction_id = f"auction_{task['id']}"
        self.auctions[auction_id] = {
            "task": task,
            "initiator": initiator,
            "bids": {},
            "deadline": asyncio.get_event_loop().time() + 10
        }
        
    async def submit_bid(self, auction_id: str, bidder: Agent, bid: Dict):
        """Submit bid for auction."""
        if auction_id in self.auctions:
            self.auctions[auction_id]["bids"][bidder.agent_id] = {
                "bidder": bidder,
                "proposal": bid
            }
            
    def select_winner(self, auction_id: str) -> Agent:
        """Select winning bid based on criteria."""
        auction = self.auctions[auction_id]
        bids = auction["bids"]
        
        # Simple: select first bid
        # Real implementation: evaluate based on cost, capability, etc.
        if bids:
            winner_id = list(bids.keys())[0]
            return bids[winner_id]["bidder"]
        return None

class ConsensusProtocol:
    """
    Reach agreement among agents.
    """
    
    def __init__(self, agents: List[Agent]):
        self.agents = agents
        self.votes = {}
        
    def propose(self, proposer: Agent, proposal: Dict):
        """Agent proposes a value."""
        self.votes[proposer.agent_id] = proposal
        
    async def vote(self, agent: Agent, value: Any):
        """Agent votes on proposal."""
        self.votes[agent.agent_id] = value
        
    def majority_vote(self) -> Any:
        """Calculate majority vote."""
        votes = list(self.votes.values())
        if not votes:
            return None
            
        # Count votes
        counts = {}
        for vote in votes:
            counts[vote] = counts.get(vote, 0) + 1
            
        # Return most voted
        return max(counts.items(), key=lambda x: x[1])[0]
    
    def weighted_vote(self, weights: Dict[str, float]) -> Any:
        """Weighted voting based on agent trust/capability."""
        weighted_scores = {}
        
        for agent_id, vote in self.votes.items():
            weight = weights.get(agent_id, 1.0)
            weighted_scores[vote] = weighted_scores.get(vote, 0) + weight
            
        return max(weighted_scores.items(), key=lambda x: x[1])[0]

class HierarchicalOrchestrator:
    """
    Hierarchical agent organization with supervisor agents.
    """
    
    def __init__(self):
        self.supervisors = {}  # supervisor_id -> subordinate_ids
        self.workers = {}  # worker_id -> Agent
        
    def add_supervisor(self, agent: Agent, subordinates: List[Agent]):
        """Add supervisor with direct reports."""
        self.supervisors[agent.agent_id] = {
            "agent": agent,
            "subordinates": subordinates
        }
        
        for sub in subordinates:
            self.workers[sub.agent_id] = sub
            
    def coordinate(self, task: Dict) -> Any:
        """Coordinate task through hierarchy."""
        # Find root supervisor
        root = self._find_root_supervisor()
        
        # Delegate down the hierarchy
        return self._delegate_task(root, task)
        
    def _find_root_supervisor(self) -> Agent:
        """Find top-level supervisor."""
        for sup_id, info in self.supervisors.items():
            if not self._is_subordinate(sup_id):
                return info["agent"]
        return None
        
    def _is_subordinate(self, agent_id: str) -> bool:
        """Check if agent is subordinate to anyone."""
        for info in self.supervisors.values():
            if agent_id in [s.agent_id for s in info["subordinates"]]:
                return True
        return False
        
    def _delegate_task(self, supervisor: Agent, task: Dict) -> Any:
        """Delegate task down the hierarchy."""
        # Supervisor decomposes task
        subtasks = supervisor.decompose(task)
        
        # Assign to subordinates
        results = []
        for subtask in subtasks:
            worker = self._find_worker_for_task(subtask)
            result = worker.execute(subtask)
            results.append(result)
            
        # Supervisor synthesizes results
        return supervisor.synthesize(results)

class MarketOrchestrator:
    """
    Market-based coordination with bidding agents.
    """
    
    def __init__(self):
        self.tasks = []
        self.agents = {}  # agent_id -> {capabilities, price, availability}
        
    def register_agent(self, agent_id: str, capabilities: List[str], 
                      price_per_task: float):
        self.agents[agent_id] = {
            "capabilities": set(capabilities),
            "price": price_per_task,
            "availability": True
        }
        
    def submit_task(self, task: Dict):
        """Submit task to marketplace."""
        self.tasks.append(task)
        
    def match_tasks(self) -> Dict[Agent, List[Dict]]:
        """Match tasks to agents via bidding."""
        assignments = {agent_id: [] for agent_id in self.agents}
        
        for task in self.tasks:
            required = set(task.get("required_capabilities", []))
            
            # Find capable agents
            capable = [
                (aid, info) for aid, info in self.agents.items()
                if required.issubset(info["capabilities"]) and info["availability"]
            ]
            
            if capable:
                # Select cheapest
                selected = min(capable, key=lambda x: x[1]["price"])
                assignments[selected[0]].append(task)
                
                # Mark agent as unavailable
                selected[1]["availability"] = False
                
        return assignments

class SwarmCoordinator:
    """
    Swarm-based coordination for large agent populations.
    """
    
    def __init__(self):
        self.agents = []
        
    def add_agent(self, agent: Agent):
        self.agents.append(agent)
        
    def emergent_coordination(self, global_goal: Dict):
        """
        Stigmergy-based coordination: agents coordinate through environment.
        """
        for agent in self.agents:
            # Agent senses local environment
            local_view = self._get_local_view(agent)
            
            # Agent applies simple rules
            action = self._stigmergy_rule(agent, local_view, global_goal)
            
            # Agent acts
            agent.execute(action)
            
    def _get_local_view(self, agent: Agent) -> Dict:
        """Get local environment view for agent."""
        # Return nearby agents and environment state
        pass
        
    def _stigmergy_rule(self, agent: Agent, local_view: Dict, 
                        goal: Dict) -> Any:
        """
        Simple rule for stigmergic coordination.
        Example: move toward goal, avoid collisions, stay near neighbors.
        """
        # Simplified: move toward goal
        return {"action": "move_toward", "target": goal.get("location")}

class CodeGenerationTeam:
    """
    Multi-agent team for code generation.
    """
    
    def __init__(self):
        self.planner = Agent("planner", "planner")
        self.coders = [
            Agent(f"coder_{i}", "coder") for i in range(3)
        ]
        self.reviewer = Agent("reviewer", "reviewer")
        self.tester = Agent("tester", "tester")
        
        self.all_agents = [self.planner] + self.coders + [self.reviewer, self.tester]
        
    async def generate_code(self, specification: str) -> str:
        # Step 1: Planner decomposes task
        plan = await self._planning_phase(specification)
        
        # Step 2: Coders implement in parallel
        modules = await self._coding_phase(plan)
        
        # Step 3: Reviewer checks code
        feedback = await self._review_phase(modules)
        
        # Step 4: Iterative refinement
        for _ in range(3):
            if not feedback:
                break
            modules = await self._refine_phase(modules, feedback)
            feedback = await self._review_phase(modules)
            
        # Step 5: Tester validates
        test_results = await self._testing_phase(modules)
        
        return self._combine_modules(modules)
    
    async def _planning_phase(self, spec: str) -> Dict:
        """Planner creates implementation plan."""
        prompt = f"""Create a detailed implementation plan for:
{spec}

Return: list of modules with their responsibilities"""
        # Use LLM to create plan
        return {"modules": ["module1", "module2", "module3"]}
    
    async def _coding_phase(self, plan: Dict) -> List[str]:
        """Coders implement in parallel."""
        tasks = []
        for module in plan["modules"]:
            tasks.append(self._implement_module(module))
        return await asyncio.gather(*tasks)
    
    async def _implement_module(self, module: str) -> str:
        """Single coder implements module."""
        return f"# Implementation of {module}\npass"
    
    async def _review_phase(self, modules: List[str]) -> List[str]:
        """Reviewer provides feedback."""
        return []  # No issues found
    
    async def _refine_phase(self, modules: List[str], 
                           feedback: List[str]) -> List[str]:
        """Refine code based on feedback."""
        return modules
    
    async def _testing_phase(self, modules: List[str]) -> Dict:
        """Tester runs tests."""
        return {"passed": True}
    
    def _combine_modules(self, modules: List[str]) -> str:
        return "\n\n".join(modules)

class ResearchTeam:
    """
    Multi-agent research assistant team.
    """
    
    def __init__(self):
        self.coordinator = Agent("coordinator", "coordinator")
        self.researchers = [
            Agent(f"researcher_{i}", "researcher") for i in range(2)
        ]
        self.analyst = Agent("analyst", "analyst")
        self.writer = Agent("writer", "writer")
        
    async def conduct_research(self, topic: str) -> str:
        # Coordinator distributes research tasks
        subtopics = await self._decompose_topic(topic)
        
        # Researchers gather information in parallel
        findings = await self._gather_findings(subtopics)
        
        # Analyst synthesizes findings
        analysis = await self._analyze_findings(findings)
        
        # Writer produces final report
        report = await self._write_report(topic, analysis)
        
        return report
    
    async def _decompose_topic(self, topic: str) -> List[str]:
        """Break topic into subtopics."""
        return ["aspect1", "aspect2", "aspect3"]
    
    async def _gather_findings(self, subtopics: List[str]) -> List[Dict]:
        """Researchers gather information."""
        tasks = [self._research(subtopic) for subtopic in subtopics]
        return await asyncio.gather(*tasks)
    
    async def _research(self, subtopic: str) -> Dict:
        """Research single aspect."""
        return {"topic": subtopic, "findings": "..."}
    
    async def _analyze_findings(self, findings: List[Dict]) -> Dict:
        """Analyze combined findings."""
        return {"summary": "...", "insights": []}
    
    async def _write_report(self, topic: str, analysis: Dict) -> str:
        """Write final report."""
        return f"# Research Report: {topic}\n\n{analysis['summary']}"

class MCPProtocol:
    """
    Model Context Protocol for agent communication.
    """
    
    @staticmethod
    def create_request(tool_name: str, params: Dict) -> Dict:
        return {
            "jsonrpc": "2.0",
            "method": f"tools/{tool_name}",
            "params": params,
            "id": 1
        }
    
    @staticmethod
    def create_response(result: Any, request_id: int) -> Dict:
        return {
            "jsonrpc": "2.0",
            "result": result,
            "id": request_id
        }
    
    @staticmethod
    def create_error(code: int, message: str, request_id: int) -> Dict:
        return {
            "jsonrpc": "2.0",
            "error": {
                "code": code,
                "message": message
            },
            "id": request_id
        }

class A2AProtocol:
    """
    Agent-to-Agent protocol for agent interoperability.
    """
    
    @staticmethod
    def create_agent_card(agent: Agent) -> Dict:
        """Describe agent capabilities."""
        return {
            "name": agent.agent_id,
            "role": agent.role,
            "capabilities": agent.get_capabilities(),
            "version": "1.0"
        }
    
    @staticmethod
    def create_task_request(task: Dict) -> Dict:
        """Request task execution from agent."""
        return {
            "taskId": task.get("id"),
            "message": {
                "role": "user",
                "parts": [{"type": "text", "text": task["description"]}]
            }
        }

def design_agent_roles(task: Dict) -> List[Dict]:
    """
    Design appropriate agent roles for a task.
    """
    # Analyze task requirements
    requirements = task.get("requirements", [])
    
    roles = []
    
    # Coordinator for task management
    if len(requirements) > 3:
        roles.append({"role": "coordinator", "responsibility": "task_decomposition"})
    
    # Specialists for each requirement
    for req in requirements:
        roles.append({"role": f"specialist_{req}", "responsibility": req})
    
    # Validator for quality assurance
    roles.append({"role": "validator", "responsibility": "verification"})
    
    return roles

class ConflictResolver:
    """
    Resolve disagreements between agents.
    """
    
    @staticmethod
    def voting(agent_opinions: Dict[str, Any]) -> Any:
        """Resolve via majority vote."""
        votes = {}
        for opinion in agent_opinions.values():
            votes[opinion] = votes.get(opinion, 0) + 1
        return max(votes.items(), key=lambda x: x[1])[0]
    
    @staticmethod
    def authority(agent_hierarchy: Dict[str, int], 
                agent_opinions: Dict[str, Any]) -> Any:
        """Defer to highest authority."""
        max_authority = -1
        winner = None
        
        for agent_id, opinion in agent_opinions.items():
            authority = agent_hierarchy.get(agent_id, 0)
            if authority > max_authority:
                max_authority = authority
                winner = opinion
                
        return winner
    
    @staticmethod
    def arbitration(arbitrator: Agent, 
                   agent_opinions: Dict[str, Any]) -> Any:
        """Let designated arbitrator decide."""
        return arbitrator.judge(agent_opinions)