Skip to main content

harn_guard/
catalog.rs

1//! The built-in catalog of downloadable injection-detection models.
2//!
3//! Catalog entries are *pointers* to already-hosted upstream model repositories
4//! (Hugging Face). Harn hosts nothing and bundles no weights — `harn guard
5//! install` fetches from these URLs on the user's machine, at the user's
6//! request. Integrity-critical weight files carry a pinned SHA-256 (the upstream
7//! Git-LFS object id); small config/tokenizer files are recorded
8//! trust-on-install.
9
10use serde::{Deserialize, Serialize};
11
12/// On-disk weight format a model ships in.
13#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
14#[serde(rename_all = "kebab-case")]
15pub enum ModelFormat {
16    /// ONNX graph, run by an ONNX runtime backend.
17    Onnx,
18    /// `safetensors` weights, run by a native-Rust backend (candle).
19    Safetensors,
20}
21
22impl ModelFormat {
23    pub fn as_str(self) -> &'static str {
24        match self {
25            Self::Onnx => "onnx",
26            Self::Safetensors => "safetensors",
27        }
28    }
29}
30
31/// One file to fetch for a model package.
32#[derive(Clone, Copy, Debug)]
33pub struct CatalogFile {
34    /// Path of the file within the upstream repository (the download source).
35    pub remote: &'static str,
36    /// Filename to store it under inside the model's install directory.
37    pub dest: &'static str,
38    /// Pinned SHA-256 (hex) of the file, or `None` to record-on-install (TOFU).
39    /// Pinned only where the upstream exposes a content digest (LFS objects).
40    pub sha256: Option<&'static str>,
41    /// Expected size in bytes, when known. Advisory (drives progress + a cheap
42    /// pre-download sanity check); integrity rests on `sha256`.
43    pub size: Option<u64>,
44}
45
46/// A catalog entry: an upstream, already-hosted model the user may install.
47#[derive(Clone, Copy, Debug)]
48pub struct CatalogModel {
49    /// Stable catalog id used on the CLI and in config (`guard_model`).
50    pub name: &'static str,
51    /// Human-facing name.
52    pub display_name: &'static str,
53    /// One-line description.
54    pub description: &'static str,
55    /// Upstream repository id (e.g. a Hugging Face `org/repo`).
56    pub repo: &'static str,
57    /// Base URL files are resolved against (`{base}/{remote}`).
58    pub base_url: &'static str,
59    /// SPDX-ish license identifier.
60    pub license_id: &'static str,
61    /// URL where the license / acceptable-use terms can be read.
62    pub license_url: &'static str,
63    /// `true` when the upstream requires authentication + license acceptance
64    /// (e.g. a gated Hugging Face repo). Install then requires `HF_TOKEN`.
65    pub gated: bool,
66    /// Weight format.
67    pub format: ModelFormat,
68    /// Files to fetch. The first whose `dest` ends in the format extension is the
69    /// weight file; the rest are tokenizer/config.
70    pub files: &'static [CatalogFile],
71}
72
73impl CatalogModel {
74    /// Resolve the absolute download URL for `file`.
75    pub fn url_for(&self, file: &CatalogFile) -> String {
76        format!("{}/{}", self.base_url.trim_end_matches('/'), file.remote)
77    }
78
79    /// Total advertised download size, when every file advertises one.
80    pub fn total_size(&self) -> Option<u64> {
81        self.files.iter().map(|f| f.size).sum()
82    }
83}
84
85/// The default model: ungated, permissively licensed, zero-friction.
86pub const DEFAULT_MODEL: &str = "deberta-v3-prompt-injection-v2";
87
88// ProtectAI DeBERTa-v3 prompt-injection classifier — Apache-2.0, ungated. The
89// ONNX weight + tokenizer + config are everything the inference backend needs.
90const DEBERTA_FILES: &[CatalogFile] = &[
91    CatalogFile {
92        remote: "onnx/model.onnx",
93        dest: "model.onnx",
94        sha256: Some("f0ea7f239f765aedbde7c9e163a7cb38a79c5b8853d3f76db5152172047b228c"),
95        size: Some(738_563_188),
96    },
97    CatalogFile {
98        remote: "onnx/tokenizer.json",
99        dest: "tokenizer.json",
100        sha256: None,
101        size: Some(8_648_886),
102    },
103    CatalogFile {
104        remote: "onnx/config.json",
105        dest: "config.json",
106        sha256: None,
107        size: Some(1_014),
108    },
109];
110
111// Meta Llama Prompt Guard 2 (community ONNX export) — gated (Llama Community
112// License + Hugging Face access request). Listed as an opt-in; install requires
113// the user's own HF_TOKEN and explicit license acceptance. We pin nothing we
114// cannot fetch anonymously, so the weight digests are recorded on install.
115const PROMPT_GUARD_2_86M_FILES: &[CatalogFile] = &[
116    CatalogFile {
117        remote: "model.onnx",
118        dest: "model.onnx",
119        sha256: None,
120        size: None,
121    },
122    CatalogFile {
123        remote: "tokenizer.json",
124        dest: "tokenizer.json",
125        sha256: None,
126        size: None,
127    },
128    CatalogFile {
129        remote: "config.json",
130        dest: "config.json",
131        sha256: None,
132        size: None,
133    },
134];
135
136const CATALOG: &[CatalogModel] = &[
137    CatalogModel {
138        name: DEFAULT_MODEL,
139        display_name: "ProtectAI DeBERTa-v3 prompt-injection v2",
140        description: "Apache-2.0, ungated. Recommended default — no account or license gate.",
141        repo: "protectai/deberta-v3-base-prompt-injection-v2",
142        base_url:
143            "https://huggingface.co/protectai/deberta-v3-base-prompt-injection-v2/resolve/main",
144        license_id: "Apache-2.0",
145        license_url: "https://huggingface.co/protectai/deberta-v3-base-prompt-injection-v2",
146        gated: false,
147        format: ModelFormat::Onnx,
148        files: DEBERTA_FILES,
149    },
150    CatalogModel {
151        name: "llama-prompt-guard-2-86m",
152        display_name: "Meta Llama Prompt Guard 2 (86M, ONNX)",
153        description: "Higher recall, but GATED — needs Hugging Face access + HF_TOKEN.",
154        repo: "gravitee-io/Llama-Prompt-Guard-2-86M-onnx",
155        base_url: "https://huggingface.co/gravitee-io/Llama-Prompt-Guard-2-86M-onnx/resolve/main",
156        license_id: "Llama-Community",
157        license_url: "https://huggingface.co/meta-llama/Llama-Prompt-Guard-2-86M",
158        gated: true,
159        format: ModelFormat::Onnx,
160        files: PROMPT_GUARD_2_86M_FILES,
161    },
162];
163
164/// All catalog models.
165pub fn all() -> &'static [CatalogModel] {
166    CATALOG
167}
168
169/// Find a catalog model by its stable `name`.
170pub fn find(name: &str) -> Option<&'static CatalogModel> {
171    CATALOG.iter().find(|m| m.name == name)
172}
173
174/// The recommended default model entry.
175pub fn default_model() -> &'static CatalogModel {
176    find(DEFAULT_MODEL).expect("default model is always present in the catalog")
177}
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182
183    #[test]
184    fn default_model_is_ungated_and_present() {
185        let model = default_model();
186        assert_eq!(model.name, DEFAULT_MODEL);
187        assert!(!model.gated, "the recommended default must be ungated");
188        assert_eq!(model.license_id, "Apache-2.0");
189    }
190
191    #[test]
192    fn default_model_matches_runtime_default_selector() {
193        // `harn-vm` carries the default guard-model selector independently (to
194        // stay free of a dependency on `harn-guard`); the two must not drift.
195        assert_eq!(DEFAULT_MODEL, harn_vm::config::DEFAULT_GUARD_MODEL);
196    }
197
198    #[test]
199    fn catalog_entries_are_well_formed() {
200        for model in all() {
201            assert!(!model.name.is_empty());
202            assert!(!model.files.is_empty(), "{} has no files", model.name);
203            assert!(
204                model.base_url.starts_with("https://"),
205                "{} base_url must be https",
206                model.name
207            );
208            // Any pinned digest must be a 64-char lowercase hex string.
209            for file in model.files {
210                if let Some(sha) = file.sha256 {
211                    assert_eq!(sha.len(), 64, "{}/{} sha len", model.name, file.dest);
212                    assert!(
213                        sha.bytes()
214                            .all(|b| b.is_ascii_hexdigit() && !b.is_ascii_uppercase()),
215                        "{}/{} sha must be lowercase hex",
216                        model.name,
217                        file.dest
218                    );
219                }
220            }
221        }
222    }
223
224    #[test]
225    fn url_for_joins_base_and_remote() {
226        let model = default_model();
227        let file = &model.files[0];
228        assert_eq!(
229            model.url_for(file),
230            "https://huggingface.co/protectai/deberta-v3-base-prompt-injection-v2/resolve/main/onnx/model.onnx"
231        );
232    }
233
234    #[test]
235    fn gated_model_is_opt_in_only() {
236        let gated = find("llama-prompt-guard-2-86m").expect("present");
237        assert!(gated.gated);
238        // We never pin digests we cannot fetch anonymously.
239        assert!(gated.files.iter().all(|f| f.sha256.is_none()));
240    }
241}