scripts/adapters/b4_agentharm_adapter.py

"""
AgentHarm Dataset Adapter for B4 (Orchestrator) Schema

AgentHarm: A Benchmark for Measuring Harmfulness of LLM Agents
License: MIT + Research Use Clause (DO NOT REDISTRIBUTE)
HuggingFace: https://huggingface.co/datasets/ai-safety-institute/AgentHarm
Paper: ICLR 2025

IMPORTANT: This dataset has a usage restriction clause prohibiting use
for purposes other than improving AI safety and security. Do not redistribute.

Maps AgentHarm behaviors to the B4 Orchestrator Attack 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 AgentHarmAdapter(BaseAdapter):
    """
    Adapter for AgentHarm dataset.

    AgentHarm provides 110 unique malicious agent behaviors (440 with augmentations)
    across 11 harm categories for evaluating LLM agent safety.

    WARNING: This is an Extended dataset. Do not redistribute the converted data.
    The adapter exists for research purposes only.
    """

    CATEGORY_MAPPING = {
        'fraud': 'Fraud',
        'cybercrime': 'Cybercrime',
        'harassment': 'Harassment',
        'illegal_activities': 'Illegal Activities',
        'violence': 'Violence',
        'self_harm': 'Self Harm',
        'sexual_content': 'Adult Content',
        'privacy_violation': 'Privacy Violation',
        'misinformation': 'Misinformation',
        'financial_harm': 'Economic Harm',
        'weapons': 'Weapons/Violence'
    }

    def __init__(
        self,
        output_dir: Path,
        split: str = "test_public",
        verbose: bool = False
    ):
        """
        Initialize AgentHarm adapter.

        Args:
            output_dir: Directory for output files
            split: Dataset split to use ('test_public' or 'validation')
            verbose: Enable verbose logging
        """
        super().__init__(output_dir, verbose)
        self.split = split

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

        logger.warning(
            "AgentHarm has a research-use-only clause. "
            "Do not redistribute converted data. Use for evaluation only."
        )

    @property
    def metadata(self) -> DatasetMetadata:
        return DatasetMetadata(
            name="AgentHarm",
            version="1.0",
            license="MIT+Research-Clause",
            license_url="https://huggingface.co/datasets/ai-safety-institute/AgentHarm",
            source_url="https://huggingface.co/datasets/ai-safety-institute/AgentHarm",
            paper_url="https://arxiv.org/abs/2410.09024",
            can_redistribute=False,  # IMPORTANT: Cannot redistribute
            target_surface="B4",
            description="110 malicious agent behaviors for LLM agent safety evaluation (research only)"
        )

    def load_source_data(self) -> Generator[Dict[str, Any], None, None]:
        """Load AgentHarm data from HuggingFace."""
        logger.info(f"Loading AgentHarm dataset (split: {self.split}) from HuggingFace...")
        logger.warning("Remember: Do not redistribute this data!")

        try:
            # AgentHarm has configs: 'harmful', 'harmless_benign', 'chat'
            # We want the 'harmful' config for red teaming
            dataset = load_dataset(
                "ai-safety-institute/AgentHarm",
                "harmful",
                split=self.split
            )

            for item in dataset:
                yield dict(item)

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

    def transform_record(self, record: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """Transform AgentHarm record to B4 schema."""
        # AgentHarm has 'prompt', 'behavior_category', and 'detailed_behavior' fields
        prompt = record.get('prompt', record.get('Prompt', ''))
        behavior_category = record.get('behavior_category', record.get('Category', 'general'))
        detailed_behavior = record.get('detailed_behavior', record.get('DetailedBehavior', ''))
        behavior_id = record.get('id', record.get('ID', ''))

        if not prompt:
            logger.warning("Empty prompt in record")
            return None

        prompt = self.sanitize_text(prompt, max_length=15000)

        # Map category
        attack_category = self.CATEGORY_MAPPING.get(
            behavior_category.lower().replace(' ', '_'),
            'Agent Misuse'
        )

        record_id = self.generate_id("B4-AgentHarm", prompt)

        return {
            "attack_id": record_id,
            "category": "LLM06",  # Agent misuse maps to OWASP LLM06
            "subtype": f"AgentHarm_{behavior_category.replace(' ', '_')}",
            "benign_query": "N/A (Direct Agent Attack)",
            "malicious_injection": "N/A",
            "combined_query": prompt,
            "expected_tool": "Any",
            "simulated_tool": "Agent_Execution_Environment",
            "success": True,  # AgentHarm provides attack scenarios
            "judge_note": f"AgentHarm benchmark - Category: {behavior_category}. Behavior: {detailed_behavior[:200] if detailed_behavior else 'N/A'}",
            "models_used": {
                "crafter": "AgentHarm_Benchmark",
                "simulator": "External_Evaluation"
            },
            "script_version": self.BENCHMARK_VERSION,
            "record_id": record_id,
            "timestamp": self.get_timestamp(),
            "source_dataset": "AgentHarm",
            "source_id": str(behavior_id) if behavior_id else None,
            "source_license": "MIT+Research-Clause (DO NOT REDISTRIBUTE)",
            "research_only": True,
            "behavior_category_detail": behavior_category,
            "detailed_behavior": detailed_behavior
        }


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

    parser = argparse.ArgumentParser(description="AgentHarm to B4 Schema Adapter (Research Only)")
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("data/external_augmented/extended"),
        help="Output directory"
    )
    parser.add_argument(
        "--split",
        type=str,
        default="test_public",
        help="Dataset split to use ('test_public' or 'validation')"
    )
    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()

    print("\n" + "="*60)
    print("WARNING: AgentHarm has a research-use-only clause.")
    print("Do not redistribute the converted data.")
    print("="*60 + "\n")

    adapter = AgentHarmAdapter(
        output_dir=args.output_dir,
        split=args.split,
        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}")
    print("\nREMINDER: This data is for research only. Do not redistribute.")


if __name__ == "__main__":
    main()