1use serde::{Deserialize, Serialize};
11
12#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
14#[serde(rename_all = "kebab-case")]
15pub enum ModelFormat {
16 Onnx,
18 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#[derive(Clone, Copy, Debug)]
33pub struct CatalogFile {
34 pub remote: &'static str,
36 pub dest: &'static str,
38 pub sha256: Option<&'static str>,
41 pub size: Option<u64>,
44}
45
46#[derive(Clone, Copy, Debug)]
48pub struct CatalogModel {
49 pub name: &'static str,
51 pub display_name: &'static str,
53 pub description: &'static str,
55 pub repo: &'static str,
57 pub base_url: &'static str,
59 pub license_id: &'static str,
61 pub license_url: &'static str,
63 pub gated: bool,
66 pub format: ModelFormat,
68 pub files: &'static [CatalogFile],
71}
72
73impl CatalogModel {
74 pub fn url_for(&self, file: &CatalogFile) -> String {
76 format!("{}/{}", self.base_url.trim_end_matches('/'), file.remote)
77 }
78
79 pub fn total_size(&self) -> Option<u64> {
81 self.files.iter().map(|f| f.size).sum()
82 }
83}
84
85pub const DEFAULT_MODEL: &str = "deberta-v3-prompt-injection-v2";
87
88const 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
111const 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
164pub fn all() -> &'static [CatalogModel] {
166 CATALOG
167}
168
169pub fn find(name: &str) -> Option<&'static CatalogModel> {
171 CATALOG.iter().find(|m| m.name == name)
172}
173
174pub 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 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 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 assert!(gated.files.iter().all(|f| f.sha256.is_none()));
240 }
241}