modelforge/
lib.rs

1//! Model provider and target parsing primitives for `thesa`.
2//!
3//! Made by Trevor Knott for Knott Dynamics.
4//!
5//! ModelForge normalizes CLI/user input for Hugging Face, Ollama, and CivitAI
6//! archive targets. It deliberately does not call provider APIs or download
7//! models; callers use it to turn user-facing strings into stable target enums.
8//! The crate keeps parsing deterministic and side-effect free.
9//!
10//! ```
11//! assert_eq!(
12//!     modelforge::parse_model_target_for_provider(
13//!         "https://huggingface.co/openai-community/gpt2",
14//!         modelforge::ModelProvider::Hf,
15//!     ),
16//!     Ok(modelforge::ModelTarget::Model("openai-community/gpt2".to_string()))
17//! );
18//! assert_eq!(
19//!     modelforge::parse_model_target_for_provider("newest", modelforge::ModelProvider::Civitai),
20//!     Ok(modelforge::ModelTarget::Latest)
21//! );
22//! ```
23
24use std::error::Error;
25use std::fmt;
26
27/// Integration guide for agents and other automated consumers.
28pub const AGENT_GUIDE: &str = include_str!("../AGENT_GUIDE.md");
29
30/// Result type used by ModelForge APIs.
31pub type Result<T> = std::result::Result<T, ModelForgeError>;
32
33/// Supported model/archive providers.
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub enum ModelProvider {
36    /// Hugging Face Hub.
37    Hf,
38    /// Local Ollama registry/runtime.
39    Ollama,
40    /// CivitAI model catalog.
41    Civitai,
42}
43
44impl ModelProvider {
45    /// Stable provider list in CLI display order.
46    pub const ALL: [Self; 3] = [Self::Hf, Self::Ollama, Self::Civitai];
47
48    /// Return all supported providers in stable CLI display order.
49    pub fn all() -> &'static [Self] {
50        &Self::ALL
51    }
52
53    /// Human-readable provider label.
54    pub fn label(self) -> &'static str {
55        match self {
56            Self::Hf => "Hugging Face",
57            Self::Ollama => "Ollama",
58            Self::Civitai => "CivitAI",
59        }
60    }
61
62    /// Short UI label used by compact archive displays.
63    pub fn short_label(self) -> &'static str {
64        match self {
65            Self::Hf => "HF",
66            Self::Ollama => "OLM",
67            Self::Civitai => "CIV",
68        }
69    }
70
71    /// Stable lowercase provider slug.
72    pub fn slug(self) -> &'static str {
73        match self {
74            Self::Hf => "hf",
75            Self::Ollama => "ollama",
76            Self::Civitai => "civitai",
77        }
78    }
79
80    /// Parse a provider slug accepted by thesa-style CLIs.
81    pub fn from_slug(value: &str) -> Option<Self> {
82        let normalized = normalized_provider_slug(value);
83        match normalized.as_str() {
84            "hf" | "huggingface" => Some(Self::Hf),
85            "ollama" | "olm" => Some(Self::Ollama),
86            "civitai" | "civit" | "civ" => Some(Self::Civitai),
87            _ => None,
88        }
89    }
90}
91
92fn normalized_provider_slug(value: &str) -> String {
93    value
94        .trim()
95        .chars()
96        .filter(|c| !matches!(c, '-' | '_' | ' ' | '\t' | '\n' | '\r'))
97        .flat_map(char::to_lowercase)
98        .collect()
99}
100
101impl fmt::Display for ModelProvider {
102    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
103        match self {
104            Self::Hf | Self::Ollama | Self::Civitai => write!(f, "{}", self.slug()),
105        }
106    }
107}
108
109/// Normalized archive target requested by a user or CLI.
110#[derive(Debug, Clone, PartialEq, Eq)]
111pub enum ModelTarget {
112    /// Provider namespace, organization, owner, or search keyword.
113    Org(String),
114    /// Concrete model id.
115    Model(String),
116    /// Provider-specific top/trending preset.
117    Top,
118    /// Provider-specific latest/newest preset.
119    Latest,
120}
121
122impl ModelTarget {
123    /// Return true for namespace, organization, owner, or search keyword targets.
124    pub fn is_org(&self) -> bool {
125        matches!(self, Self::Org(_))
126    }
127
128    /// Return true for concrete model id targets.
129    pub fn is_model(&self) -> bool {
130        matches!(self, Self::Model(_))
131    }
132
133    /// Return true for `Top` or `Latest` preset targets.
134    pub fn is_preset(&self) -> bool {
135        matches!(self, Self::Top | Self::Latest)
136    }
137
138    /// Borrow the namespace/search value for `Org` targets.
139    pub fn org(&self) -> Option<&str> {
140        match self {
141            Self::Org(value) => Some(value),
142            Self::Model(_) | Self::Top | Self::Latest => None,
143        }
144    }
145
146    /// Borrow the concrete model id for `Model` targets.
147    pub fn model(&self) -> Option<&str> {
148        match self {
149            Self::Model(value) => Some(value),
150            Self::Org(_) | Self::Top | Self::Latest => None,
151        }
152    }
153
154    /// Stable target kind label.
155    pub fn kind(&self) -> &'static str {
156        match self {
157            Self::Org(_) => "org",
158            Self::Model(_) => "model",
159            Self::Top => "top",
160            Self::Latest => "latest",
161        }
162    }
163
164    /// Borrow the contained target id for `Org` or `Model` targets.
165    pub fn value(&self) -> Option<&str> {
166        match self {
167            Self::Org(value) | Self::Model(value) => Some(value),
168            Self::Top | Self::Latest => None,
169        }
170    }
171}
172
173/// Errors returned while parsing model targets.
174#[derive(Debug, Clone, PartialEq, Eq)]
175pub enum ModelForgeError {
176    /// The provided target string could not be normalized.
177    InvalidTarget(String),
178}
179
180impl ModelForgeError {
181    /// The invalid target payload or reason string.
182    pub fn target(&self) -> &str {
183        match self {
184            Self::InvalidTarget(target) => target,
185        }
186    }
187}
188
189impl fmt::Display for ModelForgeError {
190    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
191        match self {
192            Self::InvalidTarget(target) => write!(f, "invalid model target '{target}'"),
193        }
194    }
195}
196
197impl Error for ModelForgeError {}
198
199/// Parse a Hugging Face-style target into `Org` or `Model`.
200///
201/// This is provider-neutral only for simple `owner` and `owner/model` forms.
202/// Prefer `parse_model_target_for_provider` when the provider is known.
203pub fn parse_model_target(input: &str) -> Result<ModelTarget> {
204    let cleaned = input.trim();
205    if cleaned.is_empty() {
206        return Err(ModelForgeError::InvalidTarget(
207            "empty model target".to_string(),
208        ));
209    }
210
211    if let Some(path) = strip_any_prefix_ignore_ascii_case(
212        cleaned,
213        &[
214            "https://huggingface.co/",
215            "http://huggingface.co/",
216            "https://www.huggingface.co/",
217            "http://www.huggingface.co/",
218            "https://hf.co/",
219            "http://hf.co/",
220            "https://www.hf.co/",
221            "http://www.hf.co/",
222        ],
223    ) {
224        return parse_model_path(path);
225    }
226
227    let parts = cleaned
228        .split('/')
229        .filter(|part| !part.is_empty())
230        .collect::<Vec<_>>();
231    match parts.len() {
232        0 => Err(ModelForgeError::InvalidTarget(cleaned.to_string())),
233        1 => Ok(ModelTarget::Org(parts[0].to_string())),
234        _ => Ok(ModelTarget::Model(format!("{}/{}", parts[0], parts[1]))),
235    }
236}
237
238/// Parse a model target using provider-specific rules.
239///
240/// Accepted presets for all providers are `top`, `trending`, `popular`,
241/// `latest`, `newest`, and `new`.
242///
243/// CivitAI additionally accepts numeric ids and `https://civitai.com/models/<id>`
244/// URLs as concrete model targets.
245pub fn parse_model_target_for_provider(
246    input: &str,
247    provider: ModelProvider,
248) -> Result<ModelTarget> {
249    let normalized = input.trim();
250    let canonical = normalized.to_lowercase();
251
252    match canonical.as_str() {
253        "top" | "trending" | "popular" => return Ok(ModelTarget::Top),
254        "latest" | "newest" | "new" => return Ok(ModelTarget::Latest),
255        _ => {}
256    }
257
258    if normalized.is_empty() {
259        return Err(ModelForgeError::InvalidTarget(
260            "empty model target".to_string(),
261        ));
262    }
263
264    if let ModelProvider::Civitai = provider {
265        if let Some(path) = strip_any_prefix_ignore_ascii_case(
266            normalized,
267            &[
268                "https://civitai.com/models/",
269                "http://civitai.com/models/",
270                "https://www.civitai.com/models/",
271                "http://www.civitai.com/models/",
272            ],
273        ) {
274            let id = path.split(['/', '?', '#']).next().unwrap_or("").trim();
275            if !id.is_empty() && id.chars().all(|c| c.is_ascii_digit()) {
276                return Ok(ModelTarget::Model(id.to_string()));
277            }
278        }
279
280        if normalized.chars().all(|c| c.is_ascii_digit()) {
281            return Ok(ModelTarget::Model(normalized.to_string()));
282        }
283    }
284
285    if let ModelProvider::Ollama = provider {
286        if let Some(path) = strip_any_prefix_ignore_ascii_case(
287            normalized,
288            &[
289                "https://ollama.com/library/",
290                "http://ollama.com/library/",
291                "https://www.ollama.com/library/",
292                "http://www.ollama.com/library/",
293            ],
294        ) {
295            let id = first_clean_path_segment(path);
296            if !id.is_empty() {
297                return Ok(ModelTarget::Model(id.to_string()));
298            }
299        }
300    }
301
302    let model_target = parse_model_target(normalized)?;
303    match provider {
304        ModelProvider::Civitai => match model_target {
305            ModelTarget::Model(model) if model.chars().all(|c| c.is_ascii_digit()) => {
306                Ok(ModelTarget::Model(model))
307            }
308            other => Ok(other),
309        },
310        _ => Ok(model_target),
311    }
312}
313
314fn parse_model_path(path: &str) -> Result<ModelTarget> {
315    let cleaned = path.split(['?', '#']).next().unwrap_or(path);
316    let parts = cleaned
317        .trim_matches('/')
318        .split('/')
319        .filter(|part| !part.is_empty())
320        .collect::<Vec<_>>();
321    match parts.len() {
322        0 => Err(ModelForgeError::InvalidTarget(path.to_string())),
323        1 => Ok(ModelTarget::Org(parts[0].to_string())),
324        _ => Ok(ModelTarget::Model(format!("{}/{}", parts[0], parts[1]))),
325    }
326}
327
328fn first_clean_path_segment(path: &str) -> &str {
329    path.split(['/', '?', '#']).next().unwrap_or("").trim()
330}
331
332fn strip_any_prefix_ignore_ascii_case<'a>(value: &'a str, prefixes: &[&str]) -> Option<&'a str> {
333    prefixes
334        .iter()
335        .find_map(|prefix| strip_prefix_ignore_ascii_case(value, prefix))
336}
337
338fn strip_prefix_ignore_ascii_case<'a>(value: &'a str, prefix: &str) -> Option<&'a str> {
339    let candidate = value.as_bytes().get(..prefix.len())?;
340    if candidate.eq_ignore_ascii_case(prefix.as_bytes()) {
341        value.get(prefix.len()..)
342    } else {
343        None
344    }
345}
346
347#[cfg(test)]
348mod tests {
349    use super::*;
350
351    #[test]
352    fn parses_hugging_face_targets() {
353        assert_eq!(
354            parse_model_target_for_provider("meta-llama/Llama-2-7b", ModelProvider::Hf),
355            Ok(ModelTarget::Model("meta-llama/Llama-2-7b".to_string()))
356        );
357        assert_eq!(
358            parse_model_target_for_provider("top", ModelProvider::Hf),
359            Ok(ModelTarget::Top)
360        );
361        assert_eq!(
362            parse_model_target_for_provider(
363                "https://huggingface.co/openai-community/gpt2",
364                ModelProvider::Hf,
365            ),
366            Ok(ModelTarget::Model("openai-community/gpt2".to_string()))
367        );
368        assert_eq!(
369            parse_model_target_for_provider(
370                "https://hf.co/openai-community/gpt2",
371                ModelProvider::Hf
372            ),
373            Ok(ModelTarget::Model("openai-community/gpt2".to_string()))
374        );
375        assert_eq!(
376            parse_model_target_for_provider(
377                "HTTPS://WWW.HUGGINGFACE.CO/openai-community/gpt2",
378                ModelProvider::Hf,
379            ),
380            Ok(ModelTarget::Model("openai-community/gpt2".to_string()))
381        );
382    }
383
384    #[test]
385    fn parses_civitai_targets() {
386        assert_eq!(
387            parse_model_target_for_provider("827184", ModelProvider::Civitai),
388            Ok(ModelTarget::Model("827184".to_string()))
389        );
390        assert_eq!(
391            parse_model_target_for_provider("newest", ModelProvider::Civitai),
392            Ok(ModelTarget::Latest)
393        );
394        assert_eq!(
395            parse_model_target_for_provider(
396                "https://www.civitai.com/models/827184",
397                ModelProvider::Civitai,
398            ),
399            Ok(ModelTarget::Model("827184".to_string()))
400        );
401        assert_eq!(
402            parse_model_target_for_provider("realistic-vision", ModelProvider::Civitai),
403            Ok(ModelTarget::Org("realistic-vision".to_string()))
404        );
405    }
406
407    #[test]
408    fn parses_ollama_targets() {
409        assert_eq!(
410            parse_model_target_for_provider("llama3", ModelProvider::Ollama),
411            Ok(ModelTarget::Org("llama3".to_string()))
412        );
413        assert_eq!(
414            parse_model_target_for_provider(
415                "HTTPS://WWW.OLLAMA.COM/library/llama3.2/tags",
416                ModelProvider::Ollama,
417            ),
418            Ok(ModelTarget::Model("llama3.2".to_string()))
419        );
420        assert_eq!(
421            parse_model_target_for_provider("trending", ModelProvider::Ollama),
422            Ok(ModelTarget::Top)
423        );
424    }
425
426    #[test]
427    fn strips_url_query_and_fragment_noise() {
428        assert_eq!(
429            parse_model_target_for_provider(
430                "https://huggingface.co/openai-community/gpt2?library=transformers#files",
431                ModelProvider::Hf,
432            ),
433            Ok(ModelTarget::Model("openai-community/gpt2".to_string()))
434        );
435        assert_eq!(
436            parse_model_target_for_provider(
437                "https://ollama.com/library/deepseek-r1?sort=popular",
438                ModelProvider::Ollama,
439            ),
440            Ok(ModelTarget::Model("deepseek-r1".to_string()))
441        );
442    }
443
444    #[test]
445    fn provider_slugs_round_trip() {
446        assert_eq!(ModelProvider::from_slug("hf"), Some(ModelProvider::Hf));
447        assert_eq!(
448            ModelProvider::from_slug("hugging-face"),
449            Some(ModelProvider::Hf)
450        );
451        assert_eq!(
452            ModelProvider::from_slug("hugging face"),
453            Some(ModelProvider::Hf)
454        );
455        assert_eq!(
456            ModelProvider::from_slug("civit_ai"),
457            Some(ModelProvider::Civitai)
458        );
459        assert_eq!(ModelProvider::from_slug("olm"), Some(ModelProvider::Ollama));
460        assert_eq!(
461            ModelProvider::from_slug("civ"),
462            Some(ModelProvider::Civitai)
463        );
464        assert_eq!(ModelProvider::from_slug("unknown"), None);
465        assert_eq!(ModelProvider::Civitai.to_string(), "civitai");
466    }
467
468    #[test]
469    fn target_helpers_describe_targets() {
470        let model = ModelTarget::Model("owner/model".to_string());
471        assert!(model.is_model());
472        assert!(!model.is_org());
473        assert!(!model.is_preset());
474        assert_eq!(model.kind(), "model");
475        assert_eq!(model.model(), Some("owner/model"));
476        assert_eq!(model.org(), None);
477        assert_eq!(model.value(), Some("owner/model"));
478
479        let org = ModelTarget::Org("owner".to_string());
480        assert!(org.is_org());
481        assert!(!org.is_model());
482        assert_eq!(org.org(), Some("owner"));
483        assert_eq!(org.model(), None);
484
485        let latest = ModelTarget::Latest;
486        assert!(latest.is_preset());
487        assert_eq!(latest.kind(), "latest");
488        assert_eq!(latest.value(), None);
489    }
490
491    #[test]
492    fn providers_iterate_in_stable_display_order() {
493        assert_eq!(ModelProvider::all(), &ModelProvider::ALL);
494        assert_eq!(ModelProvider::all()[0], ModelProvider::Hf);
495        assert_eq!(ModelProvider::all()[1], ModelProvider::Ollama);
496        assert_eq!(ModelProvider::all()[2], ModelProvider::Civitai);
497    }
498}
modelforge/lib.rs

modelforge/
lib.rs
