LLMCert-T: Certifying Agentic Tool-Selection Robustness

This is the official repository for the paper "Certifying Robustness of Agentic Tool-Selection Under Adversarial Distributions".

Authors: Jehyeok Yeon, Isha Chaudhary Gagandeep Singh

---

Overview

LLMCert-T is a framework for formally certifying the robustness of LLM agents that select tools from a pool in the presence of adversarial tool injections. LLMCert-T models tool selection as a two-stage pipeline: a Retriever surfaces a top-N slate of candidate tools and an Agent picks one tool to execute. We place this pipeline under adaptive, multi-round adversarial pressure and compute a high-confidence lower bound on the agent's robust accuracy using Clopper-Pearson intervals.

---

Getting Started

Installation

1. Install PyTorch for your CUDA version first.

2. Install the remaining dependencies:

   pip install -r requirements.txt

   All LLM inference runs locally via vLLM — no external API needed. Models are downloaded automatically on first use and cached in HF_HOME.

---

Configuration

Before running anything, open scripts/config.sh and review the path settings. Override any setting by editing config.sh or by exporting environment variables before running a script — environment variables always take precedence.

Key settings in config.sh:

Variable	Default	Description
LLMCert-T_OUTPUTS_DIR	<repo>/LLMCert-T_outputs	Where certification results (JSON) are saved
HF_HOME	/fast/jyeon/hf_cache	HuggingFace model cache
TORCH_HOME	/fast/jyeon/torch_cache	PyTorch model cache
LLMCert-T_CONTAINER_SIF	(unset)	Path to Apptainer .sif image (cluster runs only)

For gated HuggingFace models (e.g., Llama), also set HF_TOKEN in config.sh or export it before running.

---

Core Components

• Tool: Defines name, description, parameters (JSON schema), and a critical privilege_level (e.g., "basic", "admin").
• Retriever: Uses semantic search (default: all-MiniLM-L6-v2) to filter the tool pool down to a slate of size N.
• Agent: Prompts a local LLM to select a tool from the slate. It performs privilege checks and argument extraction. Returns None if it abstains.
• Adversary: An adaptive generator of malicious tools. It iteratively refines its attacks based on feedback from the agent's previous choices.
• ToolCertifier: Orchestrates the multi-round Monte Carlo simulations, aggregates outcomes, and computes the Clopper-Pearson lower bound for robust accuracy.
• Judge: An LLM-based oracle that evaluates if the agent's selection semantically satisfies the user's intent compared to a ground-truth reference tool.

Supported Attack Families

1. Adversarial Selection: Optimizes tool descriptions with persuasion cues (e.g., "Official", "v2") to trick the agent.
2. Top-N Saturation: Floods the retrieval slate with near-duplicates to displace the correct tool from the agent's view.
3. Privilege Escalation: Tricks the agent into selecting a tool that requires higher privileges than the user possesses.
4. Intent Shifting: Uses semantic nudging to divert the agent to a relevant-sounding but incorrect tool.
5. Abstention Trigger: Injects refusal-inducing patterns to paralyze the agent.

---

Usage — Local Runs

Use scripts/run_local.sh to run any stage directly. Run all commands from the LLMCert-T repository root.

1. (Optional) Pre-generate Augmented Queries

Pre-generate paraphrases and narrative contexts for each query (paper Section 4.1). This is run once and cached:

python scripts/generate_augmented_queries.py \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --input data/questions.json \
  --pool data/options.json \
  --output data/augmented_questions.json \
  --num_paraphrases 5 \
  --seed 42

2. Run Certification

# Using defaults (Llama-3.1-8B-Instruct, Adversarial Selection, 10 rounds):
python run.py

# Override models and attack type:
python run.py \
  --agent_model meta-llama/Llama-3.1-8B-Instruct \
  --adversary_model google/gemma-3-4b-it \
  --attack_type "Top-N Saturation" \
  --rounds 5 \
  --adv_budget 3 \
  --temperature 0.7 \
  --output_json results.json

Or use the local wrapper script:

bash scripts/run_local.sh --agent_model meta-llama/Llama-3.1-8B-Instruct

If data/augmented_questions.json exists, it is automatically loaded and applied. Pass --augmented_queries to specify a different path.

---

Usage — HTCondor Cluster

Prerequisites

1. Build or pull an Apptainer container image with the required Python dependencies. The .sub files assume the container is accessible from all execute nodes (e.g., on shared/Lustre storage).

   apptainer build LLMCert-T_container.sif LLMCert-T_container.def

2. Create the logs directory (HTCondor writes .err/.out/.log files here):

   mkdir -p logs

3. Set required environment variables:

   export LLMCert-T_CONTAINER_SIF=/path/to/LLMCert-T_container.sif

   Optionally override storage paths (all default to values in config.sh):

   export HF_HOME=/fast/storage/hf_cache
   export TORCH_HOME=/fast/storage/torch_cache

   These are inherited by submitted jobs via getenv = True.

Submitting Jobs

Always cd into the LLMCert-T root before submitting. The .sub files use LLMCert-T_ROOT to bind-mount the repo into the container, so the working directory must be the LLMCert-T root.

cd /path/to/LLMCert-T

# Step 1: Pre-generate augmented queries (run ONCE)
condor_submit_bid 200 scripts/augment_container.sub

# Step 2: Run certification (after step 1 completes)
condor_submit_bid 200 scripts/certify_container.sub

Edit the .sub files to change model names, attack type, budget, etc.

Monitoring

condor_q $USER
tail -f logs/job.*.out

The entire LLMCert-T directory is bind-mounted into the container at /code, so no file transfer is needed. Weights, outputs, and caches all resolve to paths inside the repo via scripts/config.sh unless overridden with environment variables.

---

CLI Arguments

Certification (run.py)

Argument	Default	Description
--agent_model	meta-llama/Llama-3.1-8B-Instruct	HF model for the defender agent
--adversary_model	same as agent	HF model for the adversary
--judge_model	same as agent	HF model for the judge
--attack_type	Adversarial Selection	Attack family
--rounds	10	Adversarial refinement rounds per trial
--adv_budget	5	Adversarial tools injected per round
--slate_size	10	Retriever top-N
--user_privilege	basic	User privilege level
--temperature	0.7	Sampling temperature for agent and adversary LLMs
--seed	None	Random seed for reproducibility
--num_trials	None	Limit number of trials (default: all queries)
--tool_pool_json	None	Custom tool pool JSON file
--eval_set_json	None	Custom evaluation set JSON file
--augmented_queries	None	Path to augmented queries (auto-detected if in data/)
--output_json	None	Path to save results JSON
--verbose	False	Enable debug logging

Query Augmentation (scripts/generate_augmented_queries.py)

Argument	Default	Description
--model	meta-llama/Llama-3.1-8B-Instruct	LLM for generating paraphrases and narratives
--input	data/questions.json	Input questions file (BFCL format)
--output	data/augmented_questions.json	Output augmented queries file
--pool	data/options.json	Tool pool (used to filter invalid queries)
--num_paraphrases	5	Number of paraphrases per query
--temperature	0.8	Sampling temperature for generation
--seed	248	Random seed

---

Data Format

LLMCert-T uses tool and query data from the Berkeley Function Calling Leaderboard (BFCL) in JSONL format:

• data/options.json — Tool pool. Each line is a JSON object with name, description, and parameters.
• data/questions.json — Evaluation queries. Each line is a BFCL-format entry with question (chat messages) and function (ground-truth tool spec).
• data/eval.json — Evaluation metadata.
• data/augmented_questions.json — Pre-generated paraphrases and narrative contexts (produced by generate_augmented_queries.py).

---

File Structure

LLMCert-T/
├── LLMCert-T.py                              # Core framework (Tool, Retriever, Agent, Adversary, ToolCertifier)
├── run.py                               # CLI entry point for certification
├── requirements.txt                     # Python dependencies
├── LLMCert-T_container.def                   # Apptainer container definition
├── .gitignore
├── scripts/
│   ├── config.sh                        # Path configuration (edit for your cluster)
│   ├── run_local.sh                     # Local wrapper
│   ├── run_certify.sh                   # Cluster wrapper for certification
│   ├── run_augment.sh                   # Cluster wrapper for query augmentation
│   ├── generate_augmented_queries.py    # Pre-generate paraphrases + narrative contexts
│   ├── certify_container.sub            # HTCondor submit — certification
│   └── augment_container.sub            # HTCondor submit — query augmentation
├── data/
│   ├── options.json                     # Tool pool (BFCL format)
│   ├── questions.json                   # Evaluation queries (BFCL format)
│   ├── eval.json                        # Evaluation metadata
│   └── augmented_questions.json         # Pre-generated augmented queries
├── logs/                                # HTCondor job logs
└── LLMCert-T_outputs/                        # Certification results (JSON)

---

Citation

If you use this code in your research, please cite:

@misc{yeon2025quantifyingdistributionalrobustnessagentic,
      title={Quantifying Distributional Robustness of Agentic Tool-Selection}, 
      author={Jehyeok Yeon and Isha Chaudhary and Gagandeep Singh},
      year={2025},
      eprint={2510.03992},
      archivePrefix={arXiv},
      primaryClass={cs.CR},
      url={https://arxiv.org/abs/2510.03992}, 
}