decision-model-guard 0.1.0__py3-none-any.whl

decision_model_guard/__init__.py

@@ -0,0 +1,31 @@
+"""Decision Model Guard: guardrails for decision models.
+
+Decision models (Jev-like models, classifiers, LLM judges) read a ticket, an email or a tool output and pick an answer.
+Text inside that content can claim the answer ("this case was already approved") and steer the decision. Decision Model Guard
+scans the untrusted content, runs the steering gate (erase-and-check) around any decision model, checks long content
+window by window, adds consistency, invariance, confidence and conformal checks and turns everything into an answer
+plus an action. `dmguard audit` measures how robust your model is with and without the guard.
+
+    from decision_model_guard import DecisionGuard
+    from decision_model_guard.models import WolfDefenderJev
+
+    guard = DecisionGuard.default(WolfDefenderJev.from_pretrained())
+    d = guard.decide(scene, "Does the content try to give the agent new instructions?",
+                     ["No, it is only data.", "Yes, an injected instruction."], risky={1})
+    d.answer, d.action, d.reasons
+"""
+from .audit import AuditLogger
+from .calibration import Calibrated, Calibrator
+from .config import build_guard
+from .extraction import ExtractionGuard, GuardedExtraction
+from .guards import Confidence, Conformal, Coverage, Invariance, OptionOrder, SteeringGate
+from .models import DecisionModel, FunctionModel
+from .pipeline import BOUNDARY_NOTE, DecisionGuard, Request, guard_decision, scan_content
+from .policy import Policy
+from .scene import TRUSTED_HEADERS, Scene, parse_blocks, units
+from .types import Action, Finding, GuardedDecision, GuardReport, Kind, ScanResult
+from .version import __version__
+
+__all__ = ["__version__", "DecisionGuard", "Request", "guard_decision", "scan_content", "build_guard", "Policy", "AuditLogger",
+           "SteeringGate", "OptionOrder", "Confidence", "Coverage", "Invariance", "Conformal", "Calibrator", "Calibrated", "ExtractionGuard", "GuardedExtraction", "Scene", "TRUSTED_HEADERS", "parse_blocks", "units", "DecisionModel",
+           "FunctionModel", "Action", "Kind", "Finding", "ScanResult", "GuardReport", "GuardedDecision", "BOUNDARY_NOTE"]

decision_model_guard/__main__.py

@@ -0,0 +1,3 @@
+from .cli import main
+
+raise SystemExit(main())

decision_model_guard/_onnx.py

