first commit

jailbreak_toolbox/core/async_orchestrator.py

@@ -0,0 +1,469 @@
+"""
+AsyncIO-based Orchestrator for Parallel Attack Execution
+
+This orchestrator provides high-performance parallel execution of attacks using
+asyncio instead of threading, eliminating resource contention and providing
+better concurrency for I/O-bound API operations.
+
+Features:
+- AsyncIO-based parallel execution with semaphore rate limiting
+- Factory pattern for fresh attack instances
+- Model-specific and query-specific log directories
+- Real-time progress tracking with tqdm
+- Comprehensive result tracking and evaluation
+"""
+
+import asyncio
+import json
+import os
+import time
+import traceback
+from dataclasses import asdict
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+from tqdm import tqdm
+
+try:
+    from tqdm.asyncio import tqdm as asyncio_tqdm
+except ImportError:
+    # Fallback if tqdm.asyncio is not available
+    from tqdm import tqdm
+    asyncio_tqdm = None
+
+from ..models.base_model import BaseModel
+from ..datasets.base_dataset import BaseDataset
+from ..attacks.base_attack import BaseAttack, AttackResult
+from ..evaluators.base_evaluator import BaseEvaluator, EvaluationMetrics
+
+
+class AsyncOrchestrator:
+    """
+    AsyncIO-based orchestrator for high-performance parallel attack execution.
+
+    Unlike the traditional threading-based Orchestrator, this class:
+    - Uses asyncio for true parallel I/O operations
+    - Creates fresh attack instances for each query (factory pattern)
+    - Provides model-specific and query-specific log directories
+    - Real-time progress tracking with tqdm
+    - Avoids resource contention and deadlocks
+    - Provides better performance for API-heavy workloads
+    """
+
+    def __init__(
+        self,
+        model: BaseModel,
+        dataset: BaseDataset,
+        attack_class: BaseAttack,
+        evaluator: BaseEvaluator,
+        max_concurrent_queries: int = 10,
+        base_logs_dir: str = "./async_logs",
+        enable_progress_bars: bool = True,
+        model_name: str = None,
+        attack_name: str = None,
+        attack_kwargs: dict = None
+    ):
+        """
+        Initialize the async orchestrator.
+
+        Args:
+            model: Target model to attack
+            dataset: Dataset of queries to attack
+            attack_class: Attack class (not instance - will be instantiated per query)
+            evaluator: Evaluator for results
+            max_concurrent_queries: Maximum number of concurrent queries (controls API rate limiting)
+            base_logs_dir: Base directory for all logs
+            enable_progress_bars: Whether to show tqdm progress bars
+            model_name: Custom name for model logs (defaults to model_name if provided)
+            attack_name: Custom name for attack logs (defaults to attack_name if provided)
+            attack_kwargs: Additional keyword arguments for attack initialization
+        """
+        self.model = model
+        self.dataset = dataset
+        self.attack_class = attack_class
+        self.attack_kwargs = attack_kwargs or {}
+        self.evaluator = evaluator
+        self.max_concurrent_queries = max_concurrent_queries
+        self.base_logs_dir = Path(base_logs_dir)
+        self.enable_progress_bars = enable_progress_bars
+        self.model_name = model_name or getattr(model, 'model_name', model.__class__.__name__)
+        self.attack_name = attack_name or getattr(attack_class, '__name__', 'Attack')
+
+        # Create base logs directory
+        self.base_logs_dir.mkdir(parents=True, exist_ok=True)
+
+        # Create model-specific log directory
+        model_name_safe = "".join(c for c in self.model_name if c.isalnum() or c in ('-', '_')).strip()
+        self.model_logs_dir = self.base_logs_dir / model_name_safe
+        self.model_logs_dir.mkdir(parents=True, exist_ok=True)
+
+    def _create_query_specific_log_dir(self, query_index: int) -> Path:
+        """
+        Create a query-specific log directory.
+
+        Directory structure: {base_logs_dir}/{model_name}/{attack_name}/query_{index:03d}/
+
+        Args:
+            query_index: Index of query
+
+        Returns:
+            Path to query-specific log directory
+        """
+        # Create safe attack name
+        attack_name_safe = "".join(c for c in self.attack_name if c.isalnum() or c in ('-', '_')).strip()
+
+        # Create query-specific log directory
+        query_log_dir = self.model_logs_dir / attack_name_safe / f"query_{query_index:03d}"
+        query_log_dir.mkdir(parents=True, exist_ok=True)
+
+        return query_log_dir
+
+    async def run_single_attack_async(self, query: str, query_index: int, pbar: Optional[tqdm] = None) -> Tuple[int, str, AttackResult, str]:
+        """
+        Run a single attack asynchronously with individual logging.
+
+        Args:
+            query: The target query to attack
+            query_index: Index of query for tracking
+            pbar: Optional tqdm progress bar to update
+
+        Returns:
+            Tuple of (query_index, query, attack_result, query_log_dir)
+        """
+        try:
+            # Update progress bar description
+            if pbar:
+                pbar.set_description(f"Query {query_index + 1}: {query[:30]}...")
+                pbar.set_postfix({"Status": "Starting"})
+
+            # Create query-specific log directory
+            query_log_dir = self._create_query_specific_log_dir(query_index)
+
+            # Save query metadata
+            query_metadata = {
+                "query_index": query_index,
+                "query": query,
+                "timestamp": datetime.now().isoformat(),
+                "model_name": self.model_name,
+                "attack_name": self.attack_name,
+                "evaluator_class": self.evaluator.__class__.__name__,
+                "model_logs_dir": str(self.model_logs_dir)
+            }
+
+            metadata_path = Path(query_log_dir) / "query_metadata.json"
+            with open(metadata_path, "w", encoding="utf-8") as f:
+                json.dump(query_metadata, f, ensure_ascii=False, indent=2)
+
+            # Create fresh attack instance (Factory Pattern)
+            attack_instance = self._create_attack_instance(query_log_dir)
+
+            # Update progress bar
+            if pbar:
+                pbar.set_postfix({"Status": "Running"})
+
+            # Run attack asynchronously
+            start_time = time.time()
+            result = await self._run_attack_async(attack_instance, query)
+            duration = time.time() - start_time
+
+            # Update progress bar based on result
+            if pbar:
+                status = "✅ Success" if result.success else "❌ Failed"
+                pbar.set_postfix({"Status": status, "Duration": f"{duration:.1f}s"})
+
+            # Save attack result metadata
+            result_metadata = {
+                "query_index": query_index,
+                "query": query,
+                "attack_success": result.success,
+                "final_prompt": result.final_prompt,
+                "output_length": len(result.output_text),
+                "duration_seconds": duration,
+                "method": result.method,
+                "cost": result.cost,
+                "timestamp": datetime.now().isoformat()
+            }
+
+            result_path = Path(query_log_dir) / "attack_result.json"
+            with open(result_path, "w", encoding="utf-8") as f:
+                json.dump(result_metadata, f, ensure_ascii=False, indent=2)
+
+            # Save detailed conversation history
+            if result.history:
+                history_path = Path(query_log_dir) / "conversation_history.json"
+                with open(history_path, "w", encoding="utf-8") as f:
+                    json.dump(result.history, f, ensure_ascii=False, indent=2)
+
+            print(f"✅ Finished attack {query_index + 1} in {duration:.1f}s - Success: {result.success}")
+            print(f"📁 Logs saved to: {query_log_dir}")
+
+            return query_index, query, result, str(query_log_dir)
+
+        except Exception as e:
+            # Update progress bar for error
+            if pbar:
+                pbar.set_postfix({"Status": "💥 Error", "Error": str(e)[:20]})
+            print(f"❌ Failed attack {query_index + 1}: {e}")
+
+            # Create failure result with error logging
+            failure_result = AttackResult(
+                target=query,
+                success=False,
+                final_prompt=query,
+                output_text="",
+                method="async_orchestrator",
+                history=[{"error": str(e), "traceback": traceback.format_exc()}]
+            )
+
+            # Try to save error info if log directory was created
+            error_log_dir = None
+            try:
+                if 'query_log_dir' in locals():
+                    error_log_dir = query_log_dir
+                    error_metadata = {
+                        "query_index": query_index,
+                        "query": query,
+                        "error": str(e),
+                        "traceback": traceback.format_exc(),
+                        "timestamp": datetime.now().isoformat()
+                    }
+                    error_path = Path(error_log_dir) / "error.json"
+                    with open(error_path, "w", encoding="utf-8") as f:
+                        json.dump(error_metadata, f, ensure_ascii=False, indent=2)
+            except Exception as save_error:
+                print(f"⚠️  Could not save error log: {save_error}")
+
+            return query_index, query, failure_result, error_log_dir or ""
+
+    def _create_attack_instance(self, query_log_dir: Path):
+        """
+        Create a fresh attack instance for specific query.
+
+        This method should be overridden by subclasses or attack should
+        support passing a logs_dir parameter during initialization.
+
+        Args:
+            query_log_dir: Directory for query-specific logging
+
+        Returns:
+            Attack instance configured for query
+        """
+        # Default implementation - create instance with model and attack_kwargs
+        # Override if attack needs special configuration
+        try:
+            # Try to pass logs_dir to attack constructor if it accepts it
+            import inspect
+            attack_sig = inspect.signature(self.attack_class.__init__)
+
+            # Base initialization kwargs with model
+            init_kwargs = {'model': self.model}
+
+            # Add logs_dir if attack constructor accepts it
+            if 'logs_dir' in attack_sig.parameters:
+                init_kwargs['logs_dir'] = str(query_log_dir)
+
+            # Add all attack_kwargs, excluding any that conflict with required params
+            if self.attack_kwargs:
+                for key, value in self.attack_kwargs.items():
+                    if key != 'model' and key not in init_kwargs:  # Don't override model/logs_dir
+                        init_kwargs[key] = value
+
+            return self.attack_class(**init_kwargs)
+        except Exception as e:
+            # Fallback to basic instantiation
+            try:
+                return self.attack_class(model=self.model)
+            except Exception as fallback_e:
+                # Final fallback - try without any parameters
+                return self.attack_class()
+
+    async def _run_attack_async(self, attack_instance: BaseAttack, query: str) -> AttackResult:
+        """
+        Run attack asynchronously.
+
+        This method should be overridden if the attack supports async execution.
+        Default implementation runs the synchronous attack in a thread executor.
+
+        Args:
+            attack_instance: Attack instance to run
+            query: Query to attack
+
+        Returns:
+            Attack result
+        """
+        # Default implementation - run sync attack in thread pool
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(None, attack_instance.attack, query)
+
+    async def run_parallel_attacks(self) -> Tuple[List[AttackResult], Dict]:
+        """
+        Run multiple attacks in parallel using asyncio.gather with tqdm progress.
+
+        Returns:
+            Tuple of (results, performance_metrics)
+        """
+        print(f"🎯 Starting {len(self.dataset)} parallel attacks...")
+        print(f"🔧 Max concurrent queries: {self.max_concurrent_queries}")
+        print(f"📁 Model logs directory: {self.model_logs_dir}")
+        print(f"⚙️  Attack: {self.attack_name}")
+        print(f"📊 Evaluator: {self.evaluator.__class__.__name__}")
+        print()
+
+        # Create semaphore to limit concurrent operations
+        semaphore = asyncio.Semaphore(self.max_concurrent_queries)
+
+        async def run_with_semaphore(query_idx_query, pbar=None):
+            query_idx, query = query_idx_query
+            async with semaphore:
+                return await self.run_single_attack_async(query, query_idx, pbar)
+
+        start_time = time.time()
+ 
+        # Prepare tasks
+        tasks_with_indices = [(i, query) for i, query in enumerate(self.dataset)]
+
+        # Create progress bar if enabled
+        if self.enable_progress_bars and asyncio_tqdm:
+            with asyncio_tqdm(total=len(tasks_with_indices), desc="Processing Queries",
+                              unit="query", ncols=100, colour='green') as pbar:
+
+                # Run all attacks concurrently with rate limiting
+                tasks = [run_with_semaphore(task, pbar) for task in tasks_with_indices]
+                attack_results = await asyncio.gather(*tasks, return_exceptions=True)
+        else:
+            # Run without progress bars (fallback if asyncio_tqdm not available)
+            tasks = [run_with_semaphore(task, None) for task in tasks_with_indices]
+            attack_results = await asyncio.gather(*tasks, return_exceptions=True)
+
+        total_duration = time.time() - start_time
+
+        # Process results and maintain order
+        results = [None] * len(self.dataset)
+        query_log_dirs = [None] * len(self.dataset)
+        successful_attacks = 0
+        failed_attacks = 0
+
+        # Process results with optional progress bar
+        if self.enable_progress_bars:
+            with tqdm(total=len(attack_results), desc="Processing Results",
+                     unit="result", ncols=100, colour='blue') as summary_pbar:
+
+                for attack_result in attack_results:
+                    success, failed = self._process_single_result(
+                        attack_result, results, query_log_dirs
+                    )
+                    successful_attacks += success
+                    failed_attacks += failed
+
+                    summary_pbar.set_postfix({
+                        "Success": f"{successful_attacks}/{len(self.dataset)}",
+                        "Failed": f"{failed_attacks}/{len(self.dataset)}"
+                    })
+                    summary_pbar.update(1)
+        else:
+            for attack_result in attack_results:
+                success, failed = self._process_single_result(
+                    attack_result, results, query_log_dirs
+                )
+                successful_attacks += success
+                failed_attacks += failed
+
+        # Performance metrics
+        performance_metrics = {
+            "total_queries": len(self.dataset),
+            "successful_attacks": successful_attacks,
+            "failed_attacks": failed_attacks,
+            "success_rate": successful_attacks / len(self.dataset) if self.dataset else 0.0,
+            "total_duration": total_duration,
+            "avg_duration_per_attack": total_duration / len(self.dataset) if self.dataset else 0.0,
+            "max_concurrent_queries": self.max_concurrent_queries,
+            "model_name": self.model_name,
+            "attack_name": self.attack_name,
+            "evaluator_class": self.evaluator.__class__.__name__,
+            "model_logs_dir": str(self.model_logs_dir),
+            "query_log_directories": query_log_dirs
+        }
+
+        print(f"\n📊 Parallel Attack Summary:")
+        print(f"   Target model: {self.model_name}")
+        print(f"   Total queries: {len(self.dataset)}")
+        print(f"   Successful: {successful_attacks} 🎯")
+        print(f"   Failed: {failed_attacks} ❌")
+        print(f"   Success rate: {performance_metrics['success_rate']:.2%}")
+        print(f"   Total time: {total_duration:.1f}s")
+        print(f"   Avg per attack: {performance_metrics['avg_duration_per_attack']:.1f}s")
+        print(f"   All logs in: {self.model_logs_dir}")
+
+        return results, performance_metrics
+
+    def _process_single_result(self, attack_result, results: List, query_log_dirs: List) -> Tuple[int, int]:
+        """
+        Process a single attack result and update counters.
+
+        Args:
+            attack_result: Result from a single attack
+            results: List to populate with results
+            query_log_dirs: List to populate with log directories
+
+        Returns:
+            Tuple of (successful_count, failed_count)
+        """
+        if isinstance(attack_result, Exception):
+            print(f"⚠️  Exception in attack: {attack_result}")
+            return 0, 1
+
+        query_index, query, result, log_dir = attack_result
+        results[query_index] = result
+        query_log_dirs[query_index] = log_dir
+
+        if result.success:
+            return 1, 0
+        else:
+            return 0, 1
+
+    async def run(self) -> Tuple[EvaluationMetrics, List[AttackResult]]:
+        """
+        Execute complete experiment workflow asynchronously.
+
+        Returns:
+            Tuple of (evaluation_metrics, detailed_results)
+        """
+        # Print experiment configuration
+        print("--- Starting Async Experiment ---")
+        print(f"Model: {self.model_name}")
+        print(f"Dataset: {self.dataset.__class__.__name__} (Size: {len(self.dataset)})")
+        print(f"Attack: {self.attack_name}")
+        print(f"Evaluator: {self.evaluator.__class__.__name__}")
+        print(f"Max concurrent queries: {self.max_concurrent_queries}")
+        print(f"Logs directory: {self.base_logs_dir}")
+        print("--------------------------------")
+
+        # Run parallel attacks
+        results, performance_metrics = await self.run_parallel_attacks()
+
+        # Evaluate results
+        print("\n--- Evaluating Async Experiment Results ---")
+        metrics = self.evaluator.evaluate(results)
+
+        # Add performance metrics to evaluation if supported
+        if isinstance(metrics, EvaluationMetrics):
+            metrics.async_performance = performance_metrics
+            metrics.insucess_rate = performance_metrics['success_rate']
+            print(f"Async Performance:")
+            print(f"  Total queries: {performance_metrics['total_queries']}")
+            print(f"  Success rate: {performance_metrics['success_rate']:.2%}")
+            print(f"  Total duration: {performance_metrics['total_duration']:.1f}s")
+            print(f"  Avg per attack: {performance_metrics['avg_duration_per_attack']:.1f}s")
+        elif isinstance(metrics, dict) and "overall" in metrics and isinstance(metrics["overall"], EvaluationMetrics):
+            metrics["overall"].async_performance = performance_metrics
+            metrics['overall'].insucess_rate = performance_metrics['success_rate']
+            print(f"Async Performance:")
+            print(f"  Total queries: {performance_metrics['total_queries']}")
+            print(f"  Success rate: {performance_metrics['success_rate']:.2%}")
+            print(f"  Total duration: {performance_metrics['total_duration']:.1f}s")
+            print(f"  Avg per attack: {performance_metrics['avg_duration_per_attack']:.1f}s")
+        print("Metric to save: ",metrics)
+        print("--- Async Experiment Finished ---\n")
+
+        return metrics, results

jailbreak_toolbox/core/registry.py

@@ -0,0 +1,34 @@
+from typing import Any, Callable, Dict, Type, List
+
+class Registry:
+    def __init__(self, name: str):
+        self._name = name
+        self._registry: Dict[str, Any] = {}
+
+    def register(self, name: str) -> Callable:
+        def decorator(cls: Type) -> Type:
+            if name in self._registry:
+                print(f"Warning: '{name}' is being overwritten in the {self._name} registry.")
+            self._registry[name] = cls
+            return cls
+        return decorator
+
+    def get(self, name: str) -> Any:
+        if name not in self._registry:
+            available = list(self._registry.keys())
+            raise ValueError(f"'{name}' not found in {self._name} registry. Available: {available}")
+        return self._registry[name]
+
+    def list(self) -> List[str]:
+        return list(self._registry.keys())
+
+# --- 全局注册表实例 ---
+attack_registry = Registry("Attacks")
+model_registry = Registry("Models")
+dataset_registry = Registry("Datasets")
+evaluator_registry = Registry("Evaluators")
+
+# --- 策略专属注册表 ---
+judge_registry = Registry("Judges")
+advancer_registry = Registry("Advancers")
+propagator_registry = Registry("Propagators")