limit_llm/
config.rs

1use crate::error::ConfigError;
2use serde::Deserialize;
3use std::collections::HashMap;
4use std::{env, fs, io};
5
6#[derive(Debug, Deserialize, PartialEq, Clone)]
7pub struct Config {
8    pub provider: String,
9    pub providers: HashMap<String, ProviderConfig>,
10}
11
12#[derive(Debug, Deserialize, PartialEq, Clone)]
13pub struct ProviderConfig {
14    pub api_key: Option<String>,
15    pub model: String,
16    #[serde(default)]
17    pub base_url: Option<String>,
18    #[serde(default = "default_max_tokens")]
19    pub max_tokens: u32,
20    #[serde(default = "default_timeout")]
21    pub timeout: u64,
22    /// Maximum iterations for agent loop (0 = unlimited, default: 100)
23    #[serde(default = "default_max_iterations")]
24    pub max_iterations: usize,
25    /// Enable thinking/reasoning mode (Z.AI only, default: false)
26    #[serde(default)]
27    pub thinking_enabled: bool,
28    /// Preserve thinking between turns (Z.AI only, default: true)
29    /// Set to false for Preserved Thinking in multi-turn conversations
30    #[serde(default = "default_clear_thinking")]
31    pub clear_thinking: bool,
32}
33
34fn default_model() -> String {
35    "claude-3-5-sonnet-20241022".to_string()
36}
37
38fn default_max_tokens() -> u32 {
39    4096
40}
41
42fn default_timeout() -> u64 {
43    60
44}
45
46fn default_max_iterations() -> usize {
47    100
48}
49
50fn default_clear_thinking() -> bool {
51    true
52}
53
54impl ProviderConfig {
55    pub fn api_key_or_env(&self, provider: &str) -> Option<String> {
56        if let Some(key) = &self.api_key {
57            return Some(key.clone());
58        }
59
60        match provider {
61            "anthropic" => env::var("ANTHROPIC_API_KEY").ok(),
62            "openai" => env::var("OPENAI_API_KEY")
63                .ok()
64                .or_else(|| env::var("ZAI_API_KEY").ok()),
65            "zai" => env::var("ZAI_API_KEY").ok(),
66            // Local providers don't require API key, return placeholder
67            "local" | "ollama" | "lmstudio" | "vllm" => Some("local".to_string()),
68            _ => None,
69        }
70    }
71}
72
73impl Config {
74    pub fn validate(&self) -> Result<(), ConfigError> {
75        // Valid provider names (includes local LLM aliases)
76        let valid_providers = [
77            "anthropic",
78            "openai",
79            "zai",
80            "local",
81            "ollama",
82            "lmstudio",
83            "vllm",
84        ];
85
86        // Check provider field is valid
87        if !valid_providers.contains(&self.provider.as_str()) {
88            return Err(ConfigError::InvalidProvider(self.provider.clone()));
89        }
90
91        // Check provider config exists
92        if !self.providers.contains_key(&self.provider) {
93            return Err(ConfigError::MissingProvider(self.provider.clone()));
94        }
95
96        // Local providers don't require API key
97        let local_providers = ["local", "ollama", "lmstudio", "vllm"];
98        if local_providers.contains(&self.provider.as_str()) {
99            return Ok(());
100        }
101
102        // Check active provider has required fields
103        let provider_config = self.providers.get(&self.provider).unwrap();
104        if provider_config.api_key.is_none() {
105            // Check if env var exists
106            let env_var = match self.provider.as_str() {
107                "anthropic" => "ANTHROPIC_API_KEY",
108                "openai" => "OPENAI_API_KEY",
109                "zai" => "ZAI_API_KEY",
110                _ => "API_KEY",
111            };
112            if env::var(env_var).is_err() {
113                // For openai, also check ZAI_API_KEY as fallback
114                let env_var_display = if self.provider == "openai" {
115                    "OPENAI_API_KEY or ZAI_API_KEY"
116                } else {
117                    env_var
118                };
119                return Err(ConfigError::MissingApiKey {
120                    provider: self.provider.clone(),
121                    env_var: env_var_display.to_string(),
122                });
123            }
124        }
125
126        Ok(())
127    }
128}
129
130impl Config {
131    pub fn load() -> Result<Self, io::Error> {
132        let config_path = config_path();
133
134        if !config_path.exists() {
135            return Ok(Config::default());
136        }
137
138        let config_content = fs::read_to_string(&config_path)?;
139
140        // Check for old format (detect by presence of 'api_key' at top level)
141        if config_content.contains("api_key") && !config_content.contains("[providers.") {
142            return Err(io::Error::new(
143                io::ErrorKind::InvalidData,
144                "Old config format detected. Please update to multi-provider format.\n\nNew format:\nprovider = \"anthropic\"\n\n[providers.anthropic]\nmodel = \"claude-3-5-sonnet-20241022\"\napi_key = \"...\"  # optional, falls back to ANTHROPIC_API_KEY env var"
145            ));
146        }
147
148        let config: Config = toml::from_str(&config_content)
149            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
150
151        config
152            .validate()
153            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
154
155        Ok(config)
156    }
157}
158
159impl Default for Config {
160    fn default() -> Self {
161        let mut providers = HashMap::new();
162        providers.insert(
163            "anthropic".to_string(),
164            ProviderConfig {
165                api_key: None,
166                model: default_model(),
167                base_url: None,
168                max_tokens: default_max_tokens(),
169                timeout: default_timeout(),
170                max_iterations: default_max_iterations(),
171                thinking_enabled: false,
172                clear_thinking: true,
173            },
174        );
175        Config {
176            provider: "anthropic".to_string(),
177            providers,
178        }
179    }
180}
181
182fn config_path() -> std::path::PathBuf {
183    let home_dir = dirs::home_dir().expect("Failed to get home directory");
184    home_dir.join(".limit").join("config.toml")
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190
191    #[test]
192    fn test_load_from_actual_config() {
193        // Load config from actual path (tests loading with existing file)
194        let config = Config::load().unwrap();
195
196        // Should have loaded a valid config (provider and config entry exist)
197        assert!(!config.provider.is_empty());
198        assert!(config.providers.contains_key(&config.provider));
199    }
200
201    #[test]
202    fn test_load_valid_config() {
203        let config_content = r#"
204provider = "anthropic"
205
206[providers.anthropic]
207api_key = "sk-ant-test123"
208model = "claude-3-5-sonnet-20241022"
209"#;
210
211        let config: Config = toml::from_str(config_content).unwrap();
212
213        assert_eq!(config.provider, "anthropic");
214        assert!(config.providers.contains_key("anthropic"));
215        let anthropic = config.providers.get("anthropic").unwrap();
216        assert_eq!(anthropic.api_key, Some("sk-ant-test123".to_string()));
217        assert_eq!(anthropic.model, "claude-3-5-sonnet-20241022");
218    }
219
220    #[test]
221    fn test_load_partial_config_uses_defaults() {
222        let config_content = r#"
223provider = "anthropic"
224
225[providers.anthropic]
226api_key = "sk-ant-partial"
227model = "custom-model"
228"#;
229
230        let config: Config = toml::from_str(config_content).unwrap();
231
232        assert_eq!(config.provider, "anthropic");
233        let anthropic = config.providers.get("anthropic").unwrap();
234        assert_eq!(anthropic.api_key, Some("sk-ant-partial".to_string()));
235        assert_eq!(anthropic.model, "custom-model");
236        assert!(anthropic.base_url.is_none()); // default
237    }
238
239    #[test]
240    fn test_load_config_with_base_url() {
241        let config_content = r#"
242provider = "openai"
243
244[providers.openai]
245api_key = "sk-test123"
246model = "gpt-4"
247base_url = "https://api.z.ai/api/paas/v4/chat/completions"
248"#;
249
250        let config: Config = toml::from_str(config_content).unwrap();
251
252        assert_eq!(config.provider, "openai");
253        let openai = config.providers.get("openai").unwrap();
254        assert_eq!(openai.api_key, Some("sk-test123".to_string()));
255        assert_eq!(openai.model, "gpt-4");
256        assert_eq!(
257            openai.base_url,
258            Some("https://api.z.ai/api/paas/v4/chat/completions".to_string())
259        );
260    }
261
262    #[test]
263    fn test_load_config_without_base_url() {
264        let config_content = r#"
265provider = "anthropic"
266
267[providers.anthropic]
268api_key = "sk-ant-test456"
269model = "claude-3-5-sonnet-20241022"
270"#;
271
272        let config: Config = toml::from_str(config_content).unwrap();
273
274        assert_eq!(config.provider, "anthropic");
275        let anthropic = config.providers.get("anthropic").unwrap();
276        assert_eq!(anthropic.api_key, Some("sk-ant-test456".to_string()));
277        assert_eq!(anthropic.model, "claude-3-5-sonnet-20241022");
278        assert!(anthropic.base_url.is_none()); // default
279    }
280
281    #[test]
282    fn test_default_config() {
283        let config = Config::default();
284
285        assert_eq!(config.provider, "anthropic");
286        assert!(config.providers.contains_key("anthropic"));
287        let anthropic = config.providers.get("anthropic").unwrap();
288        assert_eq!(anthropic.model, "claude-3-5-sonnet-20241022");
289        assert!(anthropic.api_key.is_none());
290        assert!(anthropic.base_url.is_none());
291    }
292
293    #[test]
294    fn test_old_format_detection() {
295        let config_content = r#"
296api_key = "sk-ant-test123"
297model = "claude-3-5-sonnet-20241022"
298"#;
299
300        let result: Result<Config, _> = toml::from_str(config_content);
301        assert!(result.is_err(), "Old format should fail to parse");
302    }
303
304    #[test]
305    fn test_api_key_or_env_from_config() {
306        let provider_config = ProviderConfig {
307            api_key: Some("sk-from-config".to_string()),
308            model: "claude-3-5-sonnet-20241022".to_string(),
309            base_url: None,
310            max_tokens: 4096,
311            timeout: 60,
312            max_iterations: 100,
313            thinking_enabled: false,
314            clear_thinking: true,
315        };
316
317        let key = provider_config.api_key_or_env("anthropic");
318        assert_eq!(key, Some("sk-from-config".to_string()));
319    }
320
321    #[test]
322    fn test_api_key_or_env_from_env() {
323        env::set_var("ANTHROPIC_API_KEY", "sk-from-env");
324        let provider_config = ProviderConfig {
325            api_key: None,
326            model: "claude-3-5-sonnet-20241022".to_string(),
327            base_url: None,
328            max_tokens: 4096,
329            timeout: 60,
330            max_iterations: 100,
331            thinking_enabled: false,
332            clear_thinking: true,
333        };
334
335        let key = provider_config.api_key_or_env("anthropic");
336        assert_eq!(key, Some("sk-from-env".to_string()));
337        env::remove_var("ANTHROPIC_API_KEY");
338    }
339
340    #[test]
341    fn test_openai_fallback_to_zai_api_key() {
342        env::set_var("ZAI_API_KEY", "sk-zai-key");
343        let provider_config = ProviderConfig {
344            api_key: None,
345            model: "gpt-4".to_string(),
346            base_url: None,
347            max_tokens: 4096,
348            timeout: 60,
349            max_iterations: 100,
350            thinking_enabled: false,
351            clear_thinking: true,
352        };
353
354        let key = provider_config.api_key_or_env("openai");
355        assert_eq!(key, Some("sk-zai-key".to_string()));
356        env::remove_var("ZAI_API_KEY");
357    }
358
359    #[test]
360    fn test_unknown_provider_no_env_var() {
361        let provider_config = ProviderConfig {
362            api_key: None,
363            model: "test-model".to_string(),
364            base_url: None,
365            max_tokens: 4096,
366            timeout: 60,
367            max_iterations: 100,
368            thinking_enabled: false,
369            clear_thinking: true,
370        };
371
372        let key = provider_config.api_key_or_env("unknown");
373        assert_eq!(key, None);
374    }
375
376    #[test]
377    fn test_zai_config_validation() {
378        env::set_var("ZAI_API_KEY", "test-zai-key");
379        let config_content = r#"
380provider = "zai"
381
382[providers.zai]
383model = "glm-4.7"
384"#;
385        let config: Config = toml::from_str(config_content).unwrap();
386        config.validate().unwrap();
387        env::remove_var("ZAI_API_KEY");
388    }
389    #[test]
390    fn test_zai_api_key_env_var() {
391        env::set_var("ZAI_API_KEY", "test-zai-key");
392        let provider_config = ProviderConfig {
393            api_key: None,
394            model: "glm-4.7".to_string(),
395            base_url: None,
396            max_tokens: 4096,
397            timeout: 60,
398            max_iterations: 100,
399            thinking_enabled: false,
400            clear_thinking: true,
401        };
402        let key = provider_config.api_key_or_env("zai");
403        assert_eq!(key, Some("test-zai-key".to_string()));
404        env::remove_var("ZAI_API_KEY");
405    }
406}
limit_llm/config.rs

limit_llm/
config.rs