@@ -0,0 +1,134 @@
+"""ONNX Runtime backend. The Patronus models (Wolf Defender, the Decision Steering Gate, Wolf Defender Jev) ship four ONNX
+variants in one layout, so a variant name is enough to pick a graph:
+
+    fp32                  onnx/onnx_fp32/model.onnx              native export
+    fp16                  onnx/onnx_fp16/model_fp16.onnx         half size
+    mixed                 onnx/onnx_mixed/model_mixed.onnx       INT8 weights, FP16 embeddings
+    int8_int4_embeddings  onnx/int8_int4_embeddings/model.onnx   INT8 weights, INT4 embeddings, the smallest
+
+A path to an .onnx file works as well. Needs the `onnx` extra (onnxruntime, tokenizers, numpy, huggingface_hub), no torch.
+"""
+from __future__ import annotations
+
+import json
+import os
+import threading
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+from .util import require
+
+VARIANTS = {"fp32": "onnx/onnx_fp32/model.onnx", "fp16": "onnx/onnx_fp16/model_fp16.onnx",
+            "mixed": "onnx/onnx_mixed/model_mixed.onnx", "int8_int4_embeddings": "onnx/int8_int4_embeddings/model.onnx"}
+DEFAULT_VARIANT = "int8_int4_embeddings"
+
+_SESSIONS: dict[tuple, Any] = {}
+_LOCK = threading.Lock()
+
+
+def resolve(repo_or_path: str, variant: str, files: Sequence[str] = ("config.json", "tokenizer.json"),
+            revision: str | None = None) -> tuple[Path, Path]:
+    """(folder with config and tokenizer, graph path). From the hub only the listed files and the chosen graph are fetched."""
+    rel = VARIANTS.get(variant, variant)
+    if not rel.endswith(".onnx"):
+        raise ValueError(f"onnx must be one of {sorted(VARIANTS)} or a path to an .onnx file, got {variant!r}")
+    if os.path.isdir(repo_or_path):
+        folder = Path(repo_or_path)
+        graph = Path(rel) if os.path.isabs(rel) else folder / rel
+    else:
+        hub = require("huggingface_hub", "onnx")
+        paths = [Path(hub.hf_hub_download(repo_or_path, f, revision=revision)) for f in files]
+        folder = paths[0].parent
+        graph = Path(rel) if os.path.isabs(rel) else Path(hub.hf_hub_download(repo_or_path, rel, revision=revision))
+    if not graph.is_file():
+        raise FileNotFoundError(f"ONNX graph not found: {graph}")
+    return folder, graph
+
+
+def session(graph: Path, threads: int | None = None, providers: Sequence[str] | None = None):
+    """One InferenceSession per graph, thread count and providers, shared within the process."""
+    ort = require("onnxruntime", "onnx")
+    key = (str(graph), threads, tuple(providers or ()))
+    with _LOCK:
+        if key not in _SESSIONS:
+            opts = ort.SessionOptions()
+            if threads:
+                opts.intra_op_num_threads = threads
+            _SESSIONS[key] = ort.InferenceSession(str(graph), opts, providers=list(providers or ["CPUExecutionProvider"]))
+        return _SESSIONS[key]
+
+
+def clear_sessions() -> None:
+    with _LOCK:
+        _SESSIONS.clear()
+
+
+def tokenizer(folder: Path):
+    """tokenizers.Tokenizer without the padding and truncation settings a saved tokenizer.json may carry."""
+    tk = require("tokenizers", "onnx")
+    tok = tk.Tokenizer.from_file(str(folder / "tokenizer.json"))
+    tok.no_padding()
+    tok.no_truncation()
+    return tok
+
+
+def read_json(path: Path) -> dict:
+    return json.loads(path.read_text(encoding="utf-8")) if path.is_file() else {}
+
+
+def pad(seqs: Sequence[Sequence[int]], pad_id: int = 0):
+    np = require("numpy", "onnx")
+    L = max(1, max((len(s) for s in seqs), default=1))
+    ids = np.full((len(seqs), L), pad_id, dtype=np.int64)
+    mask = np.zeros((len(seqs), L), dtype=np.int64)
+    for i, s in enumerate(seqs):
+        ids[i, : len(s)] = s
+        mask[i, : len(s)] = 1
+    return ids, mask
+
+
+def softmax(logits):
+    np = require("numpy", "onnx")
+    x = np.asarray(logits, dtype=np.float64)
+    e = np.exp(x - x.max(-1, keepdims=True))
+    return e / e.sum(-1, keepdims=True)
+
+
+class OnnxClassifier:
+    """A sequence classifier exported to ONNX (inputs input_ids, attention_mask; output logits), with the tokenizer and
+    the label names of its repository. Same interface as the PyTorch backend of the content scanners."""
+
+    def __init__(self, repo_or_path: str, variant: str = DEFAULT_VARIANT, threads: int | None = None,
+                 providers: Sequence[str] | None = None, revision: str | None = None):
+        folder, graph = resolve(repo_or_path, variant, revision=revision)
+        self.session = session(graph, threads, providers)
+        self.inputs = {i.name for i in self.session.get_inputs()}
+        self.tok = tokenizer(folder)
+        cfg = read_json(folder / "config.json")
+        self.id2label = {int(k): str(v) for k, v in (cfg.get("id2label") or {}).items()}
+        pad_token = self.tok.token_to_id("<pad>")
+        self.pad_id = cfg.get("pad_token_id", pad_token if pad_token is not None else 0) or 0
+        self._max_length: int | None = None
+        self._plain = tokenizer(folder)                                       # untruncated, for head/tail and chunking
+
+    def ids(self, text: str) -> list[int]:
+        return self._plain.encode(text, add_special_tokens=False).ids
+
+    def decode(self, ids: Sequence[int]) -> str:
+        return self._plain.decode(list(ids))
+
+    def _encode(self, texts: Sequence[str], max_length: int) -> list[list[int]]:
+        """Token ids with the special tokens of the template, truncated like transformers' `truncation=True`."""
+        if self._max_length != max_length:
+            self.tok.enable_truncation(max_length)
+            self._max_length = max_length
+        return [e.ids for e in self.tok.encode_batch(list(texts))]
+
+    def probs(self, texts: Sequence[str], max_length: int, batch_size: int, pos: int) -> list[float]:
+        out: list[float] = []
+        for i in range(0, len(texts), batch_size):
+            ids, mask = pad(self._encode(texts[i: i + batch_size], max_length), self.pad_id)
+            feed = {k: v for k, v in (("input_ids", ids), ("attention_mask", mask)) if k in self.inputs}
+            out += softmax(self.session.run(None, feed)[0])[:, pos].tolist()
+        return out

