1use serde::{Deserialize, Serialize};
24use std::collections::HashMap;
25
26#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct AliasEntry {
33 pub alias: String,
35 pub target: String,
37 pub default_quant: Option<String>,
39 pub variants: HashMap<String, String>,
41 pub description: Option<String>,
43}
44
45impl AliasEntry {
46 #[must_use]
48 pub fn new(alias: impl Into<String>, target: impl Into<String>) -> Self {
49 Self {
50 alias: alias.into(),
51 target: target.into(),
52 default_quant: None,
53 variants: HashMap::new(),
54 description: None,
55 }
56 }
57
58 #[must_use]
60 pub fn with_default_quant(mut self, quant: impl Into<String>) -> Self {
61 self.default_quant = Some(quant.into());
62 self
63 }
64
65 #[must_use]
67 pub fn with_variant(mut self, tag: impl Into<String>, target: impl Into<String>) -> Self {
68 self.variants.insert(tag.into(), target.into());
69 self
70 }
71
72 #[must_use]
74 pub fn with_description(mut self, desc: impl Into<String>) -> Self {
75 self.description = Some(desc.into());
76 self
77 }
78
79 #[must_use]
81 pub fn resolve_variant(&self, variant: Option<&str>) -> &str {
82 match variant {
83 Some(v) => self.variants.get(v).map(String::as_str).unwrap_or(&self.target),
84 None => &self.target,
85 }
86 }
87}
88
89#[derive(Debug, Clone, PartialEq, Eq)]
95pub struct ParsedRef {
96 pub name: String,
98 pub variant: Option<String>,
100 pub quantization: Option<String>,
102}
103
104impl ParsedRef {
105 #[must_use]
114 pub fn parse(s: &str) -> Self {
115 if s.contains("://") {
117 return Self { name: s.to_string(), variant: None, quantization: None };
118 }
119
120 let (name_part, tag_part) =
122 if let Some(idx) = s.find(':') { (&s[..idx], Some(&s[idx + 1..])) } else { (s, None) };
123
124 let (name, name_quant) = extract_quant_suffix(name_part);
126
127 let (variant, tag_quant) =
129 if let Some(tag) = tag_part { parse_tag(tag) } else { (None, None) };
130
131 let quantization = tag_quant.or(name_quant);
133
134 Self { name: name.to_string(), variant, quantization }
135 }
136
137 #[must_use]
139 pub fn to_string_repr(&self) -> String {
140 let mut s = self.name.clone();
141 if let Some(ref v) = self.variant {
142 s.push(':');
143 s.push_str(v);
144 }
145 if let Some(ref q) = self.quantization {
146 if self.variant.is_some() {
147 s.push('-');
148 } else {
149 s.push(':');
150 }
151 s.push_str(q);
152 }
153 s
154 }
155}
156
157fn extract_quant_suffix(s: &str) -> (&str, Option<String>) {
159 let quant_patterns = [
161 "-q4", "-q5", "-q6", "-q8", "-Q4_K_M", "-Q4_K_S", "-Q5_K_M", "-Q5_K_S", "-Q6_K", "-Q8_0",
162 "-Q8_K", "-f16", "-f32", "-bf16",
163 ];
164
165 for pattern in quant_patterns {
166 if let Some(idx) = s.to_lowercase().rfind(&pattern.to_lowercase()) {
167 return (&s[..idx], Some(s[idx + 1..].to_string()));
168 }
169 }
170
171 (s, None)
172}
173
174fn parse_tag(tag: &str) -> (Option<String>, Option<String>) {
176 if let Some(idx) = tag.rfind('-') {
178 let (variant, quant) = (&tag[..idx], &tag[idx + 1..]);
179 if is_quant_tag(quant) {
180 return (Some(variant.to_string()), Some(quant.to_string()));
181 }
182 }
183
184 if is_quant_tag(tag) {
186 return (None, Some(tag.to_string()));
187 }
188
189 (Some(tag.to_string()), None)
191}
192
193fn is_quant_tag(s: &str) -> bool {
195 let lower = s.to_lowercase();
196 lower.starts_with("q") && lower.chars().nth(1).is_some_and(|c| c.is_ascii_digit())
197 || lower.starts_with("iq")
198 || lower == "f16"
199 || lower == "f32"
200 || lower == "fp16"
201 || lower == "fp32"
202 || lower == "bf16"
203}
204
205#[derive(Debug, Clone, Default)]
211pub struct AliasRegistry {
212 aliases: HashMap<String, AliasEntry>,
214}
215
216impl AliasRegistry {
217 #[must_use]
219 pub fn new() -> Self {
220 Self::default()
221 }
222
223 #[must_use]
225 pub fn with_defaults() -> Self {
226 let mut registry = Self::new();
227
228 registry.add(
230 AliasEntry::new("llama3", "hf://meta-llama/Meta-Llama-3-8B-Instruct")
231 .with_default_quant("Q4_K_M")
232 .with_variant("8b", "hf://meta-llama/Meta-Llama-3-8B-Instruct")
233 .with_variant("70b", "hf://meta-llama/Meta-Llama-3-70B-Instruct")
234 .with_description("Meta's Llama 3 family"),
235 );
236
237 registry.add(
239 AliasEntry::new("llama3.1", "hf://meta-llama/Meta-Llama-3.1-8B-Instruct")
240 .with_default_quant("Q4_K_M")
241 .with_variant("8b", "hf://meta-llama/Meta-Llama-3.1-8B-Instruct")
242 .with_variant("70b", "hf://meta-llama/Meta-Llama-3.1-70B-Instruct")
243 .with_variant("405b", "hf://meta-llama/Meta-Llama-3.1-405B-Instruct")
244 .with_description("Meta's Llama 3.1 family"),
245 );
246
247 registry.add(
249 AliasEntry::new("mistral", "hf://mistralai/Mistral-7B-Instruct-v0.3")
250 .with_default_quant("Q4_K_M")
251 .with_variant("7b", "hf://mistralai/Mistral-7B-Instruct-v0.3")
252 .with_description("Mistral AI's 7B model"),
253 );
254
255 registry.add(
257 AliasEntry::new("mixtral", "hf://mistralai/Mixtral-8x7B-Instruct-v0.1")
258 .with_default_quant("Q4_K_M")
259 .with_variant("8x7b", "hf://mistralai/Mixtral-8x7B-Instruct-v0.1")
260 .with_variant("8x22b", "hf://mistralai/Mixtral-8x22B-Instruct-v0.1")
261 .with_description("Mistral AI's Mixtral MoE"),
262 );
263
264 registry.add(
266 AliasEntry::new("phi3", "hf://microsoft/Phi-3-mini-4k-instruct")
267 .with_default_quant("Q4_K_M")
268 .with_variant("mini", "hf://microsoft/Phi-3-mini-4k-instruct")
269 .with_variant("small", "hf://microsoft/Phi-3-small-8k-instruct")
270 .with_variant("medium", "hf://microsoft/Phi-3-medium-4k-instruct")
271 .with_description("Microsoft's Phi-3 family"),
272 );
273
274 registry.add(
276 AliasEntry::new("gemma", "hf://google/gemma-7b-it")
277 .with_default_quant("Q4_K_M")
278 .with_variant("2b", "hf://google/gemma-2b-it")
279 .with_variant("7b", "hf://google/gemma-7b-it")
280 .with_description("Google's Gemma family"),
281 );
282
283 registry.add(
284 AliasEntry::new("gemma2", "hf://google/gemma-2-9b-it")
285 .with_default_quant("Q4_K_M")
286 .with_variant("2b", "hf://google/gemma-2-2b-it")
287 .with_variant("9b", "hf://google/gemma-2-9b-it")
288 .with_variant("27b", "hf://google/gemma-2-27b-it")
289 .with_description("Google's Gemma 2 family"),
290 );
291
292 registry.add(
294 AliasEntry::new("qwen2", "hf://Qwen/Qwen2-7B-Instruct")
295 .with_default_quant("Q4_K_M")
296 .with_variant("0.5b", "hf://Qwen/Qwen2-0.5B-Instruct")
297 .with_variant("1.5b", "hf://Qwen/Qwen2-1.5B-Instruct")
298 .with_variant("7b", "hf://Qwen/Qwen2-7B-Instruct")
299 .with_variant("72b", "hf://Qwen/Qwen2-72B-Instruct")
300 .with_description("Alibaba's Qwen2 family"),
301 );
302
303 registry.add(
307 AliasEntry::new("qwen3-coder", "hf://unsloth/Qwen3-Coder-30B-A3B-Instruct-GGUF")
308 .with_default_quant("Q4_K_M")
309 .with_variant("30b-a3b", "hf://unsloth/Qwen3-Coder-30B-A3B-Instruct-GGUF")
310 .with_description("Alibaba's Qwen3-Coder — agentic coding default for 24 GB GPUs"),
311 );
312
313 registry.add(
315 AliasEntry::new("codellama", "hf://codellama/CodeLlama-7b-Instruct-hf")
316 .with_default_quant("Q4_K_M")
317 .with_variant("7b", "hf://codellama/CodeLlama-7b-Instruct-hf")
318 .with_variant("13b", "hf://codellama/CodeLlama-13b-Instruct-hf")
319 .with_variant("34b", "hf://codellama/CodeLlama-34b-Instruct-hf")
320 .with_description("Meta's CodeLlama"),
321 );
322
323 registry.add(
325 AliasEntry::new("deepseek-coder", "hf://deepseek-ai/deepseek-coder-6.7b-instruct")
326 .with_default_quant("Q4_K_M")
327 .with_variant("1.3b", "hf://deepseek-ai/deepseek-coder-1.3b-instruct")
328 .with_variant("6.7b", "hf://deepseek-ai/deepseek-coder-6.7b-instruct")
329 .with_variant("33b", "hf://deepseek-ai/deepseek-coder-33b-instruct")
330 .with_description("DeepSeek AI Coder"),
331 );
332
333 registry.add(
335 AliasEntry::new("starcoder2", "hf://bigcode/starcoder2-7b")
336 .with_default_quant("Q4_K_M")
337 .with_variant("3b", "hf://bigcode/starcoder2-3b")
338 .with_variant("7b", "hf://bigcode/starcoder2-7b")
339 .with_variant("15b", "hf://bigcode/starcoder2-15b")
340 .with_description("BigCode's StarCoder 2"),
341 );
342
343 registry.add(
345 AliasEntry::new("nomic-embed", "hf://nomic-ai/nomic-embed-text-v1.5")
346 .with_description("Nomic AI embedding model"),
347 );
348
349 registry.add(
350 AliasEntry::new("bge", "hf://BAAI/bge-large-en-v1.5")
351 .with_variant("small", "hf://BAAI/bge-small-en-v1.5")
352 .with_variant("base", "hf://BAAI/bge-base-en-v1.5")
353 .with_variant("large", "hf://BAAI/bge-large-en-v1.5")
354 .with_description("BGE embedding models"),
355 );
356
357 registry
358 }
359
360 pub fn add(&mut self, entry: AliasEntry) {
362 self.aliases.insert(entry.alias.clone(), entry);
363 }
364
365 #[must_use]
367 pub fn get(&self, alias: &str) -> Option<&AliasEntry> {
368 self.aliases.get(alias)
369 }
370
371 #[must_use]
373 pub fn contains(&self, alias: &str) -> bool {
374 self.aliases.contains_key(alias)
375 }
376
377 #[must_use]
379 pub fn list(&self) -> Vec<&AliasEntry> {
380 let mut entries: Vec<_> = self.aliases.values().collect();
381 entries.sort_by(|a, b| a.alias.cmp(&b.alias));
382 entries
383 }
384
385 #[must_use]
389 pub fn resolve(&self, reference: &str) -> ResolvedAlias {
390 let parsed = ParsedRef::parse(reference);
391
392 if let Some(entry) = self.aliases.get(&parsed.name) {
394 let target = entry.resolve_variant(parsed.variant.as_deref());
395 let quant = parsed.quantization.or_else(|| entry.default_quant.clone());
396
397 ResolvedAlias { uri: target.to_string(), quantization: quant, is_alias: true }
398 } else {
399 let uri = if parsed.name.contains("://") {
405 parsed.name.clone()
406 } else if parsed.name.contains('/') {
407 format!("hf://{}", parsed.name)
409 } else {
410 format!("pacha://{}", parsed.name)
411 };
412
413 ResolvedAlias { uri, quantization: parsed.quantization, is_alias: false }
414 }
415 }
416
417 #[must_use]
419 pub fn len(&self) -> usize {
420 self.aliases.len()
421 }
422
423 #[must_use]
425 pub fn is_empty(&self) -> bool {
426 self.aliases.is_empty()
427 }
428}
429
430#[derive(Debug, Clone, PartialEq, Eq)]
432pub struct ResolvedAlias {
433 pub uri: String,
435 pub quantization: Option<String>,
437 pub is_alias: bool,
439}
440
441#[cfg(test)]
446mod tests {
447 use super::*;
448
449 #[test]
454 fn test_alias_entry_new() {
455 let entry = AliasEntry::new("llama3", "hf://meta-llama/Llama-3-8B");
456 assert_eq!(entry.alias, "llama3");
457 assert_eq!(entry.target, "hf://meta-llama/Llama-3-8B");
458 }
459
460 #[test]
461 fn test_alias_entry_builder() {
462 let entry = AliasEntry::new("llama3", "hf://meta-llama/Llama-3-8B")
463 .with_default_quant("Q4_K_M")
464 .with_variant("70b", "hf://meta-llama/Llama-3-70B")
465 .with_description("Llama 3 family");
466
467 assert_eq!(entry.default_quant, Some("Q4_K_M".to_string()));
468 assert!(entry.variants.contains_key("70b"));
469 assert!(entry.description.is_some());
470 }
471
472 #[test]
473 fn test_alias_entry_resolve_variant() {
474 let entry = AliasEntry::new("llama3", "hf://default")
475 .with_variant("8b", "hf://8b-model")
476 .with_variant("70b", "hf://70b-model");
477
478 assert_eq!(entry.resolve_variant(None), "hf://default");
479 assert_eq!(entry.resolve_variant(Some("8b")), "hf://8b-model");
480 assert_eq!(entry.resolve_variant(Some("70b")), "hf://70b-model");
481 assert_eq!(entry.resolve_variant(Some("unknown")), "hf://default");
482 }
483
484 #[test]
489 fn test_parsed_ref_name_only() {
490 let parsed = ParsedRef::parse("llama3");
491 assert_eq!(parsed.name, "llama3");
492 assert!(parsed.variant.is_none());
493 assert!(parsed.quantization.is_none());
494 }
495
496 #[test]
497 fn test_parsed_ref_with_variant() {
498 let parsed = ParsedRef::parse("llama3:8b");
499 assert_eq!(parsed.name, "llama3");
500 assert_eq!(parsed.variant, Some("8b".to_string()));
501 assert!(parsed.quantization.is_none());
502 }
503
504 #[test]
505 fn test_parsed_ref_with_variant_and_quant() {
506 let parsed = ParsedRef::parse("llama3:8b-q4");
507 assert_eq!(parsed.name, "llama3");
508 assert_eq!(parsed.variant, Some("8b".to_string()));
509 assert_eq!(parsed.quantization, Some("q4".to_string()));
510 }
511
512 #[test]
513 fn test_parsed_ref_quant_only() {
514 let parsed = ParsedRef::parse("llama3:q4");
515 assert_eq!(parsed.name, "llama3");
516 assert!(parsed.variant.is_none());
517 assert_eq!(parsed.quantization, Some("q4".to_string()));
518 }
519
520 #[test]
521 fn test_parsed_ref_uppercase_quant() {
522 let parsed = ParsedRef::parse("llama3:8b-Q4_K_M");
523 assert_eq!(parsed.name, "llama3");
524 assert_eq!(parsed.variant, Some("8b".to_string()));
525 assert_eq!(parsed.quantization, Some("Q4_K_M".to_string()));
526 }
527
528 #[test]
529 fn test_parsed_ref_to_string() {
530 let parsed = ParsedRef::parse("llama3:8b-q4");
531 assert_eq!(parsed.to_string_repr(), "llama3:8b-q4");
532
533 let parsed = ParsedRef::parse("llama3:q4");
534 assert_eq!(parsed.to_string_repr(), "llama3:q4");
535
536 let parsed = ParsedRef::parse("llama3");
537 assert_eq!(parsed.to_string_repr(), "llama3");
538 }
539
540 #[test]
541 fn test_is_quant_tag() {
542 assert!(is_quant_tag("q4"));
543 assert!(is_quant_tag("Q4_K_M"));
544 assert!(is_quant_tag("q8"));
545 assert!(is_quant_tag("f16"));
546 assert!(is_quant_tag("fp16"));
547 assert!(is_quant_tag("bf16"));
548 assert!(is_quant_tag("iq4"));
549 assert!(!is_quant_tag("8b"));
550 assert!(!is_quant_tag("70b"));
551 assert!(!is_quant_tag("instruct"));
552 }
553
554 #[test]
559 fn test_alias_registry_new() {
560 let registry = AliasRegistry::new();
561 assert!(registry.is_empty());
562 assert_eq!(registry.len(), 0);
563 }
564
565 #[test]
566 fn test_alias_registry_with_defaults() {
567 let registry = AliasRegistry::with_defaults();
568 assert!(!registry.is_empty());
569 assert!(registry.contains("llama3"));
570 assert!(registry.contains("mistral"));
571 assert!(registry.contains("phi3"));
572 }
573
574 #[test]
575 fn test_alias_registry_add() {
576 let mut registry = AliasRegistry::new();
577 registry.add(AliasEntry::new("test", "hf://test/model"));
578
579 assert!(registry.contains("test"));
580 assert_eq!(registry.len(), 1);
581 }
582
583 #[test]
584 fn test_alias_registry_get() {
585 let registry = AliasRegistry::with_defaults();
586 let entry = registry.get("llama3");
587
588 assert!(entry.is_some());
589 assert!(entry.unwrap().target.contains("meta-llama"));
590 }
591
592 #[test]
593 fn test_alias_registry_list() {
594 let mut registry = AliasRegistry::new();
595 registry.add(AliasEntry::new("zzz", "hf://zzz"));
596 registry.add(AliasEntry::new("aaa", "hf://aaa"));
597
598 let list = registry.list();
599 assert_eq!(list.len(), 2);
600 assert_eq!(list[0].alias, "aaa"); assert_eq!(list[1].alias, "zzz");
602 }
603
604 #[test]
605 fn test_alias_registry_resolve_known() {
606 let registry = AliasRegistry::with_defaults();
607
608 let resolved = registry.resolve("llama3");
609 assert!(resolved.is_alias);
610 assert!(resolved.uri.contains("meta-llama"));
611 }
612
613 #[test]
614 fn test_alias_registry_resolve_with_variant() {
615 let registry = AliasRegistry::with_defaults();
616
617 let resolved = registry.resolve("llama3:70b");
618 assert!(resolved.is_alias);
619 assert!(resolved.uri.contains("70B"));
620 }
621
622 #[test]
623 fn test_alias_registry_resolve_with_quant() {
624 let registry = AliasRegistry::with_defaults();
625
626 let resolved = registry.resolve("llama3:8b-q8");
627 assert!(resolved.is_alias);
628 assert_eq!(resolved.quantization, Some("q8".to_string()));
629 }
630
631 #[test]
632 fn test_alias_registry_resolve_default_quant() {
633 let registry = AliasRegistry::with_defaults();
634
635 let resolved = registry.resolve("llama3");
636 assert!(resolved.is_alias);
637 assert_eq!(resolved.quantization, Some("Q4_K_M".to_string()));
639 }
640
641 #[test]
642 fn test_alias_registry_resolve_unknown() {
643 let registry = AliasRegistry::with_defaults();
644
645 let resolved = registry.resolve("unknown-model");
646 assert!(!resolved.is_alias);
647 assert_eq!(resolved.uri, "pacha://unknown-model");
648 }
649
650 #[test]
651 fn test_alias_registry_resolve_huggingface_style() {
652 let registry = AliasRegistry::with_defaults();
653
654 let resolved = registry.resolve("Qwen/Qwen2.5-Coder-0.5B-Instruct");
656 assert!(!resolved.is_alias);
657 assert_eq!(resolved.uri, "hf://Qwen/Qwen2.5-Coder-0.5B-Instruct");
658
659 let resolved = registry.resolve("TheBloke/Llama-2-7B-GGUF");
661 assert!(!resolved.is_alias);
662 assert_eq!(resolved.uri, "hf://TheBloke/Llama-2-7B-GGUF");
663 }
664
665 #[test]
666 fn test_alias_registry_resolve_full_uri() {
667 let registry = AliasRegistry::with_defaults();
668
669 let resolved = registry.resolve("hf://some/model");
670 assert!(!resolved.is_alias);
671 assert_eq!(resolved.uri, "hf://some/model");
672 }
673
674 #[test]
679 fn test_alias_entry_serialization() {
680 let entry = AliasEntry::new("llama3", "hf://test")
681 .with_default_quant("Q4_K_M")
682 .with_variant("70b", "hf://test-70b");
683
684 let json = serde_json::to_string(&entry).unwrap();
685 assert!(json.contains("llama3"));
686 assert!(json.contains("Q4_K_M"));
687
688 let parsed: AliasEntry = serde_json::from_str(&json).unwrap();
689 assert_eq!(parsed.alias, "llama3");
690 assert_eq!(parsed.default_quant, Some("Q4_K_M".to_string()));
691 }
692
693 #[test]
698 fn test_parsed_ref_complex_name() {
699 let parsed = ParsedRef::parse("deepseek-coder:6.7b-q4");
700 assert_eq!(parsed.name, "deepseek-coder");
701 assert_eq!(parsed.variant, Some("6.7b".to_string()));
702 assert_eq!(parsed.quantization, Some("q4".to_string()));
703 }
704
705 #[test]
706 fn test_parsed_ref_numbers_in_name() {
707 let parsed = ParsedRef::parse("llama3.1:8b");
708 assert_eq!(parsed.name, "llama3.1");
709 assert_eq!(parsed.variant, Some("8b".to_string()));
710 }
711
712 #[test]
713 fn test_resolve_quant_override() {
714 let registry = AliasRegistry::with_defaults();
715
716 let resolved = registry.resolve("llama3:q8");
718 assert_eq!(resolved.quantization, Some("q8".to_string()));
719 }
720
721 #[test]
722 fn test_gemma_variants() {
723 let registry = AliasRegistry::with_defaults();
724
725 let resolved_2b = registry.resolve("gemma:2b");
726 assert!(resolved_2b.uri.contains("2b"));
727
728 let resolved_7b = registry.resolve("gemma:7b");
729 assert!(resolved_7b.uri.contains("7b"));
730 }
731
732 #[test]
733 fn test_embedding_models() {
734 let registry = AliasRegistry::with_defaults();
735
736 assert!(registry.contains("nomic-embed"));
737 assert!(registry.contains("bge"));
738
739 let resolved = registry.resolve("bge:large");
740 assert!(resolved.uri.contains("large"));
741 }
742}