scripts/adapters/b3_jailbreakbench_adapter.py

"""
JailbreakBench Dataset Adapter for B3 (Direct Query) Schema

JailbreakBench: An Open Robustness Benchmark for Jailbreaking LLMs
License: MIT
Source: https://github.com/JailbreakBench/jailbreakbench
HuggingFace: https://huggingface.co/datasets/JailbreakBench/JBB-Behaviors
Paper: NeurIPS 2024 Datasets and Benchmarks Track

Maps JailbreakBench behaviors to the B3 Direct Query Safety schema.
"""

import logging
from pathlib import Path
from typing import Dict, Any, Optional, Generator

try:
    from datasets import load_dataset
    HAS_DATASETS = True
except ImportError:
    HAS_DATASETS = False

from .base_adapter import BaseAdapter, DatasetMetadata

logger = logging.getLogger(__name__)


class JailbreakBenchAdapter(BaseAdapter):
    """
    Adapter for JailbreakBench dataset.

    JailbreakBench provides 100 harmful and 100 benign behaviors. This adapter
    ingests the `harmful` split (100 records) for red teaming.
    """

    CATEGORY_MAPPING = {
        'Harassment/Discrimination': 'Harassment',
        'Malware/Hacking': 'Cybercrime',
        'Physical harm': 'Violence',
        'Economic harm': 'Economic Harm',
        'Fraud/Deception': 'Fraud',
        'Disinformation': 'Misinformation',
        'Sexual/Adult content': 'Adult Content',
        'Privacy': 'Privacy Violation',
        'Expert advice': 'Unauthorized Advice',
        'Government decision-making': 'Political Manipulation'
    }

    def __init__(
        self,
        output_dir: Path,
        verbose: bool = False
    ):
        """
        Initialize JailbreakBench adapter.

        Args:
            output_dir: Directory for output files
            verbose: Enable verbose logging
        """
        super().__init__(output_dir, verbose)

        if not HAS_DATASETS:
            raise ImportError(
                "The 'datasets' library is required. Install with: pip install datasets"
            )

    @property
    def metadata(self) -> DatasetMetadata:
        return DatasetMetadata(
            name="JailbreakBench",
            version="1.0",
            license="MIT",
            license_url="https://github.com/JailbreakBench/jailbreakbench/blob/main/LICENSE",
            source_url="https://github.com/JailbreakBench/jailbreakbench",
            paper_url="https://arxiv.org/abs/2404.01318",
            can_redistribute=True,
            target_surface="B3",
            description="100 harmful behaviors from JailbreakBench for jailbreak evaluation"
        )

    def load_source_data(self) -> Generator[Dict[str, Any], None, None]:
        """Load JailbreakBench data from HuggingFace."""
        logger.info("Loading JailbreakBench from HuggingFace...")

        try:
            # JailbreakBench uses 'behaviors' config with 'harmful' and 'benign' splits
            # For red teaming, we want the 'harmful' split
            dataset = load_dataset(
                "JailbreakBench/JBB-Behaviors",
                "behaviors",
                split="harmful"
            )

            for item in dataset:
                yield dict(item)

        except Exception as e:
            logger.error(f"Failed to load JailbreakBench: {e}")
            raise

    def transform_record(self, record: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """Transform JailbreakBench record to B3 schema."""
        # JailbreakBench has 'Goal' and 'Behavior' fields
        goal = record.get('Goal', record.get('goal', ''))
        behavior = record.get('Behavior', record.get('behavior', ''))
        category = record.get('Category', record.get('category', 'General'))
        source = record.get('Source', record.get('source', 'JailbreakBench'))
        behavior_id = record.get('Index', record.get('index', ''))

        if not goal and not behavior:
            logger.warning("Empty goal/behavior in record")
            return None

        # Use behavior as query if available, otherwise goal
        query = behavior if behavior else goal
        adversarial_goal = goal if goal else behavior

        # Map category
        attack_category = self.CATEGORY_MAPPING.get(category, 'General Harm')

        record_id = self.generate_id("B3-JailbreakBench", query)

        return {
            "id": record_id,
            "attack_category": attack_category,
            "attack_subtype": f"JailbreakBench_{category.replace('/', '_').replace(' ', '_')}",
            "adversarial_goal": adversarial_goal,
            "user_query": query,
            "target_llm_output": "",
            "judge_assessment": {
                "assessment": "Pending",
                "justification": "Imported from JailbreakBench - requires evaluation",
                "confidence": "N/A"
            },
            "benign_content": None,
            "mode": None,
            "role_play": None,
            "manipulation_strategy": None,
            "communication_style": None,
            "query_structure": "direct question",
            "intent_Opacity": "explicit",
            "topic_domain": category,
            "temporal_context": None,
            "language": "English",
            "Obfuscation": None,
            "adversarial_suffix": None,
            "payload_splitting": None,
            "encoding": None,
            "target_rag_component": "Generator_LLM",
            "target_component_model": "External_Evaluation",
            "target_model_system_prompt": "",
            "all_llm_details": {},
            "agent_architecture_description": "JailbreakBench standardized evaluation",
            "attack_gen_pipeline_description": f"JailbreakBench curated behaviors (source: {source})",
            "generation_timestamp": self.get_timestamp(),
            "script_version": self.BENCHMARK_VERSION,
            "source_dataset": "JailbreakBench",
            "source_id": str(behavior_id) if behavior_id else None,
            "source_license": "MIT"
        }


def main():
    """CLI entry point for JailbreakBench adapter."""
    import argparse

    parser = argparse.ArgumentParser(description="JailbreakBench to B3 Schema Adapter")
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("data/external_augmented/core"),
        help="Output directory"
    )
    parser.add_argument(
        "--max-records",
        type=int,
        default=None,
        help="Maximum records to process"
    )
    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="Enable verbose output"
    )

    args = parser.parse_args()

    adapter = JailbreakBenchAdapter(
        output_dir=args.output_dir,
        verbose=args.verbose
    )

    result = adapter.run(max_records=args.max_records)

    print(f"\nAdapter Results:")
    print(f"  Success: {result.success}")
    print(f"  Records processed: {result.records_processed}")
    print(f"  Records failed: {result.records_failed}")
    if result.output_file:
        print(f"  Output file: {result.output_file}")


if __name__ == "__main__":
    main()