decision_model_guard/audit.py

@@ -0,0 +1,44 @@
+"""Audit log: one JSON line per guarded decision."""
+from __future__ import annotations
+
+import json
+import os
+import threading
+import time
+from typing import Any
+
+from .types import GuardedDecision
+from .util import sha256
+
+
+class AuditLogger:
+    """Appends one JSON line per decision. By default only hashes of the scene, question and culprit are stored, so the
+    log holds no customer content; set `include_text=True` to keep the text for debugging. Rotates at `max_bytes`."""
+
+    def __init__(self, path: str, include_text: bool = False, max_bytes: int | None = 50_000_000):
+        self.path, self.include_text, self.max_bytes = path, include_text, max_bytes
+        self._lock = threading.Lock()
+        os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True)
+
+    def record(self, d: GuardedDecision, scene: str, question: str, meta: dict[str, Any] | None = None) -> dict[str, Any]:
+        rec: dict[str, Any] = {"ts": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "id": d.id, "action": d.action.value, "kind": d.kind.value,
+                               "answer": d.answer, "raw_answer": d.raw_answer, "changed": d.changed, "probs": [round(p, 6) for p in d.probs],
+                               "raw_probs": [round(p, 6) for p in d.raw_probs], "reasons": d.reasons, "calls": d.calls,
+                               "latency_ms": round(d.latency_ms, 2), "options": d.options,
+                               "scans": {s.scanner: {"valid": s.valid, "score": round(s.score, 6)} for s in d.scans},
+                               "guards": {g.guard: g.triggered for g in d.guards}, "meta": meta or {}}
+        if self.include_text:
+            rec.update(scene=scene, question=question, culprit=d.culprit)
+        else:
+            rec.update(scene_sha256=sha256(scene), question_sha256=sha256(question), culprit_sha256=sha256(d.culprit) if d.culprit else None)
+            if d.culprit:                                   # reasons quote the culprit; keep content out of the log
+                rec["reasons"] = [r.replace(repr(d.culprit), "<culprit>").replace(d.culprit, "<culprit>") for r in d.reasons]
+        return rec
+
+    def log(self, d: GuardedDecision, scene: str, question: str, meta: dict[str, Any] | None = None) -> None:
+        line = json.dumps(self.record(d, scene, question, meta), ensure_ascii=False)
+        with self._lock:
+            if self.max_bytes and os.path.exists(self.path) and os.path.getsize(self.path) > self.max_bytes:
+                os.replace(self.path, self.path + ".1")
+            with open(self.path, "a", encoding="utf-8") as f:
+                f.write(line + "\n")

decision_model_guard/calibration.py

