AgentPoison - [NeurIPS 2024] Official implementation for "AgentP

项目说明

<di="center">
<img src="assets/agentpoison_logo.jpg" width="32%">
</div>

AgentPoison: Red-teaming LLM Agents or Knowledge Base Backdoor Poisoning

🔥🔥 Recent news please check Project page !

![Project Page]()
![Arxiv]()
![License: MIT]()
![GitHub Stars]()

This repository provides the official PyTorch implementation of the following paper:
> [AgentPoison: Red-teaming LLM Agents or Knowledge Base Backdoor Poisoning]() <br>
> Zhaorun Chen¹,
> Zhen Xiang²,
> Chaowei Xiao ³,
> Dawn Song ⁴,
> Bo Li^1,2
>
> ¹Uni of Chicago, ²Uni of Illinois, Urbana-Champaign <br>
³Uni of Wisconsin, Madison, ⁴Uni of California, Berkeley <br>

<di="center">
<img src="assets/method.png" width="95%">
</div>

:hammer_and_wrench: Installation

To install, run the following commands to install the required packages:

RAG Embedder Checkpoints

You can download the embedder checkpoints from the links below then specify the path to the embedder checkpoints in the algo/config.yaml file.

| Embedder | HF Checkpoints |
| -------------------- | ------------------- |
| BERT | google-bert/bert-base-uncased |
| DPR | facebook/dpr-question_encoder-single-nq-base |
| ANCE | castorini/ance-dpr-question-multi |
| BGE | BAAI/bge-large-en |
| REALM | google/realm-cc-news-pretrained-embedder |
| ORQA | google/realm-orqa-nq-openqa |

You can also use custmor embedders (e.g. fine-tuned yourself) as long as you specify their identifier and model path in the config.

:smiling_imp: Trigger Optimization

After setting up the configuration for the embedders, you can run trigger optimization for all three agents using the following command:

Specifically, the descriptions of arguments are listed below:

| Argument | Example | Description |
| -------------------- | ------------------- | ------------- |
| --agent | ad | Specify the type of agent to red-team, [ad, qa, ehr]. |
| --algo | ap | Trigger optimization algorithm to use, [ap, cpa]. |
| --model | dpr-ctx_encoder-single-nq-base | Target RAG embedder to optimize, see a complete list above. |
| --sa | ./result | Path to save the optimized trigger and procedural plots |
| --num_iter | 1000 | Number of iterations to run each gradient optimization |
| --num_grad_iter | 30 | Number of gradient accumulation steps |
| --per_gpu_e | 64 | Batch size for trigger optimization |
| --num_cand | 100 | Number of discrete tokens sampled per optimization |
| --num_ad | 10 | Number of tokens in the trigger sequence |
| --golden_trigger | False | Whether to start with a golden trigger (will o --num_ad) |
| --target_gradient_guidance | False | Whether to guide the token update with target model loss |
| --use_gpt | False | Whether to approximate target model loss via MC sampling |
| --asr_threshold | 0.8 | Target-action probability threshold |
| --ppl_filter | False | Whether to enable coherence loss filter for token sampling |
| --plot | False | Whether to plot the procedural optimization of the embeddings |
| --report_to_wandb | False | Whether to report the results to wandb |

Target-generation scoring (Eq. 9 / Eq. 11)

--target_gradient_guidance scores the target-action probability and keeps candidates that improve it or meet --asr_threshold. Use --use_gpt to estimate this probability from sampled completions.

| Argument | Example | Description |
| -------------------- | ------------------- | ------------- |
| --target_model | meta-llama-2-chat-7b | Target LLM from algo/config.py; provide a Hugging Face token for gated models. |
| --target_sample_size | 8 | Queries used to score each candidate's target action |
| --target_batch_size | 16 | Batch size for target-action scoring |
| --target_max_length | 1024 | Maximum target-model prompt length |

Per-agent target actions and prompt formats are defined in AGENT_TARGET_SPECS in algo/config.py.

