From 6a99e2ce2c09dec4c0a6a7238c0ae0f599845230 Mon Sep 17 00:00:00 2001 From: CI-DEV <154627941+IlumCI@users.noreply.github.com> Date: Mon, 18 Aug 2025 19:47:09 +0300 Subject: [PATCH] Update sequential_workflow.py --- swarms/structs/sequential_workflow.py | 112 +++++++++++++++++++++++++- 1 file changed, 111 insertions(+), 1 deletion(-) diff --git a/swarms/structs/sequential_workflow.py b/swarms/structs/sequential_workflow.py index c845d508..4f49bd1c 100644 --- a/swarms/structs/sequential_workflow.py +++ b/swarms/structs/sequential_workflow.py @@ -1,10 +1,12 @@ from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Callable, List, Optional, Union +from typing import Callable, List, Optional, Union, Tuple +from datetime import datetime from swarms.structs.agent import Agent from swarms.structs.rearrange import AgentRearrange from swarms.utils.loguru_logger import initialize_logger from swarms.utils.output_types import OutputType +from swarms.security import SwarmShieldIntegration, ShieldConfig logger = initialize_logger(log_folder="sequential_workflow") @@ -20,6 +22,9 @@ class SequentialWorkflow: max_loops (int, optional): The maximum number of times to execute the workflow. Defaults to 1. output_type (OutputType, optional): The format of the output from the workflow. Defaults to "dict". shared_memory_system (callable, optional): A callable for managing shared memory between agents. Defaults to None. + 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". *args: Additional positional arguments. **kwargs: Additional keyword arguments. @@ -36,6 +41,9 @@ class SequentialWorkflow: max_loops: int = 1, output_type: OutputType = "dict", shared_memory_system: callable = None, + shield_config: Optional[ShieldConfig] = None, + enable_security: bool = True, + security_level: str = "standard", *args, **kwargs, ): @@ -47,6 +55,9 @@ class SequentialWorkflow: self.output_type = output_type self.shared_memory_system = shared_memory_system + # Initialize SwarmShield integration + self._initialize_swarm_shield(shield_config, enable_security, security_level) + self.reliability_check() self.flow = self.sequential_flow() @@ -59,6 +70,105 @@ class SequentialWorkflow: output_type=self.output_type, ) + 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) + logger.info(f"SwarmShield initialized with {security_level} security level") + else: + self.swarm_shield = None + logger.info("SwarmShield security disabled") + + # Security methods + def validate_task_with_shield(self, task: str, agent_name: str = "default") -> Tuple[bool, str, Optional[str]]: + """Validate and sanitize task input using SwarmShield.""" + if self.swarm_shield: + return self.swarm_shield.validate_task(task, agent_name) + return True, task, None + + def validate_agent_config_with_shield(self, agent_config: dict, agent_name: str = "default") -> Tuple[bool, dict, Optional[str]]: + """Validate agent configuration using SwarmShield.""" + if self.swarm_shield: + return self.swarm_shield.validate_agent_config(agent_config, agent_name) + return True, agent_config, None + + def process_agent_communication_with_shield(self, agent_name: str, message: str, direction: str = "outbound") -> Tuple[bool, str, Optional[str]]: + """Process agent communication through SwarmShield security.""" + if self.swarm_shield: + return self.swarm_shield.process_agent_communication(agent_name, message, direction) + return True, message, None + + def check_rate_limit_with_shield(self, agent_name: str, request_size: int = 1) -> Tuple[bool, Optional[str]]: + """Check rate limits for an agent using SwarmShield.""" + if self.swarm_shield: + return self.swarm_shield.check_rate_limit(agent_name, request_size) + return True, None + + def add_secure_message(self, conversation_id: str, agent_name: str, message: str) -> bool: + """Add a message to secure conversation history.""" + if self.swarm_shield: + return self.swarm_shield.add_secure_message(conversation_id, agent_name, message) + return False + + def get_secure_messages(self, conversation_id: str) -> List[Tuple[str, str, datetime]]: + """Get secure conversation messages.""" + if self.swarm_shield: + return self.swarm_shield.get_secure_messages(conversation_id) + return [] + + def create_secure_conversation(self, name: str = "") -> Optional[str]: + """Create a secure conversation.""" + if self.swarm_shield: + return self.swarm_shield.create_secure_conversation(name) + return None + + def filter_and_protect_output(self, output_data: Union[str, dict, List], agent_name: str, output_type: str = "text") -> Tuple[bool, Union[str, dict, List], Optional[str]]: + """Filter and protect output data using SwarmShield.""" + if self.swarm_shield: + return self.swarm_shield.filter_and_protect_output(output_data, agent_name, output_type) + return True, output_data, None + + 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) + logger.info("SwarmShield configuration updated") + + 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) + logger.info("SwarmShield security enabled") + + def disable_security(self) -> None: + """Disable SwarmShield security features.""" + self.swarm_shield = None + self.enable_security = False + logger.info("SwarmShield security disabled") + + def cleanup_security(self) -> None: + """Clean up SwarmShield resources.""" + if self.swarm_shield: + self.swarm_shield.cleanup() + logger.info("SwarmShield resources cleaned up") + def sequential_flow(self): # Only create flow if agents exist if self.agents: