Vela Agent Guard

A guard for AI agents that touch wallets and API keys, built on Vela 2.0.

Community project. Not affiliated with or endorsed by KR Labs or the vLLM Semantic Router team, who build Vela 2.0.

Three checks, each a trained Vela head combined with a deterministic layer:

Check Model part Deterministic part
scan_secrets(text) span head for private_key, seed_phrase, api_key, wallet_address regex for well-known key formats; ambiguous shapes (a 64-hex string can be a tx hash) need model agreement
check_injection(request, tool_output) yes/no: does the tool output try to instruct the agent? –
check_tool_call(request, call) allow / confirm / block from the user's intent exact checks on recipients, amounts, tokens, unlimited approvals and address poisoning; the stricter verdict wins

The Vela backbone is frozen and untouched. This repo ships only the small fine-tuned heads (21 MB for 4B, 8 MB for 0.8B), swapped in for guard questions only. model.safetensors at the repo root is the 4B heads (span + candidate) merged into one file with span. / cand. prefixes, byte-identical to heads/4B/.

Results

Synthetic eval set, 300 rows per task, built from templates that never appear in training. "Hard" rows are the ones where the shape of a value is not enough: a tx hash vs. a private key, a numbered seed phrase, a transfer the user only asked about.

System Secrets char-F1 (hard) Injection acc / AUROC Tool call acc (hard) Median latency, A40
Regex (secrets) / rules (tool calls) only 0.830 (0.476) / 0.735 – 0.903 (0.805) ~0 ms
Vela 2.0 9B, release heads 0.614 (0.475) 0.883 / 0.984 0.883 (0.826) ~180 ms
Vela 2.0 4B, release heads + hybrid 0.859 (0.554) 0.793 / 0.921 0.923 (0.846) ~115 ms
Vela Agent Guard 0.8B 0.908 (0.811) 0.873 / 0.974 0.930 (0.886) ~80 ms
Vela Agent Guard 4B 0.946 (0.884) 0.970 / 0.996 1.000 (1.000) ~115 ms

In no guard configuration did a harmful tool call come back as allow.

Usage

pip install torch "transformers>=5.17" safetensors tokenizers numpy huggingface_hub mnemonic
pip install flash-linear-attention   # optional, much faster on GPU
import sys
from pathlib import Path
from huggingface_hub import snapshot_download

local = Path(snapshot_download("VelaOrg/vela-agent-guard"))
sys.path.insert(0, str(local))
from trainer.guard_model import GuardModel
from vela_guard.agent_guard import AgentGuard

guard = AgentGuard(GuardModel.load("vllm-sr/Vela-2.0-4B", local / "heads" / "4B", device="cuda"),
                   known_addresses={"0xYourRouterOrBridge"})   # contracts the agent may use unprompted

guard.scan_secrets(text)                     # {"found", "spans", "redacted"}
guard.check_injection(user_request, page)    # {"injection", "score"}
guard.check_tool_call(user_request, {"name": "transfer", "arguments": {"token": "USDC", "to": "0x...", "amount": "10"}})
# {"decision": "block", "rule_decision": "block", "model_decision": "confirm", "reasons": ["amount_mismatch"], ...}

See example.py for a runnable script. GuardModel.load pins the Vela base model to the commit the heads were trained on (base_revision in guard_config.json), so upstream code changes cannot silently alter the guard.

How it was trained

  • Data: synthetic, generated by vela_guard/gen_*.py (4,000 train rows per task, 400 dev). Fake keys, addresses and seed phrases are random but format-correct. Train and eval use disjoint templates; a test enforces this.
  • Method: the training code of Vela 2.0 is not public, so this repo follows the stage-2 recipe its model card describes. Backbone hidden states are computed once with the frozen model; copies of the release broad span head and candidate head are then fine-tuned (trainer/train_heads.py, AdamW, lr 3e-4, 8 epochs, best epoch by dev score). The span threshold is fitted on dev.
  • Compute: one A40 on RunPod, about 25 minutes for both sizes including feature extraction.

Limitations

  • All data is synthetic. The templates differ between train and eval, but they share the same value generators and a narrow domain. Real-world performance is not yet measured.
  • Dev saturates. Dev scores reach ~1.0 while eval scores are lower, so the heads partly memorise training templates and dev cannot rank epochs well.
  • 0.8B is weaker: it misses about 1 in 5 injection attacks in the eval set, over-blocks some harmless gas estimates, and can split a long secret into pieces with a word left unredacted (seen on a numbered seed phrase). Prefer 4B where leaks matter.
  • The rules are a policy, not a proof. They assume the user names every recipient, or that it is listed in known_addresses. Tools outside the built-in lists defer to the model.
  • English only. Synthetic prompt-injection payloads do not cover every attack style.

Use it as one layer of defence, not the only one.

Attribution

Built on Vela 2.0 by KR Labs and vLLM Semantic Router (Apache-2.0). Base weights are not redistributed; they load from their own repositories. See NOTICE.

surplus-launcher: 0x32f8544de85B4e9d874c7AEAda9493359D43502b

Downloads last month
-
Safetensors
Model size
5.43M params
Tensor type
F32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for VelaOrg/vela-agent-guard