Update swarm_arange.py

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",