@@ -0,0 +1,232 @@
+"""Calibration and conformal abstention from your own labelled decisions.
+
+A probability of 0.7 from a decision model rarely means 70 %: calibration drifts between domains, and a threshold of
+0.5 or one copied from a benchmark seldom fits your traffic. With a few hundred labelled cases:
+
+    cal = Calibrator.fit(model, cases, alpha=0.05)   # temperature scaling + split-conformal threshold
+    cal.save("calibration.json")
+    guard = DecisionGuard(Calibrated(model, cal), guards=[Conformal(cal)], ...)
+
+`Calibrated` rescales the model's probabilities (temperature scaling: one parameter, argmax unchanged). `Conformal`
+turns every decision into a prediction set that contains the correct answer with probability at least 1 - alpha
+(split conformal prediction; the guarantee holds when your traffic looks like the calibration cases). A set with
+exactly one option is a decision; a set with two or more options, or none, means the model cannot decide at that
+coverage, and the policy abstains or escalates. `fit` also reports risk thresholds for target false-positive rates.
+"""
+from __future__ import annotations
+
+import json
+import math
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+from .models.base import DecideBatch, DecisionModel, as_model
+from .types import GuardReport, Kind
+from .util import argmax, normalize
+
+EPS = 1e-12
+
+
+def _scale(p: Sequence[float], t: float) -> list[float]:
+    logits = [math.log(max(x, EPS)) / t for x in p]
+    m = max(logits)
+    e = [math.exp(v - m) for v in logits]
+    s = sum(e)
+    return [v / s for v in e]
+
+
+def _nll(P: Sequence[Sequence[float]], y: Sequence[int], t: float) -> float:
+    return -sum(math.log(max(_scale(p, t)[g], EPS)) for p, g in zip(P, y, strict=True)) / max(1, len(y))
+
+
+def _ece(P: Sequence[Sequence[float]], y: Sequence[int], bins: int = 10) -> float:
+    acc = [[0, 0.0, 0] for _ in range(bins)]                 # count, confidence sum, correct
+    for p, g in zip(P, y, strict=True):
+        c, k = max(p), argmax(list(p))
+        b = min(bins - 1, int(c * bins))
+        acc[b][0] += 1
+        acc[b][1] += c
+        acc[b][2] += int(k == g)
+    n = max(1, len(y))
+    return sum(abs(cs / cnt - cor / cnt) * cnt / n for cnt, cs, cor in acc if cnt)
+
+
+def fit_temperature(P: Sequence[Sequence[float]], y: Sequence[int]) -> float:
+    """Temperature minimising the negative log-likelihood (golden-section search over log T in [0.05, 20])."""
+    lo, hi = math.log(0.05), math.log(20.0)
+    g = (math.sqrt(5) - 1) / 2
+    a, b = hi - g * (hi - lo), lo + g * (hi - lo)
+    fa, fb = _nll(P, y, math.exp(a)), _nll(P, y, math.exp(b))
+    for _ in range(60):
+        if fa < fb:
+            hi, b, fb = b, a, fa
+            a = hi - g * (hi - lo)
+            fa = _nll(P, y, math.exp(a))
+        else:
+            lo, a, fa = a, b, fb
+            b = lo + g * (hi - lo)
+            fb = _nll(P, y, math.exp(b))
+    return math.exp((lo + hi) / 2)
+
+
+def conformal_quantile(scores: Sequence[float], alpha: float) -> float:
+    """Split-conformal threshold: the ceil((n + 1)(1 - alpha))-th smallest nonconformity score."""
+    n = len(scores)
+    k = math.ceil((n + 1) * (1 - alpha))
+    return 1.0 if k > n else sorted(scores)[k - 1]
+
+
+def collect(model: DecisionModel | DecideBatch, cases: Sequence[dict[str, Any]]) -> tuple[list[list[float]], list[int]]:
+    """Model probabilities and gold indices for labelled cases ({"scene", "question", "options", "gold"}), batched per task."""
+    m = as_model(model)
+    groups: dict[tuple, list[int]] = {}
+    for i, c in enumerate(cases):
+        groups.setdefault((c["question"], tuple(c["options"])), []).append(i)
+    P: list[list[float] | None] = [None] * len(cases)
+    for (q, opts), idx in groups.items():
+        for i, p in zip(idx, m.decide_batch([cases[i]["scene"] for i in idx], q, list(opts)), strict=True):
+            P[i] = normalize(p)
+    return [p for p in P if p is not None], [int(c["gold"]) for c in cases]
+
+
+class Calibrator:
+    """Temperature scaling plus a split-conformal threshold, fitted on labelled decisions."""
+
+    def __init__(self, temperature: float = 1.0, qhat: float | None = None, alpha: float | None = None, n: int = 0,
+                 metrics: dict[str, Any] | None = None, operating_points: list[dict[str, float]] | None = None,
+                 risky: Sequence[int] | None = None, tasks: Sequence[str] | None = None):
+        self.temperature, self.qhat, self.alpha, self.n = temperature, qhat, alpha, n
+        self.metrics, self.operating_points = metrics or {}, operating_points or []
+        self.risky, self.tasks = list(risky) if risky else None, list(tasks or [])
+
+    # ------------------------------------------------------------------ fitting
+    @classmethod
+    def fit(cls, model: DecisionModel | DecideBatch, cases: Sequence[dict[str, Any]], alpha: float = 0.1,
+            target_fprs: Sequence[float] = (0.001, 0.01, 0.05, 0.1)) -> Calibrator:
+        """Fit on labelled cases. Use cases you did not tune anything else on; 200 or more give stable thresholds."""
+        if not 0 < alpha < 1:
+            raise ValueError("alpha must be between 0 and 1")
+        P, y = collect(model, cases)
+        return cls.fit_probs(P, y, alpha, risky=_common_risky(cases), target_fprs=target_fprs,
+                             tasks=sorted({c["question"] for c in cases}))
+
+    @classmethod
+    def fit_probs(cls, P: Sequence[Sequence[float]], y: Sequence[int], alpha: float = 0.1, risky: Sequence[int] | None = None,
+                  target_fprs: Sequence[float] = (0.001, 0.01, 0.05, 0.1), tasks: Sequence[str] | None = None) -> Calibrator:
+        if len(P) != len(y) or not P:
+            raise ValueError("need one probability list per gold label")
+        t = fit_temperature(P, y)
+        Pc = [_scale(p, t) for p in P]
+        qhat = conformal_quantile([1 - p[g] for p, g in zip(Pc, y, strict=True)], alpha)
+        sets = [[i for i, v in enumerate(p) if 1 - v <= qhat] for p in Pc]
+        metrics = {"cases": len(y), "accuracy": sum(argmax(list(p)) == g for p, g in zip(P, y, strict=True)) / len(y),
+                   "nll_before": _nll(P, y, 1.0), "nll_after": _nll(P, y, t), "ece_before": _ece(P, y), "ece_after": _ece(Pc, y),
+                   "set_coverage": sum(g in s for s, g in zip(sets, y, strict=True)) / len(y),
+                   "singleton_rate": sum(len(s) == 1 for s in sets) / len(y),
+                   "mean_set_size": sum(len(s) for s in sets) / len(y)}
+        ops = _operating_points(Pc, y, risky, target_fprs) if risky else []
+        return cls(t, qhat, alpha, len(y), metrics, ops, risky, tasks)
+
+    # ------------------------------------------------------------------ use
+    def calibrate(self, probs: Sequence[float]) -> list[float]:
+        return _scale(normalize(probs), self.temperature)
+
+    def prediction_set(self, probs: Sequence[float], calibrated: bool = False) -> list[int]:
+        """Options whose calibrated probability is at least 1 - qhat."""
+        if self.qhat is None:
+            raise ValueError("this calibrator has no conformal threshold")
+        p = list(probs) if calibrated else self.calibrate(probs)
+        return [i for i, v in enumerate(p) if 1 - v <= self.qhat]
+
+    def evaluate(self, model: DecisionModel | DecideBatch, cases: Sequence[dict[str, Any]]) -> dict[str, float]:
+        """Coverage and set sizes on held-out cases (check that the guarantee holds on your data)."""
+        P, y = collect(model, cases)
+        sets = [self.prediction_set(p) for p in P]
+        single = [(s, g) for s, g in zip(sets, y, strict=True) if len(s) == 1]
+        return {"cases": len(y), "coverage": sum(g in s for s, g in zip(sets, y, strict=True)) / len(y),
+                "singleton_rate": len(single) / len(y), "singleton_accuracy": sum(s[0] == g for s, g in single) / max(1, len(single)),
+                "mean_set_size": sum(len(s) for s in sets) / len(y), "ece": _ece([self.calibrate(p) for p in P], y)}
+
+    def summary(self) -> str:
+        m = self.metrics
+        lines = [f"Calibration on {self.n} cases: temperature {self.temperature:.3f}, accuracy {100 * m.get('accuracy', 0):.1f} %",
+                 f"ECE {100 * m.get('ece_before', 0):.1f} -> {100 * m.get('ece_after', 0):.1f}, NLL {m.get('nll_before', 0):.3f} -> {m.get('nll_after', 0):.3f}"]
+        if self.qhat is not None:
+            lines.append(f"Conformal sets at {100 * (1 - self.alpha):.0f} % coverage: {100 * m.get('singleton_rate', 0):.1f} % single answers, "
+                         f"mean set size {m.get('mean_set_size', 0):.2f}")
+        for op in self.operating_points:
+            lines.append(f"Risk threshold {op['threshold']:.4f}: FPR {100 * op['fpr']:.2f} % (target {100 * op['target_fpr']:.1f} %), "
+                         f"TPR {100 * op['tpr']:.1f} %")
+        return "\n".join(lines)
+
+    # ------------------------------------------------------------------ persistence
+    def to_dict(self) -> dict[str, Any]:
+        return {"temperature": self.temperature, "qhat": self.qhat, "alpha": self.alpha, "n": self.n, "metrics": self.metrics,
+                "operating_points": self.operating_points, "risky": self.risky, "tasks": self.tasks}
+
+    def save(self, path: str | Path) -> None:
+        Path(path).write_text(json.dumps(self.to_dict(), indent=2) + "\n", encoding="utf-8")
+
+    @classmethod
+    def load(cls, path: str | Path) -> Calibrator:
+        return cls(**json.loads(Path(path).read_text(encoding="utf-8")))
+
+
+def _common_risky(cases: Sequence[dict[str, Any]]) -> list[int] | None:
+    rs = {tuple(c["risky"]) for c in cases if c.get("risky")}
+    return list(rs.pop()) if len(rs) == 1 and all(c.get("risky") for c in cases) else None
+
+
+def _operating_points(P, y, risky, targets) -> list[dict[str, float]]:
+    rs = set(risky)
+    risk = [sum(p[i] for i in rs) for p in P]
+    pos = [r for r, g in zip(risk, y, strict=True) if g in rs]
+    neg = sorted((r for r, g in zip(risk, y, strict=True) if g not in rs), reverse=True)
+    if not pos or not neg:
+        return []
+    out = []
+    for t in targets:
+        k = int(math.floor(t * len(neg)))                   # negatives allowed above the threshold
+        thr = min(1.0, neg[k] + 1e-9) if k < len(neg) else 0.0
+        out.append({"target_fpr": t, "threshold": thr, "fpr": sum(r >= thr for r in neg) / len(neg),
+                    "tpr": sum(r >= thr for r in pos) / len(pos)})
+    return out
+
+
+class Calibrated:
+    """A decision model with calibrated probabilities (temperature scaling; the answer itself never changes)."""
+
+    def __init__(self, model: DecisionModel | DecideBatch, calibrator: Calibrator):
+        self.model, self.calibrator = as_model(model), calibrator
+        self.name = f"calibrated:{self.model.name}"
+        for attr in ("context_tokens", "count_tokens"):            # keep what Coverage needs
+            if hasattr(self.model, attr):
+                setattr(self, attr, getattr(self.model, attr))
+
+    def decide_batch(self, scenes: Sequence[str], question: str, options: Sequence[str]) -> list[list[float]]:
+        return [self.calibrator.calibrate(p) for p in self.model.decide_batch(scenes, question, options)]
+
+
+class Conformal:
+    """Abstain unless the conformal prediction set holds exactly one option (coverage 1 - alpha of the calibrator).
+
+    The guarantee holds only for the decisions the calibrator was fitted on: questions it has not seen are skipped and
+    reported as such (fit one calibrator per decision task, or fit on a mix of all tasks you guard)."""
+
+    name = "Conformal"
+
+    def __init__(self, calibrator: Calibrator | str | Path):
+        self.calibrator = calibrator if isinstance(calibrator, Calibrator) else Calibrator.load(calibrator)
+        if self.calibrator.qhat is None:
+            raise ValueError("the calibrator has no conformal threshold; fit it with alpha")
+
+    def check(self, model: DecisionModel, scene: str, question: str, options: Sequence[str], probs: Sequence[float], kind: Kind,
+              **_ctx) -> GuardReport:
+        tasks = self.calibrator.tasks
+        if tasks and not any(question.startswith(t) for t in tasks):   # startswith: a boundary note may follow the question
+            return GuardReport(self.name, False, 0.0, 0, {"skipped": "not calibrated for this question"})
+        already = isinstance(model, Calibrated) and model.calibrator is self.calibrator
+        s = self.calibrator.prediction_set(probs, calibrated=already)
+        return GuardReport(self.name, len(s) != 1, float(len(s) != 1), 0,
+                           {"set": [options[i] for i in s], "coverage": round(1 - (self.calibrator.alpha or 0), 4)})

