# swarm_coordinator.py import asyncio import json import logging import os import psutil import time from datetime import datetime from enum import Enum from pathlib import Path from typing import Any, Callable, Dict, List, Optional # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') class AgentStatus(Enum): IDLE = "idle" WORKING = "working" FAILED = "failed" COMPLETED = "completed" class Agent: """Represents a single agent in the swarm.""" def __init__( self, agent_id: str, agent_type: str, capabilities: List[str], error_callback: Callable[[str, str], None] = None, ): """ Initializes an Agent. Args: agent_id (str): Unique identifier for the agent. agent_type (str): Type of agent (e.g., "researcher", "analyzer"). capabilities (List[str]): List of tasks the agent can perform. error_callback (Callable[[str, str], None], optional): Callback function to handle errors. Defaults to None. """ self.agent_id = agent_id self.agent_type = agent_type self.capabilities = capabilities self.status = AgentStatus.IDLE self.current_task: Optional[str] = None self.result: Optional[Any] = None self.error_callback = error_callback async def execute_task(self, task: str, shared_context: Dict) -> Any: """ Executes a given task. Args: task (str): The task to be executed. shared_context (Dict): Shared context data. Returns: Any: The result of the task execution. """ self.status = AgentStatus.WORKING self.current_task = task logging.info(f"Agent {self.agent_id} started task: {task}") try: # Simulate task execution (replace with actual logic) await asyncio.sleep(0.5) # Simulate work self.result = f"Result from {self.agent_id} for task: {task}" self.status = AgentStatus.COMPLETED logging.info(f"Agent {self.agent_id} completed task: {task}") return self.result except Exception as e: self.status = AgentStatus.FAILED logging.error(f"Agent {self.agent_id} failed task: {task} with error: {e}") if self.error_callback: self.error_callback(self.agent_id, str(e)) return None def reset(self): """Resets the agent to an idle state.""" self.status = AgentStatus.IDLE self.current_task = None self.result = None class SwarmCoordinator: """Manages the swarm of agents and coordinates task execution.""" def __init__(self, config_path: str = "swarm_config.json"): """ Initializes the SwarmCoordinator. Args: config_path (str, optional): Path to the swarm configuration file. Defaults to "swarm_config.json". """ self.agents: Dict[str, Agent] = {} self.task_queue: asyncio.Queue = asyncio.Queue() self.results: Dict[str, Any] = {} self.config_path = Path(config_path) self.load_config() self.budget = 1000 # Example budget self.tasks_completed = 0 self.failed_tasks = 0 def load_config(self): """Loads agent configurations from a JSON file.""" try: with open(self.config_path, "r") as f: config = json.load(f) for agent_data in config.get("agents", []): agent_id = agent_data["id"] self.agents[agent_id] = Agent( agent_id=agent_id, agent_type=agent_data["type"], capabilities=agent_data["capabilities"], error_callback=self.handle_agent_failure, ) logging.info("Swarm configuration loaded successfully.") except FileNotFoundError: logging.error(f"Configuration file not found: {self.config_path}") except json.JSONDecodeError: logging.error(f"Invalid JSON format in configuration file: {self.config_path}") except Exception as e: logging.error(f"Error loading configuration: {e}") def add_task(self, task: str): """Adds a task to the task queue.""" self.task_queue.put_nowait(task) logging.info(f"Task added to queue: {task}") async def assign_tasks(self): """Assigns tasks to available agents.""" while True: task = await self.task_queue.get() logging.info(f"Assigning task: {task}") assigned = False for agent_id, agent in self.agents.items(): if agent.status == AgentStatus.IDLE and self.can_agent_perform_task(agent, task): asyncio.create_task(self.execute_task_on_agent(agent, task)) assigned = True break if not assigned: logging.warning(f"No suitable agent found for task: {task}. Re-queueing.") self.task_queue.put_nowait(task) # Re-queue if no agent is available self.task_queue.task_done() await asyncio.sleep(0.1) # Prevent busy-waiting async def execute_task_on_agent(self, agent: Agent, task: str): """ Executes a task on a specific agent. Args: agent (Agent): The agent to execute the task. task (str): The task to be executed. """ result = await agent.execute_task(task, {"swarm_name": "Genesis Swarm"}) if result: self.results[task] = result self.tasks_completed += 1 self.budget -= 10 # Simulate cost logging.info(f"Task completed: {task} by agent {agent.agent_id}") else: self.handle_agent_failure(agent.agent_id, f"Task {task} failed.") agent.reset() def can_agent_perform_task(self, agent: Agent, task: str) -> bool: """ Checks if an agent has the capability to perform a given task. Args: agent (Agent): The agent to check. task (str): The task to be performed. Returns: bool: True if the agent can perform the task, False otherwise. """ # Simple capability matching return any(keyword.lower() in task.lower() for keyword in agent.capabilities) def handle_agent_failure(self, agent_id: str, error_message: str): """Handles agent failures and reassigns tasks if necessary.""" logging.error(f"Agent {agent_id} failed: {error_message}") self.failed_tasks += 1 # Re-queue the task if it failed failed_agent = self.agents.get(agent_id) if failed_agent and failed_agent.current_task: self.add_task(failed_agent.current_task) failed_agent.reset() def get_swarm_status(self) -> Dict[str, Any]: """Returns the current status of the swarm.""" agent_statuses = { agent_id: agent.status.value for agent_id, agent in self.agents.items() } return { "swarm_status": "active", "agent_statuses": agent_statuses, "tasks_completed": self.tasks_completed, "failed_tasks": self.failed_tasks, "budget": self.budget, } async def run(self): """Runs the swarm coordinator.""" logging.info("Swarm Coordinator started.") assigner_task = asyncio.create_task(self.assign_tasks()) # Example tasks self.add_task("Research the optimal configuration for MCP servers.") self.add_task("Analyze performance data of existing servers.") self.add_task("Implement security protocols on new servers.") self.add_task("Review code for MCP server implementation.") self.add_task("Document the deployment process for MCP servers.") self.add_task("Find vulnerabilities in network infrastructure.") self.add_task("Optimize database queries for faster data retrieval.") self.add_task("Test new software components for compatibility.") self.add_task("Create user manuals for new applications.") self.add_task("Monitor system logs for suspicious activity.") await self.task_queue.join() # Wait for all tasks to be processed assigner_task.cancel() logging.info("All tasks completed.") logging.info(f"Swarm Status: {self.get_swarm_status()}") async def main(): """Main function to run the swarm coordinator.""" # Example configuration config = { "agents": [ { "id": "researcher_1", "type": "researcher", "capabilities": ["research", "information gathering"], }, { "id": "analyzer_1", "type": "analyzer", "capabilities": ["analyze", "data analysis"], }, { "id": "implementer_1", "type": "implementer", "capabilities": ["implement", "code"], }, { "id": "reviewer_1", "type": "reviewer", "capabilities": ["review", "validate"], }, { "id": "documenter_1", "type": "documenter", "capabilities": ["document", "summarize"], }, { "id": "security_1", "type": "security", "capabilities": ["security", "vulnerability"], }, ] } # Create a dummy config file with open("swarm_config.json", "w") as f: json.dump(config, f, indent=4) coordinator = SwarmCoordinator() await coordinator.run() if __name__ == "__main__": asyncio.run(main())