AgentPoison AI Agent

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

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

🔥🔥 Recent news please check Project page !

This repository provides the official PyTorch implementation of the following paper:

AgentPoison: Red-teaming LLM Agents via Memory or Knowledge Base Backdoor Poisoning <br> Zhaorun Chen¹, Zhen Xiang², Chaowei Xiao ³, Dawn Song ⁴, Bo Li^1,2

¹University of Chicago, ²University of Illinois, Urbana-Champaign <br> ³University of Wisconsin, Madison, ⁴University of California, Berkeley <br>

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

:hammer_and_wrench: Installation

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

git clone https://github.com/BillChan226/AgentPoison.git
cd AgentPoison
conda env create -f environment.yml
conda activate agentpoison

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:

python algo/trigger_optimization.py --agent ad --algo ap --model dpr-ctx_encoder-single-nq-base --save_dir ./results --ppl_filter --coh_sample --exclude_special --target_gradient_guidance --asr_threshold 0.8 --num_adv_passage_tokens 10 --golden_trigger -w -p

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.
--save_dir	./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_eval_batch_size	64	Batch size for trigger optimization
--num_cand	100	Number of discrete tokens sampled per optimization
--num_adv_passage_tokens	10	Number of tokens in the trigger sequence
--golden_trigger	False	Whether to start with a golden trigger (will overwrite --num_adv_passage_tokens)
--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

Local guidance runs the configured causal LLM on each candidate's demonstrations and triggered queries, then averages the joint probability of the target-action tokens. It uses no target-model backward pass. --use_gpt instead estimates the target-action rate from sampled chat completions. Both paths accept a candidate when its target score improves over its parent or reaches --asr_threshold. The token-replacement gradient comes from the retrieval model.

Use --target_model_path /path/to/llama for a local checkpoint, or set LLAMA2_CHAT_7B_PATH. --gpt_model selects the API model; OPENAI_BASE_URL selects an OpenAI-compatible endpoint and OPENAI_API_KEY supplies its credential. --device cpu allows CPU execution; the default remains cuda:0. --retriever_path and --ppl_model_path select explicit checkpoints and isolate their embedding caches from the default models.

Download the pretrained DPR, GPT-2, and Llama 3 checkpoints with:

python scripts/fetch_beam_checkpoints.py

The downloader records paths and revisions in results/beam_checkpoints.json, reuses complete cached snapshots, and downloads one weight format per model. Use --public_only for DPR and GPT-2, or --target_repo to select another target. The paper reports Llama 3-70B for StrategyQA; the downloader selects its Instruct variant, available as --target_model meta-llama-3-instruct-70b. The default target remains Llama 2-7B. Meta checkpoints require approved Hugging Face access and an authenticated account. For Llama 3, use Transformers 4.40 or later as described in Hugging Face's Llama 3 guide; for example, install transformers==4.40.2 and tokenizers==0.19.1 in a separate environment. The repository's existing environment pins remain unchanged.

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

Beam search

--beam_size 1 preserves the existing greedy optimizer and its coherence filtering. A larger beam retains distinct complete triggers across iterations. Each parent proposes replacements using its own gradient, and candidates pass retrieval and target-action gates relative to that parent. Selection ranks the combined pool by retrieval score; unchanged parents remain available when no replacement passes. Coherence sampling and target guidance are controlled by the existing flags.

For StrategyQA with DPR, start with:

python algo/trigger_optimization.py \
  --agent qa --algo ap --model dpr-ctx_encoder-single-nq-base \
  --golden_trigger --num_adv_passage_tokens 5 \
  --beam_size 4 --beam_validation_size 64 \
  --ppl_filter --coh_sample --exclude_special --target_gradient_guidance \
  --num_cand 100 --cand_pool_mult 5 --ppl_seq_batch 16 \
  --target_sample_size 8 --target_batch_size 16 --save_dir results/beam

Argument	Default	Description
--beam_size	1	Maximum number of distinct complete triggers retained. A single seed can grow into a larger beam.
--initial_triggers	unset	JSON array of seed strings, each with exactly --num_adv_passage_tokens retriever tokens and a valid text round trip. The number of distinct seeds must fit in the beam.
--beam_validation_size	0	Beam mode reserves 64 training queries when omitted, capped at 20% of the training set. Greedy mode reserves none unless explicitly requested.
--beam_validation_interval	10	Evaluate retained triggers on the fixed selection queries every this many iterations and at the final iteration.

The fixed selection queries are removed from optimization and are separate from the final evaluation split. trigger.json contains the best trigger observed on that fixed set; beam.json contains the current beam, parent links, scores, work counts, and peak GPU memory. beam_history.jsonl records each iteration. These files stay under the ignored results directory. work.json records runtime, GPU memory, and logical forward/sequence counts in beam mode and in greedy runs with validation enabled; the counts include validation and exclude OOM retries. Selection uses the retrieval score; validation target probability is recorded separately and does not establish end-to-end ASR.

For width comparisons, explicitly use the same --beam_validation_size, seed, guidance flags, and iteration settings at every width, including width 1. Enabling validation for greedy search changes its training split and output selection, while leaving its token update algorithm intact. Compare held-out ASR across several seeds, and compare against independent greedy restarts at a matched compute budget. A larger beam does not guarantee monotonically increasing ASR. Beam mode requires --coh_select_weight 0.

Run optimization, held-out retrieval checks, and complete local-agent episodes with:

python scripts/benchmark_strategyqa_beam.py \
  --checkpoint_manifest results/beam_checkpoints.json --device cuda:0