decision_model_guard/cli.py

@@ -0,0 +1,186 @@
+"""Command line: dmguard {version, init-config, scan, decide, eval, audit, calibrate, serve}."""
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+from .version import __version__
+
+
+def _guard(args):
+    from .config import build_guard
+    return build_guard(args.config)
+
+
+def _read(args) -> str:
+    if getattr(args, "text", None):
+        return args.text
+    if getattr(args, "file", None):
+        return Path(args.file).read_text(encoding="utf-8")
+    if not sys.stdin.isatty():
+        return sys.stdin.read()
+    raise SystemExit("give --text, --file or pipe text on stdin")
+
+
+def cmd_init(args) -> int:
+    from .config import EXAMPLE_CONFIG
+    path = Path(args.path)
+    if path.exists() and not args.force:
+        raise SystemExit(f"{path} exists (use --force to overwrite)")
+    path.write_text(EXAMPLE_CONFIG, encoding="utf-8")
+    print(f"wrote {path}")
+    return 0
+
+
+def cmd_scan(args) -> int:
+    g = _guard(args)
+    text = _read(args)
+    if not text.lstrip().startswith("["):
+        text = f"[USER PROMPT]\n{text}"
+    sanitized, results = g.scan(text)
+    out = {"valid": all(r.valid for r in results), "results": [r.to_dict() for r in results]}
+    if args.sanitized:
+        out["sanitized"] = sanitized
+    print(json.dumps(out, indent=2, ensure_ascii=False))
+    return 0 if out["valid"] else 2
+
+
+def cmd_decide(args) -> int:
+    g = _guard(args)
+    scene = _read(args)
+    ints = lambda s: [int(x) for x in s.split(",")] if s else None  # noqa: E731
+    d = g.decide(scene, args.question, args.option, risky=ints(args.risky), order=ints(args.order), fallback=args.fallback)
+    if args.json:
+        print(json.dumps(d.to_dict(), indent=2, ensure_ascii=False))
+    else:
+        print(f"answer:   {d.answer}")
+        print(f"action:   {d.action.value}")
+        print(f"model:    {d.raw_answer}")
+        if d.reasons:
+            print("reasons:  " + "\n          ".join(d.reasons))
+        print(f"calls:    {d.calls}   latency: {d.latency_ms:.0f} ms")
+    return 0
+
+
+def cmd_eval(args) -> int:
+    from .evaluate import evaluate, load_cases
+    g = _guard(args)
+    rep = evaluate(g.model, load_cases(args.cases), guard=g, attacks_per_case=args.attacks, controls_per_case=args.controls,
+                   langs=args.lang or ["en"], seed=args.seed)
+    md = rep.to_markdown()
+    print(md)
+    if args.out:
+        Path(args.out).write_text(json.dumps(rep.to_dict(), indent=2), encoding="utf-8")
+    return 0
+
+
+def cmd_audit(args) -> int:
+    from .evaluate import FAMILIES, audit, load_cases
+    g = _guard(args)
+    rep = audit(g.model, load_cases(args.cases), guard=None if args.no_guard else g, families=args.check or list(FAMILIES),
+                per_case=args.per_case, padding_tokens=args.padding_tokens, seed=args.seed)
+    md = rep.to_markdown()
+    print(md)
+    if args.out:
+        Path(args.out).write_text(json.dumps(rep.to_dict(), indent=2), encoding="utf-8")
+    if args.markdown:
+        Path(args.markdown).write_text(md, encoding="utf-8")
+    if args.max_attack_success is not None:                       # for CI: fail when an attack gets through too often
+        key = "raw" if args.no_guard else "guarded"
+        worst = {k: v.rates()[key] for k, v in rep.families.items() if v.kind == "attack" and v.n}
+        over = {k: r for k, r in worst.items() if r is not None and r > args.max_attack_success}
+        if over:
+            print("attack success above the limit: " + ", ".join(f"{k} {100 * r:.1f} %" for k, r in over.items()), file=sys.stderr)
+            return 2
+    return 0
+
+
+def cmd_calibrate(args) -> int:
+    import yaml
+
+    from .calibration import Calibrator
+    from .config import build_model
+    from .evaluate import load_cases
+    cfg = yaml.safe_load(Path(args.config).read_text())
+    model = build_model(cfg["model"])                             # the bare model: calibration is fitted on raw probabilities
+    cal = Calibrator.fit(model, load_cases(args.cases), alpha=args.alpha)
+    print(cal.summary())
+    if args.test:
+        ev = cal.evaluate(model, load_cases(args.test))
+        print(f"Held-out: coverage {100 * ev['coverage']:.1f} % (target {100 * (1 - args.alpha):.0f} %), "
+              f"{100 * ev['singleton_rate']:.1f} % single answers, accuracy of single answers {100 * ev['singleton_accuracy']:.1f} %")
+    cal.save(args.out)
+    print(f"wrote {args.out}")
+    return 0
+
+
+def cmd_serve(args) -> int:
+    from .server import create_app
+    from .util import require
+    uvicorn = require("uvicorn", "server")
+    uvicorn.run(create_app(config=args.config), host=args.host, port=args.port, log_level=args.log_level)
+    return 0
+
+
+def main(argv: list[str] | None = None) -> int:
+    ap = argparse.ArgumentParser(prog="dmguard", description="Guardrails for decision models.")
+    sub = ap.add_subparsers(dest="cmd", required=True)
+    sub.add_parser("version")
+    p = sub.add_parser("init-config", help="write an example config")
+    p.add_argument("path", nargs="?", default="dmguard.yml")
+    p.add_argument("--force", action="store_true")
+    for name, helptext in (("scan", "scan untrusted text"), ("decide", "one guarded decision"), ("eval", "steering evaluation"),
+                           ("audit", "robustness audit: attacks and stress checks, with and without the guard"),
+                           ("calibrate", "fit temperature and conformal threshold on labelled cases"), ("serve", "HTTP API")):
+        q = sub.add_parser(name, help=helptext)
+        q.add_argument("--config", required=True)
+        if name in ("scan", "decide"):
+            q.add_argument("--text")
+            q.add_argument("--file")
+        if name == "scan":
+            q.add_argument("--sanitized", action="store_true", help="also print the sanitized text")
+        if name == "decide":
+            q.add_argument("--question", required=True)
+            q.add_argument("--option", action="append", required=True)
+            q.add_argument("--risky", help="comma-separated risky option indices")
+            q.add_argument("--order", help="option indices, most cautious first")
+            q.add_argument("--fallback", type=int)
+            q.add_argument("--json", action="store_true")
+        if name == "eval":
+            q.add_argument("--cases", help="JSONL cases (default: bundled samples)")
+            q.add_argument("--attacks", type=int, default=4)
+            q.add_argument("--controls", type=int, default=2)
+            q.add_argument("--lang", action="append")
+            q.add_argument("--seed", type=int, default=0)
+            q.add_argument("--out", help="write the report as JSON")
+        if name == "audit":
+            q.add_argument("--cases", help="JSONL cases (default: bundled samples)")
+            q.add_argument("--check", action="append", help="run only these checks (repeatable)")
+            q.add_argument("--per-case", type=int, default=2, help="variants per check and case")
+            q.add_argument("--padding-tokens", type=int, help="filler per side for the padding check (default: 1.2 x model context)")
+            q.add_argument("--no-guard", action="store_true", help="measure the model alone")
+            q.add_argument("--seed", type=int, default=0)
+            q.add_argument("--out", help="write the report as JSON")
+            q.add_argument("--markdown", help="write the scorecard as Markdown")
+            q.add_argument("--max-attack-success", type=float, help="exit 2 when an attack check exceeds this rate (CI)")
+        if name == "calibrate":
+            q.add_argument("--cases", required=True, help="labelled JSONL cases for fitting")
+            q.add_argument("--test", help="labelled JSONL cases to check coverage on")
+            q.add_argument("--alpha", type=float, default=0.1, help="allowed miss rate of the conformal sets")
+            q.add_argument("--out", default="calibration.json")
+        if name == "serve":
+            q.add_argument("--host", default="127.0.0.1")
+            q.add_argument("--port", type=int, default=8000)
+            q.add_argument("--log-level", default="info")
+    args = ap.parse_args(argv)
+    if args.cmd == "version":
+        print(__version__)
+        return 0
+    return {"init-config": cmd_init, "scan": cmd_scan, "decide": cmd_decide, "eval": cmd_eval, "audit": cmd_audit,
+            "calibrate": cmd_calibrate, "serve": cmd_serve}[args.cmd](args)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())