oxideshield_guard/benchmark/

datasets.rs

//! Benchmark Datasets
//!
//! Standardized datasets for evaluating LLM security tools.
//! Includes both attack datasets and benign datasets for measuring
//! detection rate and false positive rate.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// A benchmark sample with ground truth label
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchmarkSample {
    /// Unique identifier
    pub id: String,
    /// The input text/prompt
    pub text: String,
    /// Ground truth: true if this is an attack, false if benign
    pub is_attack: bool,
    /// Attack category (if applicable)
    pub category: Option<String>,
    /// Subcategory or attack type
    pub subcategory: Option<String>,
    /// Source dataset
    pub source: String,
    /// Additional metadata
    pub metadata: HashMap<String, String>,
}

/// A benchmark dataset
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchmarkDataset {
    /// Dataset name
    pub name: String,
    /// Dataset version
    pub version: String,
    /// Description
    pub description: String,
    /// Samples
    pub samples: Vec<BenchmarkSample>,
    /// Source URL or reference
    pub source: String,
}

impl BenchmarkDataset {
    /// Create a new empty dataset
    pub fn new(name: impl Into<String>, version: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            version: version.into(),
            description: String::new(),
            samples: Vec::new(),
            source: String::new(),
        }
    }

    /// Add description
    pub fn with_description(mut self, desc: impl Into<String>) -> Self {
        self.description = desc.into();
        self
    }

    /// Add source
    pub fn with_source(mut self, source: impl Into<String>) -> Self {
        self.source = source.into();
        self
    }

    /// Add a sample
    pub fn add_sample(mut self, sample: BenchmarkSample) -> Self {
        self.samples.push(sample);
        self
    }

    /// Add multiple samples
    pub fn add_samples(mut self, samples: Vec<BenchmarkSample>) -> Self {
        self.samples.extend(samples);
        self
    }

    /// Get attack samples only
    pub fn attack_samples(&self) -> Vec<&BenchmarkSample> {
        self.samples.iter().filter(|s| s.is_attack).collect()
    }

    /// Get benign samples only
    pub fn benign_samples(&self) -> Vec<&BenchmarkSample> {
        self.samples.iter().filter(|s| !s.is_attack).collect()
    }

    /// Total sample count
    pub fn len(&self) -> usize {
        self.samples.len()
    }

    /// Check if empty
    pub fn is_empty(&self) -> bool {
        self.samples.is_empty()
    }

    /// Get samples by category
    pub fn samples_by_category(&self, category: &str) -> Vec<&BenchmarkSample> {
        self.samples
            .iter()
            .filter(|s| s.category.as_deref() == Some(category))
            .collect()
    }
}

/// Create the OxideShield standard benchmark dataset
/// Combines attack samples and benign samples from the adversarial module.
/// Requires the `adversarial-samples` feature.
#[cfg(feature = "adversarial-samples")]
pub fn oxideshield_benchmark_dataset() -> BenchmarkDataset {
    use crate::adversarial::{all_attack_samples, all_benign_samples};

    let mut dataset = BenchmarkDataset::new("OxideShield Standard", "1.0.0")
        .with_description("Standard benchmark dataset for OxideShield evaluation")
        .with_source("OxideShield internal");

    // Add attack samples
    for sample in all_attack_samples() {
        dataset.samples.push(BenchmarkSample {
            id: sample.id,
            text: sample.prompt,
            is_attack: true,
            category: Some(format!("{:?}", sample.attack_type)),
            subcategory: Some(sample.description),
            source: sample.source,
            metadata: HashMap::new(),
        });
    }

    // Add benign samples
    for sample in all_benign_samples() {
        dataset.samples.push(BenchmarkSample {
            id: sample.id,
            text: sample.prompt,
            is_attack: false,
            category: Some(format!("{:?}", sample.category)),
            subcategory: Some(sample.description),
            source: "OxideShield benign".into(),
            metadata: HashMap::new(),
        });
    }

    dataset
}

/// Create a JailbreakBench-based dataset.
/// Requires the `adversarial-samples` feature.
#[cfg(feature = "adversarial-samples")]
pub fn jailbreakbench_dataset() -> BenchmarkDataset {
    use crate::adversarial::jailbreakbench_behaviors;

    let mut dataset = BenchmarkDataset::new("JailbreakBench", "1.0.0")
        .with_description("Behaviors from JailbreakBench (jailbreakbench.github.io)")
        .with_source("https://jailbreakbench.github.io/");

    for behavior in jailbreakbench_behaviors() {
        dataset.samples.push(BenchmarkSample {
            id: behavior.id,
            text: behavior.prompt,
            is_attack: true,
            category: Some(behavior.category.to_string()),
            subcategory: Some(behavior.behavior),
            source: "JailbreakBench".into(),
            metadata: HashMap::new(),
        });
    }

    dataset
}

