wang-yanting/agentwatcher
AgentWatcher는 LLM 에이전트에서 간접 프롬프트 인젝션(indirect prompt injection)에 대응하는 탐지 기반 방어 기법입니다. 먼저 신뢰할 수 없는 컨텍스트에 대해 인과적 컨텍스트 기여도 분석(causal context attribution) 을 실행하여 가장 영향력이 큰 컨텍스트를 찾은 다음, 모니터 LLM이 그 컨텍스트를 명시적이고 사용자 정의 가능한 규칙에 따라 분류합니다. 완전한 블랙박스 탐지기와 비교했을 때, 이 파이프라인은 해석이 더 쉽습니다: 기여도 분석은 모델이 어디에 주목했는지 보여주고, 모니터 LLM의 판단은 규칙에 근거한 추론에 기반합니다. 모니터 LLM 출력의 예시는 아래와 같습니다:
이 규칙 기반 탐지는 모니터 LLM의 추론 능력을 활용하여 유용성(utility)과 견고성(robustness) 사이에서 더 나은 균형을 달성할 수 있습니다. AgentWatcher는 AgentDyn과 같은 LLM 에이전트 벤치마크에서 최첨단(state-of-the-art) 성능을 달성합니다:
이 저장소는 논문의 주요 실험을 재현하기 위한 것입니다.
pip install -r requirements.txt
huggingface-cli login # if needed for gated models
학습된 모니터는 LoRA 어댑터(PEFT)입니다. 기본적으로 main.py는 SecureLLMSys의 Hub 체크포인트 SecureLLMSys/AgentWatcher-Qwen3-4B-Instruct-2507 을 자동으로 로드합니다(--monitor_llm으로 재정의 가능).
모든 실행 스크립트는 scripts/에 있습니다. AgentWatcher 저장소 루트에서 실행한다고 가정합니다(또는 먼저 cd로 이동). 먼저 export OPENAI_API_KEY=<YOUR_KEY>로 OpenAI 키를 설정하세요. 실행 전에 Python 실행 스크립트(all_datasets, all_defenses, gpus, name, 모델 등)를 스윕(sweep)에 맞게 편집하세요.
main.py)배치 작업(로컬 GPU 또는 Slurm):
python scripts/run_long_context.py
이 스크립트는 vLLM(--use_vllm)으로 main.py를 구동하고 logs/main_logs/<name>/... 아래에 로그를 기록합니다. 로컬 GPU가 없으면 scripts/main.sh를 통해 Slurm 작업을 제출합니다.
실행 스크립트 없이 단일 실행:
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
main.py의 선택적 기여도 분석 재정의: --w_s, --w_l, --w_r, --K, --attribution_model.
AgentDojo/AgentDyn을 실행하려면 먼저 PIArena(https://github.com/sleeepeer/PIArena)에 따라 환경을 설정하세요:
cd agents/agentdojo && pip install -e . && cd ../..
그런 다음 실행:
python scripts/run_agentdojo.py
이 스크립트는 main_agentdojo.py를 호출하고 logs/agentdojo_logs/... 아래에 로그를 기록합니다. 로컬 GPU가 없으면 Slurm용으로 scripts/main_agentdojo.sh를 사용합니다. 수동으로 단일 실행하려면 저장소 루트에서 python main_agentdojo.py --model ... --defense agentwatcher --monitor_llm ...을 호출하면 됩니다. 래퍼(wrapper)에서 요구하는 경우 PIARENA_DEFENSE / PIARENA_MONITOR_LLM을 설정하세요. agentwatcher 방어는 PIMonitorLLMDefenseAdapter 를 사용합니다.
./scripts/run_agentdyn.sh
선택 사항: ./scripts/run_agentdyn.sh --model gpt-4o-2024-08-06은 gpt-4o를 백본 LLM으로 사용하여 agentwatcher를 실행합니다. 이 스크립트는 agentdojo conda 환경을 활성화하고 agents/agentdyn/src에서 실행됩니다. 실행이 중단되지 않도록 tmux를 사용하는 것이 좋습니다.
python scripts/run_injecagent.py
이 스크립트는 main_injecagent.py를 실행하고 logs/main_logs/... 아래에 로그를 기록합니다. Slurm 진입점: scripts/main_injecagent.sh.
실행 스크립트 없이 직접 호출:
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
이 코드를 사용하신다면 AgentWatcher 논문(공개 시)을 인용해 주시기 바랍니다.