litellm_rust/
router.rs

1use crate::config::{Config, ProviderConfig, ProviderKind};
2use crate::error::{LiteLLMError, Result};
3
4#[derive(Debug, Clone)]
5pub struct ResolvedModel {
6    pub provider: String,
7    pub model: String,
8    pub config: ProviderConfig,
9}
10
11pub fn resolve_model(model: &str, config: &Config) -> Result<ResolvedModel> {
12    let (maybe_provider, maybe_model) = split_model(model);
13
14    if let Some(provider) = maybe_provider {
15        let provider_cfg = config
16            .providers
17            .get(provider)
18            .cloned()
19            .or_else(|| default_provider_config(provider));
20        let provider_cfg =
21            provider_cfg.ok_or_else(|| LiteLLMError::ProviderNotFound(provider.to_string()))?;
22        return Ok(ResolvedModel {
23            provider: provider.to_string(),
24            model: maybe_model.unwrap_or(model).to_string(),
25            config: provider_cfg,
26        });
27    }
28
29    let default_provider = config.default_provider.as_deref().unwrap_or("openai");
30    let provider_cfg = config
31        .providers
32        .get(default_provider)
33        .cloned()
34        .or_else(|| default_provider_config(default_provider))
35        .ok_or_else(|| LiteLLMError::ProviderNotFound(default_provider.to_string()))?;
36
37    Ok(ResolvedModel {
38        provider: default_provider.to_string(),
39        model: model.to_string(),
40        config: provider_cfg,
41    })
42}
43
44fn split_model(model: &str) -> (Option<&str>, Option<&str>) {
45    if let Some((provider, model_name)) = model.split_once('/') {
46        return (Some(provider), Some(model_name));
47    }
48    (None, None)
49}
50
51fn default_provider_config(provider: &str) -> Option<ProviderConfig> {
52    let mut cfg = ProviderConfig::default();
53    match provider {
54        "openai" => {
55            cfg.base_url = Some("https://api.openai.com/v1".to_string());
56            cfg.api_key_env = Some("OPENAI_API_KEY".to_string());
57            cfg.kind = ProviderKind::OpenAICompatible;
58            Some(cfg)
59        }
60        "openrouter" => {
61            cfg.base_url = Some("https://openrouter.ai/api/v1".to_string());
62            cfg.api_key_env = Some("OPENROUTER_API_KEY".to_string());
63            cfg.kind = ProviderKind::OpenAICompatible;
64            Some(cfg)
65        }
66        "anthropic" => {
67            cfg.base_url = Some("https://api.anthropic.com".to_string());
68            cfg.api_key_env = Some("ANTHROPIC_API_KEY".to_string());
69            cfg.kind = ProviderKind::Anthropic;
70            Some(cfg)
71        }
72        "gemini" => {
73            cfg.base_url = Some("https://generativelanguage.googleapis.com/v1beta".to_string());
74            cfg.api_key_env = Some("GEMINI_API_KEY".to_string());
75            cfg.kind = ProviderKind::Gemini;
76            Some(cfg)
77        }
78        "xai" => {
79            cfg.base_url = Some("https://api.x.ai/v1".to_string());
80            cfg.api_key_env = Some("XAI_API_KEY".to_string());
81            cfg.kind = ProviderKind::OpenAICompatible;
82            Some(cfg)
83        }
84        "zai" => {
85            cfg.base_url = Some("https://api.z.ai/api/paas/v4".to_string());
86            cfg.api_key_env = Some("ZAI_API_KEY".to_string());
87            cfg.kind = ProviderKind::OpenAICompatible;
88            Some(cfg)
89        }
90        _ => None,
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97    use crate::config::Config;
98
99    #[test]
100    fn resolve_with_prefix() {
101        let config = Config::default();
102        let resolved = resolve_model("openai/gpt-4o", &config).unwrap();
103        assert_eq!(resolved.provider, "openai");
104        assert_eq!(resolved.model, "gpt-4o");
105    }
106
107    #[test]
108    fn resolve_with_default_provider() {
109        let config = Config {
110            default_provider: Some("openai".to_string()),
111            ..Config::default()
112        };
113        let resolved = resolve_model("gpt-4o", &config).unwrap();
114        assert_eq!(resolved.provider, "openai");
115        assert_eq!(resolved.model, "gpt-4o");
116    }
117
118    #[test]
119    fn resolve_zai_with_builtin_openai_compatible_config() {
120        let resolved = resolve_model("zai/glm-5.2", &Config::default()).unwrap();
121
122        assert_eq!(resolved.provider, "zai");
123        assert_eq!(resolved.model, "glm-5.2");
124        assert_eq!(resolved.config.kind, ProviderKind::OpenAICompatible);
125        assert_eq!(
126            resolved.config.base_url.as_deref(),
127            Some("https://api.z.ai/api/paas/v4")
128        );
129        assert_eq!(resolved.config.api_key_env.as_deref(), Some("ZAI_API_KEY"));
130    }
131}
litellm_rust/router.rs

litellm_rust/
router.rs
