wang-yanting/agentwatcherAtribuição causal de contexto e monitor baseado em regras para defesa de LLM contra injeção indireta de prompt em agentes LLM, alcançando desempenho de última geração no AgentDyn e em outros benchmarks.
AgentWatcher é uma defesa baseada em detecção contra injeção indireta de prompts em agentes LLM. Primeiro, executa atribuição causal de contexto sobre contextos não confiáveis para encontrar os contextos mais influentes, depois aplica um LLM monitor que classifica esses contextos sob regras explícitas e personalizáveis. Comparado com detectores totalmente caixa-preta, este pipeline é mais fácil de interpretar: a atribuição mostra onde o modelo focou, e o julgamento do LLM monitor é baseado em raciocínio fundamentado em regras. Um exemplo da saída do LLM monitor é mostrado abaixo:
Esta detecção baseada em regras pode alavancar a capacidade de raciocínio do LLM monitor para alcançar um melhor equilíbrio entre utilidade e robustez. AgentWatcher alcança desempenho de ponta em benchmarks de agentes LLM como AgentDyn:
Este repositório destina-se a reproduzir os principais experimentos do artigo.
pip install -r requirements.txt
huggingface-cli login # if needed for gated models
O monitor treinado é um adaptador LoRA (PEFT). Por padrão, main.py carrega automaticamente o checkpoint do Hub SecureLLMSys/AgentWatcher-Qwen3-4B-Instruct-2507 no SecureLLMSys (substitua com --monitor_llm).
Todos os lançadores estão em scripts/. Eles assumem que você executa a partir da raiz do repositório AgentWatcher (ou execute cd para lá primeiro). Primeiro, defina a chave OpenAI com export OPENAI_API_KEY=<SUA_CHAVE>. Edite os lançadores Python (all_datasets, all_defenses, gpus, name, modelos, etc.) para corresponder à sua varredura antes de executar.
main.py)Trabalhos em lote (GPUs locais ou Slurm):
python scripts/run_long_context.py
Isso aciona main.py com vLLM (--use_vllm) e escreve logs em logs/main_logs/<nome>/.... Sem GPU local, submete Slurm via scripts/main.sh.
Execução única sem o lançador:
python main.py \
--dataset lcc_long \
--attack combined \
--defense agentwatcher \
--backend_llm Qwen/Qwen3-4B-Instruct-2507 \
--monitor_llm SecureLLMSys/AgentWatcher-Qwen3-4B-Instruct-2507 \
--use_vllm \
--name my_run
Substituições de atribuição opcionais em main.py: --w_s, --w_l, --w_r, --K, --attribution_model.
Para executar AgentDojo/AgentDyn, primeiro configure o ambiente seguindo PIArena (https://github.com/sleeepeer/PIArena):
cd agents/agentdojo && pip install -e . && cd ../..
Depois execute:
python scripts/run_agentdojo.py
Isso chama main_agentdojo.py e escreve logs em logs/agentdojo_logs/.... Sem GPU local, usa scripts/main_agentdojo.sh para Slurm. Para uma execução manual única, você pode chamar python main_agentdojo.py --model ... --defense agentwatcher --monitor_llm ... a partir da raiz do repositório; defina PIARENA_DEFENSE / PIARENA_MONITOR_LLM se seu wrapper esperar por eles. A defesa agentwatcher usa PIMonitorLLMDefenseAdapter.
./scripts/run_agentdyn.sh
Opcional: ./scripts/run_agentdyn.sh --model gpt-4o-2024-08-06 para executar agentwatcher com gpt-4o como LLM de backbone. O script ativa o ambiente conda agentdojo e executa a partir de agents/agentdyn/src. Recomendamos usar tmux para evitar que a execução seja interrompida.
python scripts/run_injecagent.py
Isso lança main_injecagent.py e registra logs em logs/main_logs/.... Ponto de entrada Slurm: scripts/main_injecagent.sh.
Invocação direta sem o lançador:
python main_injecagent.py \
--model meta-llama/Llama-3.1-8B-Instruct \
--defense agentwatcher \
--monitor_llm SecureLLMSys/AgentWatcher-Qwen3-4B-Instruct-2507 \
--name my_injecagent_run
Se você usar este código, por favor cite o artigo AgentWatcher (quando disponível).