SafeKeep

Code for Tool Specifications Matter: Uncovering and Mitigating Safety Risks in AI Agents.

We identify schema-formatted tool specifications as a primary source of agent safety degradation and propose SafeKeep to restore native refusal behavior by assessing requests with flattened textual tool specifications before execution.

This repository contains the analysis that isolates the effect (Finding/) and the method itself (Method/).

Layout

Dataset/
  build_paired_direction_data.py   build the paired harmful/harmless dataset
  check_paired_data.py             sanity checks over the released data
  paired_direction_data/
    paired_data.jsonl              400 prebuilt pairs (included)

Finding/
  schema_to_prose.py               schema neutralization (JSON -> flattened text)
  semantic_to_gibberish.py         semantic randomization (meaning removed, structure kept)
  extract_refusal_direction.py     per-layer refusal directions, chatbot vs agent
  extract_schema_direction.py      schema direction R_c (Eq. 2)
  causal_intervention.py           activation shifts at the peak refusal layer (Eq. 3, 4)
  ablation/
    input_ablation.py              chatbot vs agent, representation level
    behavior_ablation.py           chatbot vs agent, behaviour level
    component_ablation.py          which part of the agent prompt causes the drop
    tool_specification_ablation.py 2x2 over semantic content x schema elements

Method/
  safekeep_demo.py                 end-to-end demo of the two SafeKeep stages

Setup

conda create -n safekeep python=3.11
conda activate safekeep
pip install -r requirements.txt

requirements.txt pins the versions used for the reported results. transformers in particular is worth pinning: tool specifications are rendered by each model's own chat template via apply_chat_template(tools=...), and that rendering is what the analysis measures, so a different version changes the inputs themselves.

The last two entries (datasets, openai) are only needed to rebuild the dataset and can be skipped otherwise.

Scripts load models from a local directory (local_files_only=True), so download the weights first and pass the path with --model or set MODEL_PATH. Experiments in the paper use Llama-3.1-8B-Instruct and Qwen2.5-7B-Instruct.

Two scripts call an external API and read credentials from the environment:

export LLM_JUDGE_API_KEY=...    # behavior_ablation.py (refusal judge)
export OPENAI_API_KEY=...       # build_paired_direction_data.py (dataset construction only)

Both accept an OPENAI_API_BASE / LLM_JUDGE_API_BASE override for OpenAI-compatible endpoints.

Data

Dataset/paired_direction_data/paired_data.jsonl ships with the repository: 400 pairs of harmful and harmless requests, each with the original system prompt containing the tool schema.

{"harmful": "...", "harmless": "...", "system": "...you can invoke:\n[{...}]"}

Each harmless request is a minimal rewrite of its harmful counterpart — on average 7.2 of 21.5 words change — so the two differ in intent while the surrounding agent context stays fixed. Verify with:

python Dataset/check_paired_data.py

Rebuilding the data is only necessary if you want to regenerate or extend it, and additionally requires the ToolSafety dataset (Xie et al., 2025):

export TOOLSAFETY_PATH=/path/to/ToolSafety
python Dataset/build_paired_direction_data.py

Reproducing the analysis

Each script is standalone and fits its own directions on the train split, so they can be run in any order. Results are written under a directory named after the script, then the model.

Does the agent format weaken refusal? Representation level, then behaviour level:

python Finding/ablation/input_ablation.py --model $MODEL_PATH
python Finding/ablation/behavior_ablation.py --model $MODEL_PATH

Which part of the agent prompt is responsible? Components are added incrementally on top of the chatbot baseline; the 2x2 then separates semantic content from schema elements within the tool specification:

python Finding/ablation/component_ablation.py --model $MODEL_PATH
python Finding/ablation/tool_specification_ablation.py --model $MODEL_PATH
python Finding/ablation/tool_specification_ablation.py --make-table   # after several models

Is the effect causal? Extract the directions, then intervene on activations:

python Finding/extract_refusal_direction.py --format chatbot --model $MODEL_PATH
python Finding/extract_refusal_direction.py --format agent   --model $MODEL_PATH
python Finding/extract_schema_direction.py --model $MODEL_PATH

python Finding/causal_intervention.py --model $MODEL_PATH --mode restore  --alphas 0,2,4,6,8
python Finding/causal_intervention.py --model $MODEL_PATH --mode subtract --alphas 0,4,8,12,16

Both transforms can be inspected on their own, without loading a model:

python Finding/schema_to_prose.py
python Finding/semantic_to_gibberish.py

Running SafeKeep

python Method/safekeep_demo.py --model $MODEL_PATH

Prints both stages for a harmful and a benign request that share the same tools and phrasing: the neutralized tool view given to the judge, the verdict, and the resulting response. --case harmful or --case benign runs just one.

Notes

• All experiments use a fixed 70/30 split with seed 42, shared across scripts, so directions fitted in one script are evaluated on inputs unseen in the same split.
• Hidden states are always taken at the final prompt token.
• behavior_ablation.py caches per-format results in results_<format>.json and reuses them on rerun; delete that file to force regeneration.