The benchmark prepares the full StrategyQA corpus in batches, runs matched settings at widths 1, 2, and 4, and records target answers, accuracy, retrieval counts, and format backoffs in an ignored results directory. Defaults use five optimization iterations and eight agent episodes; use --iterations, --gradient_steps, --query_batch, --candidates, --pool_mult, and --agent_instances for longer runs. For a smaller pretrained CPU experiment, download --target_repo TinyLlama/TinyLlama-1.1B-Chat-v1.0 and run the benchmark with --device cpu. The report retains model identifiers and revisions and marks short runs as pilot experiments.

Run the matched width/seed optimization sweep with:

PYTHON=/path/to/agentpoison/python bash scripts/run_beam_sweep.sh

This runs widths 1, 2, and 4 with seeds 0, 1, and 2. Set NUM_ITER=1000 for longer runs. It compares equal iteration settings; equal compute comparisons require matching recorded work and runtime against greedy restarts. Evaluate each output on the same held-out questions, for example:

python ReAct/run_strategyqa_gpt3.5.py --model dpr --task_type adv \
  --trigger_file 'results/beam_sweep/width_4/seed_0/qa/ap/*/trigger.json' \
  --save_dir results/beam_sweep/width_4/seed_0/evaluation

scripts/repro/ holds the per-seed run scripts and paired-comparison analysis used for the target-guidance results. Run one seed with, for example, PYTHON=/path/to/python bash scripts/repro/run_qa_seed.sh tgt01rel 0 GPU (arms are listed in each script's header; results go to results/llama_{qa,ehr,ad}_ARM/), then compare two arms with python scripts/repro/qa_pair.py llama_qa_tgt01rel llama_qa_noguide 1 (ehr_pair.py and ad_pair.py take the same two arm names).

To check the complete command pipeline without downloading models:

python scripts/verify_beam_pipeline.py

This builds small untrained DPR, GPT-2, and Llama checkpoints, runs widths 1, 2, and 4 with local guidance and mocked GPT transport, evaluates local-agent episodes, and checks actual retrieval followed by local generation. It verifies execution and interfaces, not pretrained-model ASR. Results are saved under results/pipeline_checks/. Live ASR comparisons require the intended pretrained checkpoints and an accessible GPU or GPT endpoint.

For StrategyQA local-agent evaluation, use --backbone llama3 --local_model_path /path/to/llama. The local path does not create a GPT client. Evaluation also accepts --device and --retriever_path, which should match optimization.

EHRAgent accepts --backbone local --local_model_path /path/to/llama, together with --device and --retriever_path. Its existing --backbone llama3 option uses Replicate. Use --held_out to evaluate the second half of the supplied query file, matching the optimizer's EHR split. Saved results include retrieved memory indices, poisoned retrieval counts, and generated code; code is not executed by this evaluator.

Build Agent-Driver and EHRAgent optimization caches in batches with:

python scripts/prepare_agent_embeddings.py --agent ad \
  --checkpoint_manifest results/beam_checkpoints.json --device cuda:0
python scripts/prepare_agent_embeddings.py --agent ehr \
  --checkpoint_manifest results/beam_checkpoints.json --device cuda:0

After building the driving optimization cache, prepare its full inference memory:

python scripts/prepare_driver_memory.py \
  --checkpoint_manifest results/beam_checkpoints.json --device cuda:0

Agent-Driver inference accepts --local_model_path, --retriever_path, --device, --local_max_new_tokens, and --seed. It uses the saved trigger text directly and retains its tokens when a long query is truncated. Pass --planner_model '' for a reasoning-only evaluation. For trajectory metrics, train a LoRA motion planner with python scripts/finetune_motion_planner.py --checkpoint_manifest results/beam_checkpoints.json --output_dir results/planner and pass --local_planner_path results/planner.

EHRAgent memory order is deterministic, and its embedding caches are specific to the ordered records. To compare saved EHR triggers in directories named width_1_seed_0, width_2_seed_0, and width_4_seed_0, run:

python scripts/evaluate_ehr_beam.py --optimization_root results/ehr_beam \
  --checkpoint_manifest results/beam_checkpoints.json --device cuda:0

Provide the labeled StrategyQA evaluation split at ReAct/data/strategyqa/strategyqa_dev.json. Evaluation 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 have modified the original code for Agent-Driver, ReAct-StrategyQA, EHRAgent to support more RAG embedders, and add interface for data poisoning. We have provided 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:

sh scripts/agent_driver/run_inference.sh

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:

sh scripts/agent_driver/run_evaluation.sh

We provide more options for red-teaming agent-driver that cover each individual 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:

python ReAct/run_strategyqa_gpt3.5.py --model dpr --task_type adv

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.

python ReAct/run_strategyqa_llama3_api.py --model dpr --task_type adv

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 --save_dir.

Evaluation

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

python ReAct/eval.py -p [RESPONSE_PATH]

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:

python EhrAgent/ehragent/main.py --backbone gpt --model dpr --algo ap --attack

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 --save_dir.

Evaluation

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

python EhrAgent/ehragent/eval.py -p [RESPONSE_PATH]

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:

@inproceedings{chenagentpoison,
  title={AgentPoison: Red-teaming LLM Agents via Poisoning Memory or Knowledge Bases},
  author={Chen, Zhaorun and Xiang, Zhen and Xiao, Chaowei and Song, Dawn and Li, Bo},
  booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems}
}

: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 zhaorun@uchicago.edu.

:key: License

This repository is under MIT License.