1#![doc = include_str!("../docs/catalog.md")]
2
3use crate::providers::local::discovery::discover_local_models;
4use crate::{ReasoningDisabledSupport, ReasoningEffort};
5use std::borrow::Cow;
6use std::fmt;
7use std::str::FromStr;
8use std::sync::LazyLock;
9
10mod bedrock;
11mod model_spec;
12pub mod transport;
13
14pub use crate::usage::ModelPricing;
15pub use bedrock::BedrockModel;
16pub use model_spec::{ModelSpec, ModelSpecError, ReasoningEffortError, validate_reasoning_effort};
17pub use transport::ModelTransport;
18
19include!(concat!(env!("OUT_DIR"), "/generated.rs"));
20
21pub fn available_models() -> Vec<LlmModel> {
23 LlmModel::all()
24 .iter()
25 .filter(|m| m.required_env_var().is_none_or(|var| std::env::var(var).is_ok()))
26 .cloned()
27 .collect()
28}
29
30pub async fn get_local_models() -> Vec<LlmModel> {
32 let mut models = available_models();
33 let local = discover_local_models().await;
34 models.extend(local);
35 models
36}
37
38impl LlmModel {
39 pub(crate) fn parse_for(provider: Provider, model_id: &str) -> Result<Self, String> {
40 match provider {
41 Provider::Ollama => Ok(Self::Ollama(model_id.to_string())),
42 Provider::LlamaCpp => Ok(Self::LlamaCpp(model_id.to_string())),
43 Provider::Bedrock => model_id.parse().map(Self::Bedrock),
44 provider => Self::all()
45 .iter()
46 .find(|model| model.provider_enum() == provider && model.model_id() == model_id)
47 .cloned()
48 .ok_or_else(|| format!("Unknown {provider} model: '{model_id}'")),
49 }
50 }
51
52 pub fn model_id(&self) -> Cow<'static, str> {
54 match self.source() {
55 ModelSource::Catalog(entry) => Cow::Borrowed(entry.id),
56 ModelSource::Runtime(id) => Cow::Owned(id.to_string()),
57 }
58 }
59
60 pub fn display_name(&self) -> Cow<'static, str> {
62 match self.source() {
63 ModelSource::Catalog(entry) => Cow::Borrowed(entry.display_name),
64 ModelSource::Runtime(id) => Cow::Owned(format!("{} {id}", self.provider_display_name())),
65 }
66 }
67
68 pub fn provider(&self) -> &'static str {
70 self.provider_enum().parser_name()
71 }
72
73 pub fn provider_display_name(&self) -> &'static str {
75 self.provider_enum().display_name()
76 }
77
78 pub fn context_window(&self) -> Option<u32> {
80 self.entry().map(|entry| entry.context_window)
81 }
82
83 pub fn required_env_var(&self) -> Option<&'static str> {
85 self.provider_enum().required_env_var()
86 }
87
88 pub fn oauth_provider_id(&self) -> Option<&'static str> {
90 self.provider_enum().oauth_provider_id()
91 }
92
93 pub fn reasoning_levels(&self) -> &'static [ReasoningEffort] {
95 self.entry().map_or(&[], |entry| entry.reasoning_levels)
96 }
97
98 pub fn reasoning_disabled_support(&self) -> ReasoningDisabledSupport {
99 self.entry().map_or(ReasoningDisabledSupport::Unsupported, |entry| entry.disabled_support)
100 }
101
102 pub fn supports_reasoning_off(&self) -> bool {
104 self.reasoning_levels().contains(&ReasoningEffort::Disabled)
105 }
106
107 pub fn supports_reasoning_off_transport(&self) -> bool {
109 if !self.supports_reasoning_off() {
110 return false;
111 }
112 match self.provider_enum().info().reasoning_off {
113 ReasoningOff::Always => true,
114 ReasoningOff::Effort => self.reasoning_disabled_support() == ReasoningDisabledSupport::Effort,
115 ReasoningOff::ModelTransport => self.transport().is_some(),
116 ReasoningOff::Never => false,
117 }
118 }
119
120 pub fn effective_reasoning_levels(&self) -> Vec<ReasoningEffort> {
122 self.reasoning_levels()
123 .iter()
124 .copied()
125 .filter(|effort| *effort != ReasoningEffort::Disabled || self.supports_reasoning_off_transport())
126 .collect()
127 }
128
129 pub fn validate_reasoning_effort(&self, effort: ReasoningEffort) -> Result<(), ReasoningEffortError> {
130 let supported = self.effective_reasoning_levels();
131 if effort == ReasoningEffort::Default || supported.contains(&effort) {
132 Ok(())
133 } else {
134 Err(ReasoningEffortError::Unsupported { model: self.to_string(), effort, supported })
135 }
136 }
137
138 pub fn supports_reasoning(&self) -> bool {
140 self.reasoning_levels().iter().any(|effort| effort.is_enabled())
141 }
142
143 pub fn supports_prompt_caching(&self) -> bool {
145 self.entry().is_some_and(|entry| entry.supports_prompt_caching)
146 }
147
148 pub fn pricing(&self) -> Option<ModelPricing> {
149 self.entry().and_then(|entry| entry.pricing)
150 }
151
152 pub fn supports_image(&self) -> bool {
154 self.entry().is_some_and(|entry| entry.supports_image)
155 }
156
157 pub fn supports_audio(&self) -> bool {
159 self.entry().is_some_and(|entry| entry.supports_audio)
160 }
161
162 pub fn transport(&self) -> Option<ModelTransport> {
165 self.entry().and_then(|entry| entry.transport)
166 }
167
168 fn entry(&self) -> Option<&'static CatalogEntry> {
169 match self.source() {
170 ModelSource::Catalog(entry) => Some(entry),
171 ModelSource::Runtime(_) => None,
172 }
173 }
174}
175
176impl fmt::Display for LlmModel {
177 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
178 write!(f, "{}:{}", self.provider(), self.model_id())
179 }
180}
181
182impl FromStr for LlmModel {
183 type Err = String;
184
185 fn from_str(s: &str) -> Result<Self, Self::Err> {
187 let (provider, model_id) = s.split_once(':').unwrap_or((s, ""));
188 Self::parse_for(provider.parse()?, model_id)
189 }
190}
191
192impl Provider {
193 pub fn parser_name(self) -> &'static str {
195 self.info().parser_name
196 }
197
198 pub fn genai_provider_name(self) -> &'static str {
200 self.info().genai_provider_name
201 }
202
203 pub fn display_name(self) -> &'static str {
205 self.info().display_name
206 }
207
208 pub fn required_env_var(self) -> Option<&'static str> {
210 self.info().required_env_var
211 }
212
213 pub fn env_vars(self) -> Vec<&'static str> {
214 let overrides: &[&'static str] = match self {
215 Self::Ollama => &["OLLAMA_HOST"],
216 Self::LlamaCpp => &["LLAMA_CPP_HOST"],
217 Self::Bedrock => AWS_ENV_VARS,
218 _ => &[],
219 };
220 self.required_env_var().into_iter().chain(overrides.iter().copied()).collect()
221 }
222
223 pub fn oauth_provider_id(self) -> Option<&'static str> {
225 self.info().oauth_provider_id
226 }
227
228 pub fn is_local(self) -> bool {
231 self.info().is_local
232 }
233
234 fn info(self) -> &'static ProviderInfo {
235 &Self::INFO[self as usize]
236 }
237}
238
239impl fmt::Display for Provider {
240 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
241 f.write_str(self.parser_name())
242 }
243}
244
245impl FromStr for Provider {
246 type Err = String;
247
248 fn from_str(s: &str) -> Result<Self, Self::Err> {
249 Self::ALL
250 .iter()
251 .copied()
252 .find(|provider| provider.parser_name() == s)
253 .ok_or_else(|| format!("Unknown provider: '{s}'"))
254 }
255}
256
257const AWS_ENV_VARS: &[&str] = &[
258 "AWS_ACCESS_KEY_ID",
259 "AWS_BEARER_TOKEN_BEDROCK",
260 "AWS_DEFAULT_REGION",
261 "AWS_PROFILE",
262 "AWS_REGION",
263 "AWS_SECRET_ACCESS_KEY",
264 "AWS_SESSION_TOKEN",
265];
266
267enum ModelSource<'a> {
269 Catalog(&'static CatalogEntry),
271 Runtime(&'a str),
273}
274
275struct CatalogEntry {
277 id: &'static str,
278 display_name: &'static str,
279 context_window: u32,
280 reasoning_levels: &'static [ReasoningEffort],
281 disabled_support: ReasoningDisabledSupport,
282 supports_prompt_caching: bool,
283 pricing: Option<ModelPricing>,
284 supports_image: bool,
285 supports_audio: bool,
286 transport: Option<ModelTransport>,
287}
288
289struct ProviderInfo {
291 parser_name: &'static str,
292 genai_provider_name: &'static str,
293 display_name: &'static str,
294 required_env_var: Option<&'static str>,
295 oauth_provider_id: Option<&'static str>,
296 is_local: bool,
297 reasoning_off: ReasoningOff,
298}
299
300enum ReasoningOff {
302 Always,
304 Effort,
306 ModelTransport,
308 Never,
310}
311
312#[cfg(test)]
313mod tests {
314 use super::*;
315
316 #[test]
317 fn display_fromstr_roundtrip_all_catalog_models() {
318 for model in LlmModel::all() {
319 let s = model.to_string();
320 let parsed: LlmModel = s.parse().unwrap_or_else(|e| panic!("Failed to parse '{s}' back to LlmModel: {e}"));
321 assert_eq!(&parsed, model, "roundtrip failed for '{s}'");
322 }
323 }
324
325 #[test]
326 fn display_fromstr_roundtrip_dynamic_providers() {
327 let cases = [LlmModel::Ollama("llama3.2".to_string()), LlmModel::LlamaCpp("my-model".to_string())];
328 for model in &cases {
329 let s = model.to_string();
330 let parsed: LlmModel = s.parse().unwrap();
331 assert_eq!(&parsed, model);
332 }
333 }
334
335 #[test]
336 fn model_pricing_comes_from_provider_catalog() {
337 let model: LlmModel = "anthropic:claude-sonnet-4-5".parse().unwrap();
338
339 assert_eq!(
340 model.pricing(),
341 Some(ModelPricing {
342 input_per_million: 3.0,
343 output_per_million: 15.0,
344 cache_read_per_million: Some(0.3),
345 cache_write_per_million: Some(3.75),
346 })
347 );
348 }
349
350 #[test]
351 fn models_without_usage_pricing_do_not_inherit_unrelated_prices() {
352 let codex: LlmModel = "codex:gpt-5.5".parse().unwrap();
353 let local: LlmModel = "ollama:local-model".parse().unwrap();
354 let bedrock_profile: LlmModel = "bedrock:us.anthropic.claude-future-model-v99:0".parse().unwrap();
355
356 assert_eq!(codex.pricing(), None);
357 assert_eq!(local.pricing(), None);
358 assert_eq!(bedrock_profile.pricing(), None);
359 }
360
361 #[test]
362 fn codex_gpt56_subscription_models_have_expected_metadata() {
363 let expected_levels = &[
364 ReasoningEffort::Low,
365 ReasoningEffort::Medium,
366 ReasoningEffort::High,
367 ReasoningEffort::Xhigh,
368 ReasoningEffort::Max,
369 ];
370 for id in ["gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna"] {
371 let model: LlmModel = format!("codex:{id}").parse().unwrap();
372 assert_eq!(model.model_id(), id);
373 assert_eq!(model.context_window(), Some(272_000));
374 assert_eq!(model.reasoning_levels(), expected_levels);
375 }
376
377 assert!("codex:gpt-5.6".parse::<LlmModel>().is_err());
378 assert!("codex:gpt-5.1-codex".parse::<LlmModel>().is_err());
379 assert!("openai:gpt-5.6".parse::<LlmModel>().is_ok());
380 }
381
382 #[test]
383 fn codex_gpt6_models_are_available_with_subscription_metadata() {
384 let expected_levels = &[
385 ReasoningEffort::Low,
386 ReasoningEffort::Medium,
387 ReasoningEffort::High,
388 ReasoningEffort::Xhigh,
389 ReasoningEffort::Max,
390 ];
391 for id in ["gpt-6.1-sol", "gpt-6-sol", "gpt-6-astra", "gpt-6-luna"] {
392 let model: LlmModel = format!("codex:{id}").parse().unwrap();
393 assert_eq!(model.model_id(), id);
394 assert_eq!(model.context_window(), Some(272_000));
395 assert_eq!(model.reasoning_levels(), expected_levels);
396 assert_eq!(model.oauth_provider_id(), Some("codex"));
397 }
398 }
399
400 #[test]
401 fn codex_gpt55_uses_subscription_context_window() {
402 let model: LlmModel = "codex:gpt-5.5".parse().unwrap();
403 assert_eq!(model.context_window(), Some(272_000));
404 }
405
406 #[test]
407 fn openai_gpt55_keeps_api_context_window() {
408 let model: LlmModel = "openai:gpt-5.5".parse().unwrap();
409 assert_eq!(model.context_window(), Some(1_050_000));
410 }
411
412 #[test]
413 fn bedrock_foundation_model_parses() {
414 let model: LlmModel = "bedrock:anthropic.claude-sonnet-4-5-20250929-v1:0".parse().unwrap();
415
416 assert_eq!(model.to_string(), "bedrock:anthropic.claude-sonnet-4-5-20250929-v1:0");
417 assert_eq!(model.context_window(), Some(200_000));
418 }
419
420 #[test]
421 fn bedrock_prompt_caching_support_comes_from_catalog() {
422 let claude: LlmModel = "bedrock:anthropic.claude-sonnet-4-5-20250929-v1:0".parse().unwrap();
423 let nova: LlmModel = "bedrock:amazon.nova-lite-v1:0".parse().unwrap();
424 let profile: LlmModel = "bedrock:us.anthropic.claude-future-model-v99:0".parse().unwrap();
425
426 assert!(claude.supports_prompt_caching());
427 assert!(nova.supports_prompt_caching());
428 assert!(!profile.supports_prompt_caching());
429 }
430
431 #[test]
432 fn provider_display_name_returns_human_readable() {
433 let anthropic: LlmModel = "anthropic:claude-opus-4-6".parse().unwrap();
434 assert_eq!(anthropic.provider_display_name(), "Anthropic");
435
436 let bedrock: LlmModel = "bedrock:anthropic.claude-haiku-4-5-20251001-v1:0".parse().unwrap();
437 assert_eq!(bedrock.provider_display_name(), "AWS Bedrock");
438
439 let zai: LlmModel = "zai:glm-4.5".parse().unwrap();
440 assert_eq!(zai.provider_display_name(), "ZAI");
441
442 let ollama = LlmModel::Ollama("llama3.2".to_string());
443 assert_eq!(ollama.provider_display_name(), "Ollama");
444 }
445
446 #[test]
447 fn oauth_provider_id_is_codex_only_for_codex_provider() {
448 let codex: LlmModel = "codex:gpt-5.5".parse().unwrap();
449 assert_eq!(codex.oauth_provider_id(), Some("codex"));
450
451 for non_oauth in
452 ["anthropic:claude-opus-4-6", "openai:gpt-5.5", "bedrock:anthropic.claude-sonnet-4-5-20250929-v1:0"]
453 {
454 let model: LlmModel = non_oauth.parse().unwrap();
455 assert_eq!(model.oauth_provider_id(), None, "{non_oauth} should not have OAuth");
456 }
457 assert_eq!(LlmModel::Ollama("foo".into()).oauth_provider_id(), None);
458 }
459
460 #[test]
461 fn required_env_var_matches_provider() {
462 let cases = [
463 ("anthropic:claude-opus-4-6", Some("ANTHROPIC_API_KEY")),
464 ("openai:gpt-5.5", Some("OPENAI_API_KEY")),
465 ("deepseek:deepseek-v4-flash", Some("DEEPSEEK_API_KEY")),
466 ("gemini:gemini-2.5-pro", Some("GEMINI_API_KEY")),
467 ("openrouter:anthropic/claude-opus-4.6", Some("OPENROUTER_API_KEY")),
468 ("zai:glm-4.5", Some("ZAI_API_KEY")),
469 ("xiaomi:mimo-v2.6-pro", Some("XIAOMI_API_KEY")),
470 ("codex:gpt-5.5", None),
471 ("bedrock:anthropic.claude-sonnet-4-5-20250929-v1:0", None),
472 ];
473 for (input, expected) in cases {
474 let model: LlmModel = input.parse().unwrap();
475 assert_eq!(model.required_env_var(), expected, "{input}");
476 }
477 assert_eq!(LlmModel::Ollama("foo".into()).required_env_var(), None);
478 assert_eq!(LlmModel::LlamaCpp("foo".into()).required_env_var(), None);
479 }
480
481 #[test]
482 fn xiaomi_models_parse_with_catalog_metadata() {
483 let model: LlmModel = "xiaomi:mimo-v2.6-pro".parse().unwrap();
484 assert_eq!(model.model_id(), "mimo-v2.6-pro");
485 assert_eq!(model.context_window(), Some(1_048_576));
486 assert_eq!(model.provider_display_name(), "Xiaomi");
487 assert!(model.supports_reasoning());
488 }
489
490 #[test]
491 fn codex_reasoning_models_include_xhigh_level() {
492 let codex: LlmModel = "codex:gpt-5.5".parse().unwrap();
493 assert_eq!(
494 codex.reasoning_levels(),
495 &[ReasoningEffort::Low, ReasoningEffort::Medium, ReasoningEffort::High, ReasoningEffort::Xhigh]
496 );
497 assert!(codex.supports_reasoning());
498 }
499
500 #[test]
501 fn anthropic_reasoning_models_use_catalog_levels() {
502 let claude: LlmModel = "anthropic:claude-opus-4-6".parse().unwrap();
503 assert_eq!(
504 claude.reasoning_levels(),
505 &[ReasoningEffort::Low, ReasoningEffort::Medium, ReasoningEffort::High, ReasoningEffort::Max,]
506 );
507 }
508
509 #[test]
510 fn dynamic_provider_models_have_no_reasoning_levels() {
511 assert_eq!(LlmModel::Ollama("llama3.2".into()).reasoning_levels(), []);
512 assert!(!LlmModel::Ollama("llama3.2".into()).supports_reasoning());
513 }
514
515 #[test]
516 fn supports_reasoning_matches_reasoning_levels_emptiness() {
517 for model in LlmModel::all() {
518 assert_eq!(model.supports_reasoning(), !model.reasoning_levels().is_empty(), "{model}");
519 }
520 }
521
522 #[test]
523 fn dynamic_providers_have_no_context_window() {
524 assert_eq!(LlmModel::Ollama("foo".into()).context_window(), None);
525 assert_eq!(LlmModel::LlamaCpp("foo".into()).context_window(), None);
526 }
527
528 #[test]
529 fn bedrock_responses_models_carry_their_endpoint_and_wire_shape() {
530 let cases = [
531 ("openai.gpt-6.1-sol", "https://bedrock-mantle.${AWS_REGION}.api.aws/openai/v1"),
532 ("us.openai.gpt-6.1-sol", "https://bedrock-runtime.${AWS_REGION}.amazonaws.com/openai/v1"),
533 ("openai.gpt-5.6-luna", "https://bedrock-mantle.${AWS_REGION}.api.aws/openai/v1"),
534 ("openai.gpt-5.6-sol", "https://bedrock-mantle.${AWS_REGION}.api.aws/openai/v1"),
535 ("openai.gpt-5.6-terra", "https://bedrock-mantle.${AWS_REGION}.api.aws/openai/v1"),
536 ("openai.gpt-6-astra", "https://bedrock-mantle.${AWS_REGION}.api.aws/openai/v1"),
537 ("us.openai.gpt-6-luna", "https://bedrock-runtime.${AWS_REGION}.amazonaws.com/openai/v1"),
538 ("us.openai.gpt-6-sol", "https://bedrock-runtime.${AWS_REGION}.amazonaws.com/openai/v1"),
539 ("us.openai.gpt-6-astra", "https://bedrock-runtime.${AWS_REGION}.amazonaws.com/openai/v1"),
540 ("openai.gpt-5.5", "https://bedrock-mantle.${AWS_REGION}.api.aws/openai/v1"),
541 ("openai.gpt-5.4", "https://bedrock-mantle.${AWS_REGION}.api.aws/openai/v1"),
542 ("xai.grok-4.3", "https://bedrock-mantle.${AWS_REGION}.api.aws/openai/v1"),
543 ("openai.gpt-oss-120b", "https://bedrock-mantle.${AWS_REGION}.api.aws/v1"),
544 ("openai.gpt-oss-20b", "https://bedrock-mantle.${AWS_REGION}.api.aws/v1"),
545 ];
546
547 for (id, expected_api) in cases {
548 let model: LlmModel = format!("bedrock:{id}").parse().unwrap();
549 assert_eq!(
550 model.transport(),
551 Some(ModelTransport::OpenAiResponses { base_url_template: expected_api }),
552 "{id}"
553 );
554 }
555 }
556
557 #[test]
558 fn converse_models_and_profiles_declare_no_transport_override() {
559 for id in [
560 "anthropic.claude-opus-5",
561 "amazon.nova-lite-v1:0",
562 "us.anthropic.claude-future-model-v99:0",
563 "global.openai.gpt-6-luna",
564 "global.openai.gpt-6-sol",
565 "global.openai.gpt-6.1-sol",
566 "global.openai.gpt-6-astra",
567 ] {
568 let model: LlmModel = format!("bedrock:{id}").parse().unwrap();
569 assert_eq!(model.transport(), None, "{id}");
570 }
571 assert_eq!(LlmModel::Ollama("llama3.2".into()).transport(), None);
572 assert_eq!("anthropic:claude-opus-4-6".parse::<LlmModel>().unwrap().transport(), None);
573 }
574
575 #[test]
576 fn every_declared_transport_template_resolves() {
577 for model in LlmModel::all() {
578 let Some(ModelTransport::OpenAiResponses { base_url_template }) = model.transport() else {
579 continue;
580 };
581 let resolved = utils::variables::Vars::new()
582 .with_env_lookup(|_| Some("test-region".to_string()))
583 .expand(base_url_template)
584 .unwrap_or_else(|e| panic!("{model} has an unresolvable endpoint template: {e}"));
585 assert!(resolved.starts_with("https://"), "{model} resolved to '{resolved}'");
586 assert!(!resolved.contains("${"), "{model} left a placeholder in '{resolved}'");
587 }
588 }
589
590 #[test]
591 fn bedrock_profile_fallback_parses_arbitrary_id() {
592 let profile: LlmModel = "bedrock:us.anthropic.future-model-v99:0".parse().unwrap();
593 assert_eq!(profile.context_window(), None);
594 assert_eq!(profile.reasoning_levels(), []);
595 assert!(!profile.supports_prompt_caching());
596 }
597}