norrrrrrr-lyn/wainjectbenchBenchmarking prompt injection detections for web agents.
WAInjectBench is a comprehensive benchmark for prompt injection detection in web agents.
It covers 6 types of attacks, across two modalities: text and image.
data/
text/
benign/ → 4 categories, stored as JSONL filesmalicious/ → 8 attack types, stored as JSONL filesimage/
benign/ → 2 categories, stored in subfoldersmalicious/ → 7 attack types, stored in subfoldersClone the repo and create the environment:
git clone https://github.com/Norrrrrrr-lyn/WAInjectBench.git
cd WAInjectBench
conda env create -f environment.yml
conda activate wainjectbench
WAInjectBench supports two evaluation pipelines: text-based detection and image-based detection.
python main_text.py \
--data_dir [path to text dataset] \
--detector [detector name] \
--result_dir [output path] \
--gpu [gpu id]
Available detectors: ["kad", "promptarmor", "embedding-t", "promptguard", "datasentinel", "ensemble"]
PromptArmor → requires OPENAI_API_KEY as environment variable.
DataSentinel →
git clone https://github.com/liu00222/Open-Prompt-Injection.git
Download the pretrained model into: WAInjectBench/Open-Prompt-Injection/DataSentinel_Models Set the directory and model path in detector_text/datasentinel.py.
python main_image.py \
--data_dir [path to image dataset] \
--detector [detector name] \
--result_dir [output path] \
--gpu [gpu id]
Available detectors: ["gpt-4o-prompt", "llava-1.5-7b-prompt", "jailguard", "embedding-i", "llava-1.5-7b-ft", "ensemble"]
GPT-4o-Prompt → requires OPENAI_API_KEY as environment variable.
JailGuard →
git clone https://github.com/shiningrain/JailGuard.git
Follow its README to configure MiniGPT4.
LLaVA-1.5-7B-FT → requires downloading our finetuned model and setting its path in detector_image/llava.py.
We also provide in-domain trained versions of the Embedding-T and Embedding-I models, available in model/embedding-t/in-domain and model/embedding-i/in-domain. To use them, follow the same evaluation procedure as in the main experiments, but update the model path in detector_text/embedding-t.py and detector_image/embedding-i.py.
We provide code for training embedding-based binary classifiers for both text and image.
Text embedding classifier
python train/embedding-t.py \
--input_dir [dir with training text jsonl files] \
--output_dir [model output path]
JSONL format:
{"text": "example", "label": 1} # 1 for malicious, 0 for benign
Image embedding classifier
python train/embedding-i.py \
--input_dir [dir with training image jsonl files] \
--output_dir [model output path]
JSONL format:
{"path": "path/to/image.png", "label": 1}
Finetuning LLaVA-1.5-7B
python train.py \
--train_jsonl train.jsonl \
--val_jsonl val.jsonl \
--use_lora \
--amp_dtype bf16 \
--device_mode single \
--gpu_id 0
The JSONL files should contain image paths and labels (1 = malicious, 0 = benign). Experiments in our paper use the default hyperparameters.