diff --git a/swarms/structs/swarm_arange.py b/swarms/structs/swarm_arange.py index 7968c172..c8513b59 100644 --- a/swarms/structs/swarm_arange.py +++ b/swarms/structs/swarm_arange.py @@ -8,6 +8,7 @@ from swarms.structs.conversation import Conversation from swarms.utils.history_output_formatter import ( HistoryOutputType, ) +from swarms.security import SwarmShieldIntegration, ShieldConfig logger = initialize_logger(log_folder="swarm_arange") @@ -33,15 +34,6 @@ class SwarmRearrange: return_json (bool): A flag indicating whether to return the result in JSON format swarm_history (dict): A dictionary to keep track of the history of each swarm lock (threading.Lock): A lock for thread-safe operations - - Methods: - __init__(id: str, name: str, description: str, swarms: List[swarm], flow: str, max_loops: int, verbose: bool, - human_in_the_loop: bool, custom_human_in_the_loop: Callable, return_json: bool): Initializes the SwarmRearrange object - add_swarm(swarm: swarm): Adds an swarm to the swarm - remove_swarm(swarm_name: str): Removes an swarm from the swarm - add_swarms(swarms: List[swarm]): Adds multiple swarms to the swarm - validate_flow(): Validates the flow pattern - run(task): Runs the swarm to rearrange the tasks """ def __init__( @@ -59,6 +51,9 @@ class SwarmRearrange: ] = None, return_json: bool = False, output_type: HistoryOutputType = "dict-all-except-first", + shield_config: Optional[ShieldConfig] = None, + enable_security: bool = True, + security_level: str = "standard", *args, **kwargs, ): @@ -76,6 +71,10 @@ class SwarmRearrange: human_in_the_loop (bool): Whether human intervention is required. Defaults to False. custom_human_in_the_loop (Callable): Custom function for human intervention. Defaults to None. return_json (bool): Whether to return results as JSON. Defaults to False. + output_type (HistoryOutputType): Type of output format. Defaults to "dict-all-except-first". + shield_config (ShieldConfig, optional): Security configuration for SwarmShield integration. Defaults to None. + enable_security (bool, optional): Whether to enable SwarmShield security features. Defaults to True. + security_level (str, optional): Pre-defined security level. Options: "basic", "standard", "enhanced", "maximum". Defaults to "standard". """ self.id = id self.name = name @@ -89,6 +88,9 @@ class SwarmRearrange: self.output_type = output_type self.return_json = return_json + # Initialize SwarmShield integration + self._initialize_swarm_shield(shield_config, enable_security, security_level) + self.swarm_history = {swarm.name: [] for swarm in swarms} self.lock = threading.Lock() self.id = uuid.uuid4().hex if id is None else id @@ -351,6 +353,86 @@ class SwarmRearrange: logger.error(f"An error occurred: {e}") return str(e) + def _initialize_swarm_shield( + self, + shield_config: Optional[ShieldConfig] = None, + enable_security: bool = True, + security_level: str = "standard" + ) -> None: + """Initialize SwarmShield integration for security features.""" + self.enable_security = enable_security + self.security_level = security_level + + if enable_security: + if shield_config is None: + shield_config = ShieldConfig.get_security_level(security_level) + + self.swarm_shield = SwarmShieldIntegration(shield_config) + else: + self.swarm_shield = None + + # Security methods + def validate_task_with_shield(self, task: str) -> str: + """Validate and sanitize task input using SwarmShield.""" + if self.swarm_shield: + return self.swarm_shield.validate_and_protect_input(task) + return task + + def validate_agent_config_with_shield(self, agent_config: dict) -> dict: + """Validate agent configuration using SwarmShield.""" + if self.swarm_shield: + return self.swarm_shield.validate_and_protect_input(str(agent_config)) + return agent_config + + def process_agent_communication_with_shield(self, message: str, agent_name: str) -> str: + """Process agent communication through SwarmShield security.""" + if self.swarm_shield: + return self.swarm_shield.process_agent_communication(message, agent_name) + return message + + def check_rate_limit_with_shield(self, agent_name: str) -> bool: + """Check rate limits for an agent using SwarmShield.""" + if self.swarm_shield: + return self.swarm_shield.check_rate_limit(agent_name) + return True + + def add_secure_message(self, message: str, agent_name: str) -> None: + """Add a message to secure conversation history.""" + if self.swarm_shield: + self.swarm_shield.add_secure_message(message, agent_name) + + def get_secure_messages(self) -> List[dict]: + """Get secure conversation messages.""" + if self.swarm_shield: + return self.swarm_shield.get_secure_messages() + return [] + + def get_security_stats(self) -> dict: + """Get security statistics and metrics.""" + if self.swarm_shield: + return self.swarm_shield.get_security_stats() + return {"security_enabled": False} + + def update_shield_config(self, new_config: ShieldConfig) -> None: + """Update SwarmShield configuration.""" + if self.swarm_shield: + self.swarm_shield.update_config(new_config) + + def enable_security(self) -> None: + """Enable SwarmShield security features.""" + if not self.swarm_shield: + self._initialize_swarm_shield(enable_security=True, security_level=self.security_level) + + def disable_security(self) -> None: + """Disable SwarmShield security features.""" + self.swarm_shield = None + self.enable_security = False + + def cleanup_security(self) -> None: + """Clean up SwarmShield resources.""" + if self.swarm_shield: + self.swarm_shield.cleanup() + def swarm_arrange( name: str = "SwarmArrange-01",