/// Create a prompt injection focused dataset.
/// Requires the `adversarial-samples` feature.
#[cfg(feature = "adversarial-samples")]
pub fn prompt_injection_dataset() -> BenchmarkDataset {
    use crate::adversarial::{prompt_injection_samples, system_prompt_leak_samples};

    let mut dataset = BenchmarkDataset::new("Prompt Injection", "1.0.0")
        .with_description("Focused dataset for prompt injection and system prompt leak attacks")
        .with_source("OxideShield curated");

    for sample in prompt_injection_samples() {
        dataset.samples.push(BenchmarkSample {
            id: sample.id,
            text: sample.prompt,
            is_attack: true,
            category: Some("PromptInjection".into()),
            subcategory: Some(sample.description),
            source: sample.source,
            metadata: HashMap::new(),
        });
    }

    for sample in system_prompt_leak_samples() {
        dataset.samples.push(BenchmarkSample {
            id: sample.id,
            text: sample.prompt,
            is_attack: true,
            category: Some("SystemPromptLeak".into()),
            subcategory: Some(sample.description),
            source: sample.source,
            metadata: HashMap::new(),
        });
    }

    dataset
}

/// Create an adversarial suffix focused dataset (AutoDAN, GCG).
/// Requires the `adversarial-samples` feature.
#[cfg(feature = "adversarial-samples")]
pub fn adversarial_suffix_dataset() -> BenchmarkDataset {
    use crate::adversarial::{autodan_samples, gcg_samples};

    let mut dataset = BenchmarkDataset::new("Adversarial Suffixes", "1.0.0")
        .with_description("AutoDAN and GCG style adversarial suffix attacks")
        .with_source("Research paper simulations");

    for sample in autodan_samples() {
        dataset.samples.push(BenchmarkSample {
            id: sample.id,
            text: sample.prompt,
            is_attack: true,
            category: Some("AutoDAN".into()),
            subcategory: Some(sample.description),
            source: sample.source,
            metadata: HashMap::new(),
        });
    }

    for sample in gcg_samples() {
        dataset.samples.push(BenchmarkSample {
            id: sample.id,
            text: sample.prompt,
            is_attack: true,
            category: Some("GCG".into()),
            subcategory: Some(sample.description),
            source: sample.source,
            metadata: HashMap::new(),
        });
    }

    dataset
}

/// Combine multiple datasets
pub fn combined_dataset(datasets: Vec<BenchmarkDataset>) -> BenchmarkDataset {
    let mut combined = BenchmarkDataset::new("Combined", "1.0.0")
        .with_description("Combined dataset from multiple sources");

    let sources: Vec<String> = datasets.iter().map(|d| d.name.clone()).collect();
    combined.source = sources.join(", ");

    for dataset in datasets {
        for mut sample in dataset.samples {
            // Prefix ID with dataset name to avoid collisions
            sample.id = format!("{}_{}", dataset.name.replace(' ', "_"), sample.id);
            combined.samples.push(sample);
        }
    }

    combined
}

#[cfg(test)]
mod tests {
    use super::*;

    #[cfg(feature = "adversarial-samples")]
    #[test]
    fn test_oxideshield_dataset() {
        let dataset = oxideshield_benchmark_dataset();
        assert!(!dataset.is_empty());
        assert!(dataset.attack_samples().len() > 0);
        assert!(dataset.benign_samples().len() > 0);
    }

    #[cfg(feature = "adversarial-samples")]
    #[test]
    fn test_jailbreakbench_dataset() {
        let dataset = jailbreakbench_dataset();
        assert!(!dataset.is_empty());
        // All JailbreakBench samples are attacks
        assert_eq!(dataset.attack_samples().len(), dataset.len());
    }

    #[cfg(feature = "adversarial-samples")]
    #[test]
    fn test_prompt_injection_dataset() {
        let dataset = prompt_injection_dataset();
        assert!(dataset.len() >= 13); // 8 PI + 5 SL samples
    }

    #[cfg(feature = "adversarial-samples")]
    #[test]
    fn test_adversarial_suffix_dataset() {
        let dataset = adversarial_suffix_dataset();
        assert!(dataset.len() >= 6); // 3 AutoDAN + 3 GCG
    }

    #[cfg(feature = "adversarial-samples")]
    #[test]
    fn test_combined_dataset() {
        let d1 = prompt_injection_dataset();
        let d2 = adversarial_suffix_dataset();
        let combined = combined_dataset(vec![d1, d2]);
        assert!(combined.len() >= 19);
    }

    #[cfg(feature = "adversarial-samples")]
    #[test]
    fn test_samples_by_category() {
        let dataset = oxideshield_benchmark_dataset();
        let pi_samples = dataset.samples_by_category("PromptInjection");
        assert!(pi_samples.len() > 0);
    }

    #[test]
    fn test_benchmark_dataset_new() {
        let dataset = BenchmarkDataset::new("Test", "1.0.0");
        assert!(dataset.is_empty());
        assert_eq!(dataset.len(), 0);
    }

    #[test]
    fn test_combined_empty_datasets() {
        let combined = combined_dataset(vec![]);
        assert!(combined.is_empty());
    }
}