Provide the labeled StrategyQA e split at ReAct/data/strategyqa/strategyqa_dev.json. E questions must be separate from the queries used for trigger optimization and poisoning.
For EHRAgent, pass held-out questions with --data_path; eicu_ac.json is also used for trigger optimization.

:robot: Agent Experiment

We ha the original code for Agent-Driver, ReAct-StrategyQA, EHRAgent to support more RAG embedders, and add interface for data poisoning. We ha unified dataset access for all three agents at here. Specifically, we list the inference command for all three agents.

:car: Agent-Driver

First download the corresponding dataset from here or the original dataset host. Put the corresponding dataset in agentdriver/data.
Then put the optimized trigger tokens in here and you can also determine more attack parameters in here. Specifically, set attack_or_not to False to get the benign utility under attack.

Then run the following script for inference:

The motion planning result regarding ASR-r, ASR-a, and ACC will be printed directly at the end of the program. The planned trajectory will be saved to ./result. Run the following command to get ASR-t:

We provide more options for red-teaming agent-driver that cover each indi components of an autonomous agent, including perception APIs, memory module, ego-states, mission goal.

You need to follow the instruction here and fine-tune a motion planner based on GPT-3.5 using OpenAI's API first. As an alternative, we fine-tune a motion planner based on LLaMA-3 in here, such that the agent inference can be completely offline. Set use_local_planner in here to True to enable this.

:memo: ReAct-StrategyQA

First download the corresponding dataset from here or the StrategyQA dataset. Put the corresponding dataset in ReAct/database.
Then put the optimized trigger tokens in here. Run the following command to infer with GPT backbone:

and similarly to infer with LLaMA-3-70b backbone (you need to first obtain an API key in Replicate to access LLaMA-3) and put it here.

Specifically, set --task_type to adv to inject querries with trigger and benign to get the benign utility under attack. You can also run corresponding commands through scripts/react_strategyqa. The results will be saved to a path indicated by --sa.

#### E

To e the red-teaming performance for StrategyQA, simply run the following command:

where RESPONSE_PATH is the path to the response json file.

:man_health_worker: EHRAgent

First download the corresponding dataset from here and put it under EhrAgent/database.
Then put the optimized trigger tokens in here. Run the following command to infer with GPT/LLaMA3:

You can specify --backbone to llama3 to infer with LLaMA3, and set --attack to False to get the benign utility under attack. You can also run corresponding commands through scripts/ehragent. The results will be saved to a path indicated by --sa.

#### E

To e the red-teaming performance for EHRAgent, simply run the following command:

where RESPONSE_PATH is the path to the response json file.

Note that for each of the agent, you need to run the experiments twice, once with the trigger to get the ASR-r, ASR-a, and ASR-t, and another time without the trigger to get ACC (benign utility).

:book: Acknowledgement

Please cite the paper as follows if you use the data or code from AgentPoison:

:book: Contact

Please reach out to us if you have any suggestions or need any help in reproducing the results. You can submit an issue or pull request, or send an email to .

:key: License

This repository is under MIT License.

项目信息

开发语言：Python

Star 数：244

技术要点

该项目属于红蓝攻防方向的开源资源，代码结构清晰，适合安全研究人员学习参考。项目围绕[NeurIPS 2024] Official implementation for "AgentPoison: Red-teaming LLM Agents via Memory or Knowledge Base Backdoor Poisoning"展开，提供了可运行的实现，对理解相关安全机制、工具原理与攻防技术具有较高的研究价值。

从技术实现角度看，此类项目的核心在于对底层原理的深入理解与工程化落地。开发者在研究过程中可以结合官方文档与社区讨论，梳理项目的初始化流程、核心模块职责与数据流转，逐步掌握其中的技术细节与设计思路。

使用与学习建议

1. 下载源码压缩包，解压到本地环境。

2. 根据项目类型准备对应的运行环境与依赖。

3. 按照 README 中的步骤完成配置与部署。

4. 在隔离的测试环境中运行验证，避免影响生产系统。

免责声明

本资源来源于开源社区，仅供安全研究、技术学习与合法授权范围内的攻防演练使用，请勿用于任何非法或侵害他人权益的用途。使用本项目产生的一切后果由使用者自行承担。