interactionguard_defense

Multi-Layer Security for LLM Agents

This repository implements a research-oriented defense framework for LLM-based agents against:

• prompt injection
• tool injection
• RAG poisoning
• correlated multi-channel attacks

The system combines three layers:

• Layer 1: AutoEncoder-based anomaly screening over ModernBERT embeddings
• Layer 2: fine-tuned DeBERTa classifier for known malicious traces
• Layer 3: LLM-based semantic security auditor (single active implementation: O2)

It also includes:

• data split and preprocessing pipelines
• model training scripts (Layer 1 encoder + autoencoder, Layer 2 classifier, Layer 3 local Llama fine-tune)
• baseline defenses (RAGuard, Prompt Sentinel, Tool Result Parsing) for comparison
• ASR (Attack Success Rate) evaluation scripts, with and without defense
• figure generation for the paper/thesis
• a multi-seed "Std Testing" feature that retrains Layer 1/Layer 2 with different random seeds to report mean ± std instead of a single-run point estimate

Repository Layout

interactionguard_defense/
├── config/                       # YAML + Python configuration
│   ├── detection_config.yaml     # layer thresholds, model paths, layer3 variant selection
│   ├── test_modes.yaml           # the 9 supported detection modes (1,2,3,10-15)
│   └── std_testing/              # per-trial config snapshots generated by Std Testing
├── data/
│   ├── raw/                      # raw trace JSONL files (gitignored, must be supplied)
│   ├── splits/                   # train/val/test trace splits (generated)
│   ├── processed/                # prefix-expanded processed samples (generated)
│   ├── data_split.py             # raw -> splits
│   └── data_preprocess.py        # splits -> processed
├── eval/                         # ASR evaluation scripts (with/without defense, baseline)
├── experiments/                  # sweeps and comparison experiments
│   └── std_testing/              # multi-seed retraining + error-bar reporting
├── models/
│   ├── detectors/                # Layer 1 / 2 / 3 detector implementations + pipeline
│   ├── baseline/                 # RAGuard / Sentinel / Tool Result Parsing baselines
│   ├── layer1/, layer2/, layer3/ # trained checkpoints (gitignored, must be trained/supplied)
│   └── std_testing/              # multi-seed retrained checkpoints (gitignored)
├── results/                      # detection outputs, sweep outputs, figures (gitignored)
├── scripts/                      # menu-wired utility scripts (isolation sweeps, backups)
│   └── standalone/               # standalone research scripts not wired into main.py
├── training/                     # training scripts for Layer 1 / Layer 2 / Layer 3
├── utils/                        # data, metrics, plotting, menu helpers
└── main.py                       # main interactive entry point

Data Flow

The codebase follows this high-level workflow:

1. Raw traces are supplied in data/raw/*.jsonl (see the schema note below).
2. data/data_split.py creates train/val/test splits in data/splits/.
3. data/data_preprocess.py expands each trace into one or more prefix-level sequence samples in data/processed/.
4. Training scripts build Layer 1 and Layer 2 checkpoints under models/layer1/ and models/layer2/.
5. main.py runs detection, baseline tests, ASR analysis, and figure generation.

Raw trace schema (data/raw/*.jsonl)

Five files are required, with these exact names:

benign_traces.jsonl
prompt_injection_traces.jsonl
rag_poison_traces.jsonl
tool_injection_traces.jsonl
correlated_injection_traces.jsonl

Each line is one JSON trace. The required field is reasoning_steps: an ordered list of step dicts. Recognized node values are retrieve, agent, and tools — a tool-call step's tool name is carried in a separate "tool" field, not encoded into node itself:

{"node": "tools", "type": "tool", "content": "...", "tool": "read_docx"}

label (0/1) and attack_type are optional per-record — if omitted, data_split.py fills them in based on which file the record came from.

Generating these files: raw traces are produced by a separate data-generation testbed, not by this repository — interactionguard_testbed. It runs a real LangGraph tool-using agent (RAG + Gmail + file tools) under benign traffic and three attack families (RAG poisoning, tool-output injection, correlated multi-channel), and its data/traces/*.jsonl output maps directly, file-for-file, onto the five files listed above — copy them over unchanged (see that repo's "Output Contract for Downstream Consumers" section for the exact mapping and schema notes).

Processed samples use a sequence format such as:

[PROMPT]
...

[SEP]

[RETRIEVE]
...

[SEP]

[TOOL: get_information]
...

[SEP]

[AGENT]
...

[EOS]

Defense Modes

The detection pipeline supports 9 modes from config/test_modes.yaml:

Important:

• 10 and 15 are 2-layer pipelines; 11 and 12 are the full 3-layer pipelines.
• Figures that refer to the "3-layer" defense need results from mode 11 (or 12 for the alternate ordering).
• Only one Layer 3 implementation is actively supported: O2 (single-call channel auditor, models/detectors/layer3_detector_o2.py). A separate layer3orig implementation (models/detectors/layer3_detector.py) exists as a legacy fallback but is not part of any reported experiment.

Threshold Reference

This section states what each layer's decision threshold is for, so a reproduction attempt uses the right value. It intentionally does not include any measured results — see config/detection_config.yaml for the live values and the Std Testing feature (below) for how to characterize run-to-run variance instead of relying on a single number.

Notes:

• All three thresholds live in config/detection_config.yaml and can be edited directly; no code changes are needed to try a different value.
• layers.layer3_variant (near the top of config/detection_config.yaml) is an optional global override for any mode that uses the logical layer3 key. Leave it unset for the default (resolves to O2), or set it to orig to use the legacy implementation instead.
• Layer 2's fine-tuning step (Training -> Fine-tune Layer 2 model) always excludes attack_type == "correlated" from both train and validation, by design — this is not configurable from the menu.
• Because Layer 1 and Layer 2 are trained artifacts, re-training them (even with the same threshold and hyperparameters) produces a different model each time a different random seed is used, and therefore different precision/recall at the same threshold. This is expected, not a bug — see "Std Testing" below for the feature built specifically to quantify this.

Installation

Create a Python environment and install dependencies:

pip install -r requirements.txt

Recommended:

• Python 3.10+
• PyTorch (with CUDA if a GPU is available)
• Transformers
• scikit-learn
• tqdm
• OpenAI-compatible client dependencies (openai, python-dotenv)

Environment Variables

Create a .env file in the project root (not tracked by git). Variables actually read by this codebase (os.getenv(...)):

# Layer 3 (O2) and other OpenRouter-backed LLM calls
OPENROUTER_API_KEY=...
OPENROUTER_MODEL=openai/gpt-4.1-mini
OPENROUTER_BASE_URL=https://openrouter.ai/api/v1

# Used by some evaluation scripts (e.g. no-defense prompt-injection judging)
OPENAI_API_KEY=...

# GPU selection
CUDA_DEVICE_ML=0          # device index (or "cpu") for Layer 1 / Layer 2
CUDA_DEVICE_LLM=0         # device index for local Layer 3 LLM fine-tuning, if used

# Tool Result Parsing baseline behavior
TOOL_RESULT_EXTRACT_METHOD_LIST=ParseData,CheckTool
TOOL_RESULT_EXTRACT_PARSE_DATA_FULL_CONVERSATION=False
TOOL_RESULT_EXTRACT_CHECK_TOOL_FILTER=True

Notes:

• Layer 3 (O2) and any evaluation script that calls an LLM require a valid API key; without one, those code paths will fail at the API call, not silently produce wrong results.
• Local Ollama-backed Layer 3 variants (if you set layer3_variant: orig) read their base_url/model from config/detection_config.yaml instead of .env.

Run the Project

Launch the interactive menu:

python main.py

Main Menu Reference

main.py's top-level menu:

AGENT DEFENSE SYSTEM
1) Defense pipeline
2) Baselines
3) Evaluations
4) Figures
5) Training and threshold tools
6) Sweeps
7) Guided workflows
8) Documentation and about
9) Backup current results
10) Data preparation
0) Exit

1 Defense pipeline

• Interactive mode picker: prompts for a mode ID (see Defense Modes above) and, for modes that use the logical layer3 key, lets you pick a Layer 3 variant (default resolves to O2).
• Run mode 10/11/12/13/14/15: runs that specific mode directly on data/processed/test_processed.jsonl and saves the output to results/mode_<id>_results.json.

Use this when you want raw detection results for one pipeline mode, e.g. you need mode_11_results.json for full 3-layer ASR evaluation.

2 Baselines

• Run RAGuard / Run Tool Result Parsing / Run Prompt Sentinel: each runs its baseline over data/splits/test.jsonl and writes models/baseline/<name>_predictions.jsonl, <name>_metrics.json, <name>_thresholds.json.
• Build Combined baseline: checks whether the three prediction files exist (running any missing detector first), then merges them into models/baseline/combined_predictions.jsonl.

3 Evaluations

• Prompt attacks without defense (LLM judge): evaluates prompt-injection success with no defense in place, using an LLM judge; writes prompt failure cases used by the no-defense evaluator below.
• All attacks without defense: generates eval/eval_asr_without_defense.json (no-defense ASR for prompt, rag, tool, correlated). Depends on the prompt-attack judging step above having been run at least once.
• Defense and baseline ASR reports: evaluates the 3-layer defense and the baseline defenses, writing eval/eval_asr_within_defense.json and eval/eval_asr_within_baseline.json. The 3-layer part expects results/mode_11_results.json to already exist; the baseline part expects the baseline prediction files to already exist.
• Baseline-only ASR report: re-evaluates baseline prediction files only.

4 Figures

• Generate all figures: regenerates the saved figures in results/ from whatever evaluation/result files already exist. It does not re-run any detector itself.

5 Training and threshold tools

• Fine-tune Layer 1 encoder (ModernBERT)
• Train Layer 1 AutoEncoder (on top of the fine-tuned encoder)
• Fine-tune Layer 2 model (ProtectAI DeBERTa; automatically excludes correlated from train/val)
• Fine-tune Layer 3 local Llama 3.1 (Ollama-aligned) / deploy it to Ollama
• Find optimal Layer 1 threshold (searches for zero-false-negative threshold while minimizing false positives)
• Isolation-experiment tools (retrain + high-benign mode-11 sweep, extract FP samples, re-evaluate isolation models with the current Layer 3)

6 Sweeps

• Threshold sweep / Distribution sweep / High-benign distribution sweep (Issue 2)
• Layer 3 backend latency sweep, and its staged high-benign variant
• Layer 3 LLM backend comparison (swaps the LLM behind O2 across several Ollama/OpenRouter models, reusing a frozen Layer 1/2 routing cache)
• Layer 3 queueing decomposition sweep (fixed workload)
• Std testing (see next section)

6 -> 8 Std Testing (multi-seed L1/L2 retraining, error bars for Table 6.4)

This feature exists specifically because a single training run's point estimate is not, by itself, a reliable number to report (see "Threshold Reference" above). It retrains Layer 1 and Layer 2 with several random seeds (same data split, same hyperparameters) and evaluates all 9 modes per seed:

1. Phase 1: Train — trains Layer 1 (ModernBERT + AutoEncoder) and Layer 2 with N different seeds, saving each trial's checkpoints under models/std_testing/trial<i>_seed<seed>/.
2. Phase 2: Evaluate — runs all 9 modes per trial, saving results under results/std_testing/trial<i>_seed<seed>/.
3. Aggregate — recomputes results/std_testing/summary.json / summary.csv (mean ± std per metric per mode) from whatever trial results already exist, without re-running anything.
4. Generate figures — redraws the Issue 1 precision/recall/latency chart with error bars, sourced from summary.json.

Phase 1 and Phase 2 are resumable: each skips a trial/mode combination whose output already exists, unless that mode ID is listed under std_testing_invalidated_modes in config/detection_config.yaml (used to force re-evaluation after a config value such as a threshold has changed).

7 Guided workflows

Pre-composed sequences of the steps above for common tasks — see "Common Workflows" below for the exact menu paths.

8 Documentation and about

In-app reference text (introduction, workflow guide, author note, quick reference) — useful as a quick refresher without leaving the terminal.

9 Backup current results

Archives the current results/ directory into a timestamped zip under backups/ (via scripts/backup_current_results.sh).

10 Data preparation

• Regenerate splits from data/raw
• Regenerate processed files from data/splits
• Full regenerate: raw -> splits -> processed

Each option asks for a YES confirmation before overwriting, since it regenerates files other steps depend on.

Common Workflows

Menu paths below use the notation top-level -> submenu choice.

Build the full 3-layer result from scratch

1 -> 1 (Interactive mode picker) -> enter 11

Result: creates results/mode_11_results.json.

Refresh only baseline predictions

2 -> 1, 2 -> 2, 2 -> 3, then 2 -> 4 for the merged baseline.

Recompute baseline evaluation after changing a baseline

3 -> 4

Use this after raguard_predictions.jsonl, tool_result_parsing_predictions.jsonl, sentinel_predictions.jsonl, or combined_predictions.jsonl changed.

Generate the ASR line chart quickly from existing outputs

7 -> 4 (Guided workflows -> "Generate Issue 3 ASR line chart from ready outputs")

This is the shortest path only if results/mode_11_results.json and eval/eval_asr_without_defense.json already exist.

Generate the ASR line chart from a clean or uncertain state

7 -> 5 (Guided workflows -> "Generate Issue 3 ASR line chart from scratch")

This runs, in order: no-defense prompt judging -> no-defense ASR for all attacks -> mode 11 -> combined baseline -> defense/baseline ASR evaluation -> figure generation.

Update Issue 3 after rerunning only Tool Result Parsing

2 -> 2, 2 -> 4, 3 -> 4, 4 -> 1

Report mean ± std instead of a single run (recommended before citing a number)

6 -> 8 -> 1 (Std Testing Phase 1), then 6 -> 8 -> 2 (Phase 2), then 6 -> 8 -> 3 (Aggregate), then 6 -> 8 -> 4 (Generate figures).

Generate Issue 2 threshold/distribution figures

6 -> 1 or 6 -> 2 (choose mode 11 or 12 when prompted), then 4 -> 1.

Key Output Files

Detection pipeline outputs

• results/mode_<id>_results.json for each mode you have run (1, 2, 3, 10-15)

Baseline outputs

• models/baseline/raguard_predictions.jsonl
• models/baseline/tool_result_parsing_predictions.jsonl
• models/baseline/sentinel_predictions.jsonl
• models/baseline/combined_predictions.jsonl

Evaluation outputs

• eval/eval_asr_without_defense.json
• eval/eval_asr_within_defense.json
• eval/eval_asr_within_baseline.json

Std Testing outputs

• results/std_testing/summary.json, summary.csv
• results/std_testing/issue1_precision_recall_latency_chart_with_error_bars.png

Figure outputs

Examples:

• results/issue3_ASRline_chart.png
• threshold sweep, distribution sweep, and high-benign distribution sweep figures

Important Notes

• issue3_ASRline does not read models/baseline/tool_result_parsing_metrics.json directly — it reads eval/eval_asr_within_baseline.json. After changing a baseline prediction file, re-run baseline evaluation before regenerating figures.
• The full 3-layer ASR result depends on mode 11, not mode 10.
• Single-run result files (e.g. results/mode_2_results.json) reflect exactly one trained checkpoint. They are useful as a sanity check that the pipeline runs end-to-end, but should not be treated as the precise, reproducible number for any layer whose model is retrained — use Std Testing's mean ± std for that instead.
• data/raw/, data/splits/, data/processed/, models/layer1/, models/layer2/, models/layer3/, models/std_testing/, and results/ are all gitignored. Cloning this repository gives you the code only — raw traces, trained checkpoints, and prior results must be supplied or regenerated separately.

Entry Points

Primary entry:

python main.py

Data processing (equivalent to menu 10 -> 3):

python -m data.data_split
python -m data.data_preprocess

Repository History / Migration Note

2026-09-30: active development moved to a new, cleaned-up repository: https://github.com/MinhBang1000/interactionguard_defense. The previous repository (agent_defense, pushed to https://github.com/MinhBang1000/agent_eval.git) is frozen at its last commit and kept only as a historical archive of the full development process (all AE_000* branches, experimental Layer 3 variants, Dojo integration work, etc.). Its origin remote was intentionally removed (git remote remove origin) to prevent accidental pushes here.

To recover the link to the old GitHub remote if ever needed:

git remote add origin https://github.com/MinhBang1000/agent_eval.git

Project Status

This is a research codebase rather than a production package. Some scripts under scripts/standalone/ are experimental or exploratory and are not wired into main.py — run them directly with python -m if needed. The menu reference above is the safest way to operate the repository consistently.