# swarm_coordinator.py import asyncio import json import logging import os import random import time from datetime import datetime from enum import Enum from pathlib import Path from typing import List, Dict, Any, Optional, Callable from dataclasses import dataclass, field, asdict # 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" @dataclass class Agent: agent_id: str agent_type: str capabilities: List[str] status: AgentStatus = AgentStatus.IDLE task_id: Optional[str] = None result: Optional[Any] = None error: Optional[str] = None def to_dict(self) -> Dict: d = asdict(self) d['status'] = self.status.value return d class TaskStatus(Enum): PENDING = "pending" IN_PROGRESS = "in_progress" COMPLETED = "completed" FAILED = "failed" @dataclass class Task: task_id: str description: str required_capabilities: List[str] status: TaskStatus = TaskStatus.PENDING assigned_agent: Optional[str] = None result: Optional[Any] = None error: Optional[str] = None budget: float = 100.0 # Example budget start_time: Optional[datetime] = None end_time: Optional[datetime] = None def to_dict(self) -> Dict: d = asdict(self) d['status'] = self.status.value return d class SwarmCoordinator: """ Manages a swarm of agents, distributes tasks, and aggregates results. Implements fault tolerance and progress tracking. """ def __init__(self, config_path: str = "swarm_config.json", persist_path: str = "swarm_data"): self.agents: Dict[str, Agent] = {} self.tasks: Dict[str, Task] = {} self.config_path = Path(config_path) self.persist_path = Path(persist_path) self.persist_path.mkdir(parents=True, exist_ok=True) self._load_config() self.task_counter = 0 def _load_config(self): """Loads agent configurations from a JSON file.""" if not self.config_path.exists(): logging.warning(f"Config file not found at {self.config_path}. Using default agents.") self._create_default_agents() self.save_config() return try: with open(self.config_path, "r", encoding="utf-8") as f: config = json.load(f) for agent_data in config.get("agents", []): agent = Agent( agent_id=agent_data["agent_id"], agent_type=agent_data["agent_type"], capabilities=agent_data["capabilities"] ) self.agents[agent.agent_id] = agent logging.info(f"Loaded agent configurations from {self.config_path}") except Exception as e: logging.error(f"Error loading config: {e}") self._create_default_agents() self.save_config() def save_config(self): """Saves the current agent configuration to a JSON file.""" data = {"agents": [agent.to_dict() for agent in self.agents.values()]} try: with open(self.config_path, "w", encoding="utf-8") as f: json.dump(data, f, indent=2) logging.info(f"Saved agent configurations to {self.config_path}") except Exception as e: logging.error(f"Error saving config: {e}") def _create_default_agents(self): """Creates a default set of agents if no config is found.""" logging.info("Creating default agents.") self.agents = { "queen_core": Agent(agent_id="queen_core", agent_type="core", capabilities=["decision_making", "task_management"]), "guardian_1": Agent(agent_id="guardian_1", agent_type="guardian", capabilities=["validation", "security"]), "guardian_2": Agent(agent_id="guardian_2", agent_type="guardian", capabilities=["validation", "security"]), "guardian_3": Agent(agent_id="guardian_3", agent_type="guardian", capabilities=["validation", "security"]), "guardian_4": Agent(agent_id="guardian_4", agent_type="guardian", capabilities=["validation", "security"]), "guardian_5": Agent(agent_id="guardian_5", agent_type="guardian", capabilities=["validation", "security"]), "guardian_6": Agent(agent_id="guardian_6", agent_type="guardian", capabilities=["validation", "security"]), "processor_1": Agent(agent_id="processor_1", agent_type="processor", capabilities=["data_processing", "analysis"]), "processor_2": Agent(agent_id="processor_2", agent_type="processor", capabilities=["data_processing", "analysis"]), "processor_3": Agent(agent_id="processor_3", agent_type="processor", capabilities=["data_processing", "analysis"]), "processor_4": Agent(agent_id="processor_4", agent_type="processor", capabilities=["data_processing", "analysis"]), "processor_5": Agent(agent_id="processor_5", agent_type="processor", capabilities=["data_processing", "analysis"]), "processor_6": Agent(agent_id="processor_6", agent_type="processor", capabilities=["data_processing", "analysis"]), "processor_7": Agent(agent_id="processor_7", agent_type="processor", capabilities=["data_processing", "analysis"]), "processor_8": Agent(agent_id="processor_8", agent_type="processor", capabilities=["data_processing", "analysis"]), "processor_9": Agent(agent_id="processor_9", agent_type="processor", capabilities=["data_processing", "analysis"]), "processor_10": Agent(agent_id="processor_10", agent_type="processor", capabilities=["data_processing", "analysis"]), "worker_1": Agent(agent_id="worker_1", agent_type="worker", capabilities=["execution", "code_generation"]), "worker_2": Agent(agent_id="worker_2", agent_type="worker", capabilities=["execution", "code_generation"]), "worker_3": Agent(agent_id="worker_3", agent_type="worker", capabilities=["execution", "code_generation"]), "worker_4": Agent(agent_id="worker_4", agent_type="worker", capabilities=["execution", "code_generation"]), "worker_5": Agent(agent_id="worker_5", agent_type="worker", capabilities=["execution", "code_generation"]), } def create_task(self, description: str, required_capabilities: List[str], budget: float = 100.0) -> str: """Creates a new task and adds it to the task queue.""" self.task_counter += 1 task_id = f"task_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{self.task_counter}" task = Task( task_id=task_id, description=description, required_capabilities=required_capabilities, budget=budget ) self.tasks[task_id] = task self.save_state() logging.info(f"Created task: {task_id} - {description}") return task_id def assign_task(self, task_id: str) -> bool: """Assigns a task to a suitable agent based on capabilities.""" task = self.tasks.get(task_id) if not task: logging.error(f"Task not found: {task_id}") return False if task.status != TaskStatus.PENDING: logging.warning(f"Task {task_id} is not in PENDING state.") return False available_agents = [ agent for agent in self.agents.values() if agent.status == AgentStatus.IDLE and all(cap in agent.capabilities for cap in task.required_capabilities) ] if not available_agents: logging.warning(f"No suitable agents available for task: {task_id}") return False # Assign task to a random available agent agent = random.choice(available_agents) task.assigned_agent = agent.agent_id task.status = TaskStatus.IN_PROGRESS task.start_time = datetime.now() agent.status = AgentStatus.WORKING agent.task_id = task_id self.save_state() logging.info(f"Assigned task {task_id} to agent {agent.agent_id}") return True async def execute_task(self, task_id: str, executor: Callable = None): """Executes a task assigned to an agent.""" task = self.tasks.get(task_id) if not task: logging.error(f"Task not found: {task_id}") return agent_id = task.assigned_agent agent = self.agents.get(agent_id) if not agent: logging.error(f"Agent not found: {agent_id}") task.status = TaskStatus.FAILED task.error = f"Assigned agent {agent_id} not found." self.save_state() return try: logging.info(f"Executing task {task_id} with agent {agent_id}") if executor: result = await executor(task.description, agent.capabilities) else: # Simulate task execution (replace with actual execution logic) await asyncio.sleep(random.uniform(1, 5)) result = f"Task {task_id} completed successfully by {agent_id}." task.status = TaskStatus.COMPLETED task.result = result task.end_time = datetime.now() agent.status = AgentStatus.COMPLETED agent.result = result agent.task_id = None logging.info(f"Task {task_id} completed by agent {agent_id}") except Exception as e: logging.error(f"Task {task_id} failed: {e}") task.status = TaskStatus.FAILED task.error = str(e) agent.status = AgentStatus.FAILED agent.error = str(e) self.reassign_task(task_id) finally: self.save_state() def reassign_task(self, task_id: str): """Reassigns a failed task to another available agent.""" task = self.tasks.get(task_id) if not task: logging.error(f"Task not found: {task_id}") return task.assigned_agent = None task.status = TaskStatus.PENDING failed_agent_id = [agent_id for agent_id, agent in self.agents.items() if agent.task_id == task_id][0] failed_agent = self.agents[failed_agent_id] failed_agent.status = AgentStatus.IDLE failed_agent.task_id = None failed_agent.result = None logging.info(f"Reassigning task {task_id} after failure.") self.assign_task(task_id) def aggregate_results(self) -> Dict[str, Any]: """Aggregates results from completed tasks.""" completed_tasks = [task for task in self.tasks.values() if task.status == TaskStatus.COMPLETED] results = { "completed_tasks": len(completed_tasks), "success_rate": len(completed_tasks) / len(self.tasks) if self.tasks else 0, "task_results": [ { "task_id": task.task_id, "description": task.description, "result": task.result, "assigned_agent": task.assigned_agent } for task in completed_tasks ] } return results def track_progress(self) -> Dict[str, int]: """Tracks the progress of all tasks.""" task_counts = {} for status in TaskStatus: task_counts[status.value] = len([task for task in self.tasks.values() if task.status == status]) return task_counts def track_budget(self) -> float: """Tracks the remaining budget across all tasks.""" total_budget = sum(task.budget for task in self.tasks.values()) # In a real-world scenario, track actual costs against the budget return total_budget def save_state(self): """Persists the current state of the swarm to disk.""" state = { "agents": [agent.to_dict() for agent in self.agents.values()], "tasks": [task.to_dict() for task in self.tasks.values()] } timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") filepath = self.persist_path / f"swarm_state_{timestamp}.json" try: with open(filepath, "w", encoding="utf-8") as f: json.dump(state, f, indent=2) logging.info(f"Saved swarm state to {filepath}") except Exception as e: logging.error(f"Error saving swarm state: {e}") def load_state(self, filepath: str): """Loads the swarm state from a JSON file.""" filepath = Path(filepath) if not filepath.exists(): logging.error(f"State file not found: {filepath}") return try: with open(filepath, "r", encoding="utf-8") as f: state = json.load(f) agents_data = state.get("agents", []) tasks_data = state.get("tasks", []) self.agents = {} for agent_data in agents_data: agent = Agent(**{k: (AgentStatus(v) if k == 'status' else v) for k, v in agent_data.items()}) self.agents[agent.agent_id] = agent self.tasks = {} for task_data in tasks_data: task = Task(**{k: (TaskStatus(v) if k == 'status' else v) for k, v in task_data.items()}) self.tasks[task.task_id] = task logging.info(f"Loaded swarm state from {filepath}") except Exception as e: logging.error(f"Error loading swarm state: {e}") async def run_orchestration(self, task_description: str, required_capabilities: List[str], executor: Callable = None): """Runs the complete orchestration flow for a given task.""" task_id = self.create_task(task_description, required_capabilities) if self.assign_task(task_id): await self.execute_task(task_id, executor) else: logging.error(f"Failed to assign task {task_id}") results = self.aggregate_results() progress = self.track_progress() budget = self.track_budget() logging.info(f"Orchestration completed. Results: {results}, Progress: {progress}, Budget: {budget}") return results, progress, budget async def example_executor(task_description: str, capabilities: List[str]) -> str: """Example executor function that simulates task execution.""" await asyncio.sleep(random.uniform(0.5, 2)) return f"Task '{task_description}' executed with capabilities: {capabilities}" async def main(): """Main function to demonstrate the SwarmCoordinator.""" coordinator = SwarmCoordinator() # Example task task_description = "Analyze market trends and generate a report." required_capabilities = ["data_processing", "analysis"] # Run the orchestration results, progress, budget = await coordinator.run_orchestration(task_description, required_capabilities, executor=example_executor) print("Final Results:", results) print("Progress:", progress) print("Remaining Budget:", budget) if __name__ == "__main__": asyncio.run(main())