1use std::collections::HashMap;
4use std::sync::OnceLock;
5
6use aither_core::llm::model::Ability;
7
8use crate::convert;
9use crate::tier::{ChatPriceBreakpoints, ModelTier, classify_entry};
10use crate::types::{ModelEntry, ModelMode, Pricing, Provider};
11
12include!(concat!(env!("OUT_DIR"), "/generated.rs"));
17
18#[derive(Debug)]
20pub struct ModelRegistry {
21 entries: Vec<ModelEntry>,
22 index: HashMap<String, usize>,
23 chat_price_breakpoints: Option<ChatPriceBreakpoints>,
24}
25
26impl ModelRegistry {
27 #[must_use]
29 pub fn bundled() -> &'static Self {
30 static INSTANCE: OnceLock<ModelRegistry> = OnceLock::new();
31 INSTANCE.get_or_init(|| {
32 let entries: Vec<ModelEntry> = ENTRIES.to_vec();
33 Self::from_entries(entries)
34 })
35 }
36
37 #[must_use]
39 pub fn from_litellm_json(json: &[u8]) -> Self {
40 let converted = convert::parse_litellm_json(json);
41 let entries = converted.into_iter().filter_map(convert_to_entry).collect();
42 Self::from_entries(entries)
43 }
44
45 fn from_entries(entries: Vec<ModelEntry>) -> Self {
46 let mut index = HashMap::with_capacity(entries.len() * 2);
47 for (i, entry) in entries.iter().enumerate() {
52 if entry.litellm_id() == entry.id() {
53 index.entry(entry.id().to_lowercase()).or_insert(i);
54 }
55 }
56 for (i, entry) in entries.iter().enumerate() {
57 if is_first_party(entry) {
58 index.entry(entry.id().to_lowercase()).or_insert(i);
59 }
60 }
61 for (i, entry) in entries.iter().enumerate() {
63 index.entry(entry.litellm_id().to_lowercase()).or_insert(i);
64 }
65 let chat_price_breakpoints = ChatPriceBreakpoints::compute(entries.iter());
66 Self {
67 entries,
68 index,
69 chat_price_breakpoints,
70 }
71 }
72
73 #[must_use]
77 pub fn tier(&self, model_id: &str) -> ModelTier {
78 self.lookup(model_id)
79 .map_or(ModelTier::Balanced, |entry| self.tier_of(entry))
80 }
81
82 #[must_use]
84 pub fn tier_of(&self, entry: &ModelEntry) -> ModelTier {
85 classify_entry(entry, self.chat_price_breakpoints)
86 }
87
88 #[must_use]
99 pub fn lookup(&self, model_id: &str) -> Option<&ModelEntry> {
100 let mut id = model_id.to_lowercase();
101 loop {
102 if let Some(&idx) = self.index.get(id.as_str()) {
103 return Some(&self.entries[idx]);
104 }
105 id = next_lookup_form(&id)?;
106 }
107 }
108
109 pub fn models_for_provider(&self, provider: &Provider) -> impl Iterator<Item = &ModelEntry> {
111 let provider = provider.clone();
112 self.entries
113 .iter()
114 .filter(move |e| *e.provider() == provider)
115 }
116
117 pub fn models_by_mode(&self, mode: ModelMode) -> impl Iterator<Item = &ModelEntry> {
119 self.entries.iter().filter(move |e| e.mode() == mode)
120 }
121
122 pub fn models_with_ability(&self, ability: Ability) -> impl Iterator<Item = &ModelEntry> {
124 self.entries.iter().filter(move |e| e.has_ability(ability))
125 }
126
127 pub fn all(&self) -> impl Iterator<Item = &ModelEntry> {
129 self.entries.iter()
130 }
131
132 #[must_use]
134 pub const fn len(&self) -> usize {
135 self.entries.len()
136 }
137
138 #[must_use]
140 pub const fn is_empty(&self) -> bool {
141 self.entries.is_empty()
142 }
143}
144
145const FIRST_PARTY_PREFIXES: &[&str] = &[
150 "anthropic/",
151 "cohere/",
152 "cohere_chat/",
153 "dashscope/",
154 "deepseek/",
155 "gemini/",
156 "meta-llama/",
157 "mistral/",
158 "openai/",
159 "xai/",
160];
161
162fn is_first_party(entry: &ModelEntry) -> bool {
166 entry.litellm_id() == entry.id()
167 || FIRST_PARTY_PREFIXES
168 .iter()
169 .any(|prefix| entry.litellm_id().starts_with(prefix))
170}
171
172fn next_lookup_form(id: &str) -> Option<String> {
178 if let Some(stripped) = strip_version_suffix(id) {
179 return Some(stripped.to_owned());
180 }
181 let minor = find_v_minor(id)?;
182 if minor.sep_byte == b'-' {
183 let mut dotted = id.to_owned();
184 dotted.replace_range(minor.sep..=minor.sep, ".");
185 return Some(dotted);
186 }
187 let mut collapsed = String::with_capacity(id.len() - (minor.minor_end - minor.sep));
188 collapsed.push_str(&id[..minor.sep]);
189 collapsed.push_str(&id[minor.minor_end..]);
190 Some(collapsed)
191}
192
193fn strip_version_suffix(id: &str) -> Option<&str> {
198 let bytes = id.as_bytes();
199 let n = bytes.len();
200 if let Some((base, suffix)) = id.rsplit_once('-')
201 && let Some(rest) = suffix.strip_prefix('v')
202 && rest
203 .split(':')
204 .all(|p| !p.is_empty() && p.bytes().all(|b| b.is_ascii_digit()))
205 {
206 return Some(base);
207 }
208 if n > 9 && matches!(bytes[n - 9], b'-' | b'@') && bytes[n - 8..].iter().all(u8::is_ascii_digit)
209 {
210 return Some(&id[..n - 9]);
211 }
212 if n > 11
213 && bytes[n - 11] == b'-'
214 && bytes[n - 6] == b'-'
215 && bytes[n - 3] == b'-'
216 && bytes[n - 10..n - 6].iter().all(u8::is_ascii_digit)
217 && bytes[n - 5..n - 3].iter().all(u8::is_ascii_digit)
218 && bytes[n - 2..].iter().all(u8::is_ascii_digit)
219 {
220 return Some(&id[..n - 11]);
221 }
222 if n > 5 && bytes[n - 5] == b'-' && is_compact_month_day(&bytes[n - 4..]) {
223 return Some(&id[..n - 5]);
224 }
225 None
226}
227
228fn is_compact_month_day(digits: &[u8]) -> bool {
230 if digits.len() != 4 || !digits.iter().all(u8::is_ascii_digit) {
231 return false;
232 }
233 let month = (digits[0] - b'0') * 10 + (digits[1] - b'0');
234 let day = (digits[2] - b'0') * 10 + (digits[3] - b'0');
235 (1..=12).contains(&month) && (1..=31).contains(&day)
236}
237
238struct VMinor {
240 sep: usize,
242 minor_end: usize,
244 sep_byte: u8,
246}
247
248fn find_v_minor(id: &str) -> Option<VMinor> {
250 let bytes = id.as_bytes();
251 let mut i = 0;
252 while i < bytes.len() {
253 let at_token_start = i == 0 || bytes[i - 1] == b'-' || bytes[i - 1] == b'/';
254 if bytes[i] != b'v' || !at_token_start {
255 i += 1;
256 continue;
257 }
258 let mut j = i + 1;
259 if j >= bytes.len() || !bytes[j].is_ascii_digit() {
260 i += 1;
261 continue;
262 }
263 while j < bytes.len() && bytes[j].is_ascii_digit() {
264 j += 1;
265 }
266 let Some(&sep_byte) = bytes.get(j) else {
267 i += 1;
268 continue;
269 };
270 if sep_byte != b'-' && sep_byte != b'.' {
271 i += 1;
272 continue;
273 }
274 let mut k = j + 1;
275 let minor_start = k;
276 while k < bytes.len() && bytes[k].is_ascii_digit() {
277 k += 1;
278 }
279 let minor_len = k - minor_start;
280 if !(1..=2).contains(&minor_len) {
281 i += 1;
282 continue;
283 }
284 if k != bytes.len() && bytes[k] != b'-' {
285 i += 1;
286 continue;
287 }
288 return Some(VMinor {
289 sep: j,
290 minor_end: k,
291 sep_byte,
292 });
293 }
294 None
295}
296
297fn convert_to_entry(c: convert::ConvertedEntry) -> Option<ModelEntry> {
299 let provider = parse_provider_enum(&c.provider);
300 let mode = parse_mode_enum(&c.mode)?;
301
302 let mut pricing = crate::types::Pricing::new(c.input_cost_per_token, c.output_cost_per_token);
303 if let Some(v) = c.cache_read_per_token {
304 pricing = pricing.with_cache_read(v);
305 }
306 if let Some(v) = c.cache_write_per_token {
307 pricing = pricing.with_cache_write(v);
308 }
309 if let Some(v) = c.reasoning_per_token {
310 pricing = pricing.with_reasoning(v);
311 }
312 if let Some(v) = c.image_per_token {
313 pricing = pricing.with_image(v);
314 }
315
316 let abilities: Vec<Ability> = c
317 .abilities
318 .iter()
319 .filter_map(|a| parse_ability(a))
320 .collect();
321
322 let mut entry = ModelEntry::new(c.litellm_id, c.canonical_id, provider, mode)
323 .with_pricing(pricing)
324 .with_abilities(abilities);
325
326 if let Some(max_in) = c.max_input_tokens {
327 entry = entry.with_max_input_tokens(max_in);
328 }
329 if let Some(max_out) = c.max_output_tokens {
330 entry = entry.with_max_output_tokens(max_out);
331 }
332 if let Some(date) = c.deprecation_date {
333 entry = entry.with_deprecation_date(date);
334 }
335
336 Some(entry)
337}
338
339fn parse_provider_enum(s: &str) -> Provider {
340 match s {
341 "openai" => Provider::OpenAI,
342 "anthropic" => Provider::Anthropic,
343 "gemini" => Provider::Google,
344 "deepseek" => Provider::DeepSeek,
345 "xai" => Provider::XAI,
346 "mistral" => Provider::Mistral,
347 "meta" => Provider::Meta,
348 "alibaba" => Provider::Alibaba,
349 "bedrock" => Provider::Bedrock,
350 "azure" => Provider::Azure,
351 "vertex_ai" => Provider::VertexAI,
352 "groq" => Provider::Groq,
353 "together_ai" => Provider::Together,
354 "fireworks_ai" => Provider::Fireworks,
355 "replicate" => Provider::Replicate,
356 "cohere" => Provider::Cohere,
357 "perplexity" => Provider::Perplexity,
358 "github_copilot" => Provider::Copilot,
359 other => Provider::Other(std::borrow::Cow::Owned(other.to_string())),
360 }
361}
362
363fn parse_mode_enum(s: &str) -> Option<ModelMode> {
364 match s {
365 "chat" => Some(ModelMode::Chat),
366 "embedding" => Some(ModelMode::Embedding),
367 "image_generation" => Some(ModelMode::ImageGeneration),
368 "audio_transcription" => Some(ModelMode::AudioTranscription),
369 "audio_speech" => Some(ModelMode::AudioSpeech),
370 "rerank" => Some(ModelMode::Rerank),
371 "moderation" => Some(ModelMode::Moderation),
372 "completion" => Some(ModelMode::Completion),
373 "video_generation" => Some(ModelMode::VideoGeneration),
374 "search" => Some(ModelMode::Search),
375 "ocr" => Some(ModelMode::Ocr),
376 _ => None,
377 }
378}
379
380fn parse_ability(s: &str) -> Option<Ability> {
381 match s {
382 "ToolUse" => Some(Ability::ToolUse),
383 "Vision" => Some(Ability::Vision),
384 "Audio" => Some(Ability::Audio),
385 "AudioOutput" => Some(Ability::AudioOutput),
386 "Video" => Some(Ability::Video),
387 "WebSearch" => Some(Ability::WebSearch),
388 "Pdf" => Some(Ability::Pdf),
389 "CodeExecution" => Some(Ability::CodeExecution),
390 "Reasoning" => Some(Ability::Reasoning),
391 "ImageGeneration" => Some(Ability::ImageGeneration),
392 "ComputerUse" => Some(Ability::ComputerUse),
393 "PromptCaching" => Some(Ability::PromptCaching),
394 "AssistantPrefill" => Some(Ability::AssistantPrefill),
395 "PerMessageEffort" => Some(Ability::PerMessageEffort),
396 "MidConversationSystem" => Some(Ability::MidConversationSystem),
397 _ => None,
398 }
399}
400
401#[cfg(test)]
402pub mod tests {
403 use super::*;
404
405 pub fn fixture() -> ModelRegistry {
409 ModelRegistry::from_litellm_json(include_bytes!("../testdata/litellm_fixture.json"))
410 }
411
412 fn litellm_id_of(registry: &ModelRegistry, query: &str) -> Option<String> {
413 registry
414 .lookup(query)
415 .map(|entry| entry.litellm_id().to_owned())
416 }
417
418 #[test]
421 fn mid_conversation_system_follows_the_documented_models() {
422 for (id, expected) in [
423 ("claude-opus-5-5", true),
424 ("claude-opus-4-8", true),
425 ("claude-sonnet-5-5", true),
426 ("claude-fable-5", true),
427 ("claude-sonnet-5", false),
428 ("gpt-6-sol", false),
429 ] {
430 let entry = crate::lookup(id).unwrap_or_else(|| panic!("{id} is in the snapshot"));
431 assert_eq!(
432 entry.has_ability(Ability::MidConversationSystem),
433 expected,
434 "{id}"
435 );
436 }
437 }
438
439 #[test]
442 fn per_message_effort_follows_the_documented_models_and_platforms() {
443 use crate::convert::takes_per_message_effort;
444 for (id, expected) in [
445 ("claude-opus-5-5", true),
446 ("claude-sonnet-5-5", true),
447 ("claude-fable-5-1", true),
448 ("claude-sonnet-5", false),
449 ("claude-fable-5", false),
450 ("gpt-6-sol", true),
451 ("gpt-6.1-sol", true),
452 ("gpt-5", false),
453 ] {
454 let entry = crate::lookup(id).unwrap_or_else(|| panic!("{id} is in the snapshot"));
455 assert_eq!(
456 entry.has_ability(Ability::PerMessageEffort),
457 expected,
458 "{id}"
459 );
460 }
461 assert!(takes_per_message_effort(
462 "vertex_ai",
463 "claude-opus-5-5@default"
464 ));
465 assert!(!takes_per_message_effort(
466 "bedrock",
467 "anthropic.claude-opus-5-5"
468 ));
469 assert!(!takes_per_message_effort("azure", "gpt-6-sol"));
470 }
471
472 #[test]
473 fn fixture_skips_sample_spec_and_unknown_modes() {
474 let registry = fixture();
475 assert_eq!(registry.len(), 13);
476 assert!(registry.lookup("sample_spec").is_none());
477 assert!(registry.lookup("fixture-unknown-mode").is_none());
478 }
479
480 #[test]
481 fn exact_lookup_returns_converted_entry() {
482 let registry = fixture();
483 let entry = registry.lookup("fixture-gpt").expect("fixture-gpt exists");
484 assert_eq!(entry.id(), "fixture-gpt");
485 assert_eq!(*entry.provider(), Provider::OpenAI);
486 assert_eq!(entry.mode(), ModelMode::Chat);
487 assert_eq!(entry.max_input_tokens(), Some(128_000));
488 assert!(entry.has_ability(Ability::Vision));
489 assert!(entry.has_ability(Ability::ToolUse));
490 assert!(entry.pricing().input_per_token() > 0.0);
491 assert!(entry.pricing().output_per_token() > 0.0);
492 }
493
494 #[test]
495 fn lookup_is_case_insensitive() {
496 let registry = fixture();
497 assert_eq!(
498 litellm_id_of(®istry, "FIXTURE-GPT").as_deref(),
499 Some("fixture-gpt")
500 );
501 assert_eq!(
502 litellm_id_of(®istry, "GITHUB_COPILOT/FIXTURE-HOSTED").as_deref(),
503 Some("github_copilot/fixture-hosted")
504 );
505 }
506
507 #[test]
508 fn dated_snapshot_with_own_record_matches_exactly() {
509 let registry = fixture();
510 assert_eq!(
511 litellm_id_of(®istry, "fixture-gpt-2024-05-13").as_deref(),
512 Some("fixture-gpt-2024-05-13")
513 );
514 }
515
516 #[test]
517 fn version_suffixes_normalize_to_base_model() {
518 let registry = fixture();
519 for alias in [
520 "fixture-gpt-21000101",
521 "fixture-gpt-2099-01-01",
522 "fixture-gpt@20240513",
523 "fixture-gpt-v2",
524 "fixture-gpt-v2:1",
525 "fixture-gpt-1017",
526 ] {
527 assert_eq!(
528 litellm_id_of(®istry, alias).as_deref(),
529 Some("fixture-gpt"),
530 "{alias} should normalize to fixture-gpt"
531 );
532 }
533 assert_eq!(
535 litellm_id_of(®istry, "fixture-gpt-mini-20250101").as_deref(),
536 Some("fixture-gpt-mini")
537 );
538 }
539
540 #[test]
541 fn hyphenated_v_minor_and_compact_date_normalize_to_base() {
542 let registry = fixture();
543 for alias in [
544 "fixture-v4-flash-1017",
545 "fixture-v4-1-flash",
546 "fixture-v4.1-flash",
547 "fixture-v4-1-flash-1017",
548 ] {
549 assert_eq!(
550 litellm_id_of(®istry, alias).as_deref(),
551 Some("fixture-v4-flash"),
552 "{alias} should normalize to fixture-v4-flash"
553 );
554 assert_eq!(
555 registry
556 .lookup(alias)
557 .and_then(ModelEntry::max_input_tokens),
558 Some(1_000_000),
559 "{alias} should inherit the 1M window"
560 );
561 }
562 }
563
564 #[test]
565 fn hyphenated_v_minor_prefers_dotted_record_when_present() {
566 let registry = ModelRegistry::from_litellm_json(
567 br#"{
568 "fixture-v4-flash": {
569 "litellm_provider": "deepseek",
570 "mode": "chat",
571 "max_input_tokens": 1000000,
572 "input_cost_per_token": 3e-07,
573 "output_cost_per_token": 1.2e-06
574 },
575 "fixture-v4.1-flash": {
576 "litellm_provider": "deepseek",
577 "mode": "chat",
578 "max_input_tokens": 1048576,
579 "input_cost_per_token": 3e-07,
580 "output_cost_per_token": 1.2e-06
581 }
582 }"#,
583 );
584 assert_eq!(
585 litellm_id_of(®istry, "fixture-v4-1-flash").as_deref(),
586 Some("fixture-v4.1-flash")
587 );
588 assert_eq!(
589 registry
590 .lookup("fixture-v4-1-flash")
591 .and_then(ModelEntry::max_input_tokens),
592 Some(1_048_576)
593 );
594 }
595
596 #[test]
597 fn compact_month_day_must_be_a_valid_calendar_date() {
598 let registry = fixture();
599 assert!(registry.lookup("fixture-gpt-1301").is_none());
600 assert!(registry.lookup("fixture-gpt-0001").is_none());
601 assert!(registry.lookup("fixture-gpt-4096").is_none());
602 }
603
604 #[test]
605 fn numeric_minor_without_v_prefix_does_not_collapse() {
606 let registry = fixture();
607 assert!(registry.lookup("fixture-gpt-4-1").is_none());
608 assert!(registry.lookup("fixture-gpt-4.1").is_none());
609 }
610
611 #[test]
612 fn normalization_never_changes_model_family() {
613 let registry = fixture();
614 assert!(registry.lookup("fixture-gpt-turbo").is_none());
615 assert!(registry.lookup("fixture-gpt-mini-turbo").is_none());
616 }
617
618 #[test]
619 fn unknown_id_returns_none() {
620 let registry = fixture();
621 assert!(registry.lookup("no-such-model-xyz").is_none());
622 assert!(registry.lookup("no-such-model-xyz-20250101").is_none());
623 }
624
625 #[test]
626 fn bare_id_does_not_bind_provider_scoped_record() {
627 let registry = fixture();
628 assert!(registry.lookup("fixture-hosted").is_none());
631 assert!(registry.lookup("fixture-hosted-20250514").is_none());
632 assert!(registry.lookup("fixture-azure-only").is_none());
633 assert!(registry.lookup("fixture-bedrock-v1:0").is_none());
634 }
635
636 #[test]
637 fn provider_scoped_query_normalizes_within_its_namespace() {
638 let registry = fixture();
639 assert_eq!(
640 litellm_id_of(®istry, "github_copilot/fixture-hosted-20250514").as_deref(),
641 Some("github_copilot/fixture-hosted")
642 );
643 assert_eq!(
644 litellm_id_of(®istry, "azure/fixture-azure-only").as_deref(),
645 Some("azure/fixture-azure-only")
646 );
647 assert_eq!(
648 litellm_id_of(®istry, "bedrock/fixture-bedrock-v1:0").as_deref(),
649 Some("bedrock/fixture-bedrock-v1:0")
650 );
651 assert!(registry.lookup("azure/fixture-hosted").is_none());
653 }
654
655 #[test]
656 fn first_party_prefix_claims_bare_id() {
657 let registry = fixture();
658 assert_eq!(
659 litellm_id_of(®istry, "fixture-gemini").as_deref(),
660 Some("gemini/fixture-gemini")
661 );
662 assert_eq!(
663 litellm_id_of(®istry, "fixture-gemini-2099-01-01").as_deref(),
664 Some("gemini/fixture-gemini")
665 );
666 }
667
668 #[test]
669 fn unprefixed_record_wins_bare_id_over_first_party_prefix() {
670 let registry = fixture();
671 assert_eq!(
674 litellm_id_of(®istry, "fixture-claude-mid").as_deref(),
675 Some("fixture-claude-mid")
676 );
677 assert_eq!(
678 litellm_id_of(®istry, "anthropic/fixture-claude-mid").as_deref(),
679 Some("anthropic/fixture-claude-mid")
680 );
681 }
682
683 #[test]
684 fn query_helpers_filter_by_mode_ability_and_provider() {
685 let registry = fixture();
686 let image: Vec<_> = registry
687 .models_by_mode(ModelMode::ImageGeneration)
688 .map(ModelEntry::litellm_id)
689 .collect();
690 assert_eq!(image, ["fixture-image"]);
691 let reasoning: Vec<_> = registry
692 .models_with_ability(Ability::Reasoning)
693 .map(ModelEntry::litellm_id)
694 .collect();
695 assert_eq!(reasoning, ["fixture-flagship"]);
696 let copilot: Vec<_> = registry
697 .models_for_provider(&Provider::Copilot)
698 .map(ModelEntry::litellm_id)
699 .collect();
700 assert_eq!(copilot, ["github_copilot/fixture-hosted"]);
701 }
702}