diff --git a/swarms/structs/council_judge.py b/swarms/structs/council_judge.py index 39d47d8a..4d1bedee 100644 --- a/swarms/structs/council_judge.py +++ b/swarms/structs/council_judge.py @@ -2,7 +2,7 @@ import os import uuid from concurrent.futures import ThreadPoolExecutor, as_completed from functools import lru_cache -from typing import Dict, Optional, Tuple +from typing import Dict, Optional, Tuple, List from loguru import logger @@ -12,6 +12,7 @@ from swarms.structs.ma_utils import set_random_models_for_agents from swarms.utils.history_output_formatter import ( history_output_formatter, ) +from swarms.security import SwarmShieldIntegration, ShieldConfig class EvaluationError(Exception): @@ -225,11 +226,12 @@ def build_aggregation_prompt(rationales: Dict[str, str]) -> str: class CouncilAsAJudge: """ - A council of AI agents that evaluates task responses across multiple dimensions. + A council of specialized judge agents that evaluate task responses across multiple dimensions. - This class implements a parallel evaluation system where multiple specialized agents - evaluate different aspects of a task response, and their findings are aggregated - into a comprehensive report. + This class creates a system of expert evaluators, each focusing on a specific quality dimension + (accuracy, helpfulness, harmlessness, coherence, conciseness, instruction adherence). The judges + work in parallel to provide comprehensive assessments, which are then aggregated into a final + evaluation. Attributes: id (str): Unique identifier for the council @@ -237,10 +239,15 @@ class CouncilAsAJudge: description (str): Description of the council's purpose model_name (str): Name of the model to use for evaluations output_type (str): Type of output to return + cache_size (int): Size of the LRU cache for prompts + random_model_name (bool): Whether to use random model names + max_loops (int): Maximum number of loops for agents + aggregation_model_name (str): Model name for the aggregator agent + judge_agent_model_name (Optional[str]): Model name for judge agents + max_workers (int): Maximum number of worker threads for parallel execution judge_agents (Dict[str, Agent]): Dictionary of dimension-specific judge agents aggregator_agent (Agent): Agent responsible for aggregating evaluations conversation (Conversation): Conversation history tracker - max_workers (int): Maximum number of worker threads for parallel execution """ def __init__( @@ -255,6 +262,9 @@ class CouncilAsAJudge: max_loops: int = 1, aggregation_model_name: str = "gpt-4o-mini", judge_agent_model_name: Optional[str] = None, + shield_config: Optional[ShieldConfig] = None, + enable_security: bool = True, + security_level: str = "standard", ): """ Initialize the CouncilAsAJudge. @@ -270,6 +280,10 @@ class CouncilAsAJudge: random_model_name (bool): Whether to use random model names max_loops (int): Maximum number of loops for agents aggregation_model_name (str): Model name for the aggregator agent + judge_agent_model_name (Optional[str]): Model name for judge agents + 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 @@ -283,6 +297,9 @@ class CouncilAsAJudge: self.judge_agent_model_name = judge_agent_model_name self.max_workers = max(1, int(os.cpu_count() * 0.75)) + # Initialize SwarmShield integration + self._initialize_swarm_shield(shield_config, enable_security, security_level) + self.reliability_check() self.judge_agents = self._create_judges() @@ -499,3 +516,89 @@ class CouncilAsAJudge: raise EvaluationError( f"Evaluation process failed: {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) + 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) -> 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) + 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")