PEFT
Safetensors
English
decision-model
system-1
calibration
tool-use
agents
lora

agent-decision-9b

A fast, calibrated decision model for AI agents. Given the tools, the user's request and the agent's proposed tool calls, it answers typed questions such as are these calls correct?, what is wrong with them? and should the agent call a tool, ask the user, or decline?, returning a probability for every allowed answer in one forward pass, with no text generation.

Use it anywhere an agent's step needs a fast, calibrated check: as a hook that vets tool calls before they run (like the grounding check in grounded_ai's Claude Code plugin), as a guardrail that stops an agent from acting on a confidently wrong call, for labelling or filtering agent traces, or as the cheap first stage of an evaluation cascade where only low-confidence decisions go to an LLM judge.

This repository holds a LoRA adapter and a fine-tuned joint head for Cloudflare/clef-flash (9.4B, Apache-2.0), with the calibration temperature already folded into the head.

Results

Test set: 1,139 rows from APIs and items never seen in training, sampled to a fixed number of rows per kind of example so rare kinds are measured as precisely as common ones. "Base" is the released Clef-flash on the same rows, zero-shot. Accuracy / expected calibration error, x100.

Question Base agent-decision-9b
Are the tool calls correct and complete? (yes/no) 68.0 / 18.5 98.1 / 1.0
What is wrong with the calls? (correct / wrong tool / wrong arguments / incomplete) 56.2 / 8.0 96.4 / 1.9
Should the assistant call a tool, ask, or decline? 60.6 / 22.4 88.2 / 3.8
All 63.3 / 16.6 95.3 / 1.4

As a cascade's first stage: 91% of decisions come out at 90% confidence or more, and 97.7% of those are right (base: 37%, of which 75% are right).

Are the calls correct? by kind of call (accuracy x100):

Kind Base agent-decision-9b
Correct calls 97.3 97.3
A parameter the tool does not define 22.0 100
A value the user gave is missing 45.0 97.0
A value differs from what the user asked 72.2 96.5
The wrong tool 85.3 100
One of several requested calls is missing 87.8 96.3

Call, ask or decline, by correct answer: call 97.6 → 95.2, decline 58.7 → 87.3, ask 20.7 → 81.1.

Unseen presentation (the same test rows rendered as markdown, a layout never used in training): 93.5% overall (base 75.2%). The weak spot is "ask the user": 44.1% in this layout (base 15.3%).

Usage

Needs a GPU with bf16 and about 24 GB (L4 or larger). Qwen3.5's linear attention produces NaNs in float16, so a T4 will not work.

import sys
from huggingface_hub import snapshot_download
from peft import PeftModel
from safetensors.torch import load_file

base = snapshot_download("Cloudflare/clef-flash", revision="8b2e5fd17c09fd49bc2880b805d5323ec7ff4ff3")
sys.path.insert(0, base)
from joint_schema_model import load_release_model, systemone

model, processor = load_release_model(base, device="cuda")
ours = snapshot_download("grounded-ai/agent-decision-9b")
model.language_model = PeftModel.from_pretrained(model.language_model, ours).merge_and_unload()
model.head.load_state_dict(load_file(f"{ours}/joint_head.safetensors"))  # calibrated head

tools = [{"name": "loginuser", "description": "Log in a user.",
          "parameters": {"username": {"type": "str"}, "password": {"type": "str"}}}]
response = systemone(model, processor, {
    "model": "agent-decision-9b",
    "state": {"tools": tools, "user": "Log in a user named 'guest' with the password 'guestpass'.",
              "assistant_tool_calls": [{"name": "loginuser", "arguments": {"username": "guest789", "password": "guestpass"}}]},
    "questions": {
        "call_ok": {"type": "noul", "instructions": "Are the assistant's tool calls correct and complete for the user's request?"},
        "call_kind": {"type": "choice", "instructions": "What, if anything, is wrong with the proposed tool calls?",
                      "criteria": {"correct": "the calls are right and do everything the user asked",
                                   "wrong_tool": "a call uses a tool that cannot do what was asked",
                                   "wrong_arguments": "the right tool is called with arguments that do not match the request or the tool",
                                   "incomplete": "the calls are right but part of the request is not done"}},
    },
})
print(response["answers"])

Questions are free-form: any noul (yes/no), choice or score question can be asked about any state, as with Clef. The fine-tune targets the questions above; other questions behave like the base model.

Training data

All labels are exact, and nothing is truncated:

  • Salesforce/xlam-function-calling-60k (CC-BY-4.0). Half the rows are the execution-verified calls. The rest carry exactly one edit that makes them wrong by construction: a tool swapped for one whose schema rejects the arguments, a parameter the tool does not define, a value the user typed removed or changed to one they did not type, or one of several distinct calls dropped. Splits are by API name, so test APIs are never seen in training.
  • nvidia/When2Call (CC-BY-4.0), multiple-choice split: call a tool, ask for missing information, or decline.

52,552 training rows, rendered as JSON or as a transcript; markdown is held out for testing.

Training

LoRA rank 16, alpha 32, no dropout, on every attention, linear-attention and MLP projection of the text decoder except the two small linear-attention gates; the joint head trained fully in fp32. AdamW, learning rate 1e-4 (head 1e-5), constant schedule, effective batch 32 rows (4 length-grouped micro-batches of 8), gradient clipping 1.0, bf16 autocast, gradient checkpointing. The weights with the best validation log loss (step 125 of 150) were kept. One temperature (0.975) fitted on a held-out calibration split and folded into the head. One A100 40 GB, about 25 minutes of training. run_manifest.json holds the exact configuration.

Limitations

  • Wrong calls in training and test are exact edits of correct calls. They are certain to be wrong, but may be easier to spot than the mistakes real agents make; treat the results above as an upper bound until measured on real agent errors.
  • A call with a parameter the tool does not define counts as wrong. Some real APIs would ignore it; for a guardrail we chose strict.
  • "Ask the user" is the weakest decision, especially in unfamiliar layouts.
  • Confidence for a choice is its top probability, not normalized for the number of options.

License

Apache-2.0, like the base model. The training data is CC-BY-4.0: Salesforce xLAM function-calling dataset and NVIDIA When2Call.

Downloads last month
22
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for grounded-ai/agent-decision-9b

Finetuned
Qwen/Qwen3.5-9B
Adapter
(4)
this model

Datasets used to train grounded-ai/agent-decision-9b

Collection including grounded-ai/agent-decision-9b