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        _ => None,
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91    use crate::config::Config;
92
93    #[test]
94    fn resolve_with_prefix() {
95        let config = Config::default();
96        let resolved = resolve_model("openai/gpt-4o", &config).unwrap();
97        assert_eq!(resolved.provider, "openai");
98        assert_eq!(resolved.model, "gpt-4o");
99    }
100
101    #[test]
102    fn resolve_with_default_provider() {
103        let config = Config {
104            default_provider: Some("openai".to_string()),
105            ..Config::default()
106        };
107        let resolved = resolve_model("gpt-4o", &config).unwrap();
108        assert_eq!(resolved.provider, "openai");
109        assert_eq!(resolved.model, "gpt-4o");
110    }
111}
litellm_rust/router.rs

litellm_rust/
router.rs
