1use agent_base::{AgentError, AgentResult};
7
8const DEFAULT_MODEL: &str = "copilot";
9const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1";
10
11#[derive(Clone, Debug)]
13pub struct LlmConfig {
14 pub api_key: String,
16 pub model: String,
18 pub base_url: String,
20}
21
22pub fn resolve_llm_config(model: Option<&str>, base_url: Option<&str>) -> AgentResult<LlmConfig> {
26 let api_key =
27 super::optional_env("LLM_API_KEY").or_else(|| super::optional_env("OPENAI_API_KEY")).ok_or_else(|| {
28 AgentError::config_error("Missing environment variable LLM_API_KEY. Please configure it in .env.")
29 })?;
30
31 let resolved_model = model
32 .map(|s| s.to_string())
33 .or_else(|| super::optional_env("LLM_MODEL"))
34 .or_else(|| super::optional_env("OPENAI_MODEL"))
35 .unwrap_or_else(|| DEFAULT_MODEL.to_string());
36
37 let resolved_base_url = base_url
38 .map(|s| s.to_string())
39 .or_else(|| super::optional_env("LLM_BASE_URL"))
40 .or_else(|| super::optional_env("OPENAI_BASE_URL"))
41 .unwrap_or_else(|| DEFAULT_BASE_URL.to_string());
42
43 Ok(LlmConfig { api_key, model: resolved_model, base_url: resolved_base_url })
44}
45
46#[cfg(test)]
47mod tests {
48 use super::*;
49
50 struct EnvGuard {
51 keys: Vec<&'static str>,
52 saved: Vec<Option<String>>,
53 }
54
55 impl EnvGuard {
56 fn new(keys: &[&'static str]) -> Self {
57 let saved: Vec<Option<String>> = keys.iter().map(|k| std::env::var(k).ok()).collect();
58 for k in keys {
59 unsafe { std::env::remove_var(k) };
60 }
61 Self { keys: keys.to_vec(), saved }
62 }
63 }
64
65 impl Drop for EnvGuard {
66 fn drop(&mut self) {
67 for (i, k) in self.keys.iter().enumerate() {
68 unsafe { std::env::remove_var(k) };
69 if let Some(ref v) = self.saved[i] {
70 unsafe { std::env::set_var(k, v) };
71 }
72 }
73 }
74 }
75
76 #[test]
78 fn test_env_var_resolution_chain() {
79 let vars = &["LLM_API_KEY", "OPENAI_API_KEY", "LLM_MODEL", "OPENAI_MODEL", "LLM_BASE_URL", "OPENAI_BASE_URL"];
80 let _guard = EnvGuard::new(vars);
81
82 let set = |k: &str, v: &str| unsafe { std::env::set_var(k, v) };
83 let rm = |k: &str| unsafe { std::env::remove_var(k) };
84
85 assert!(resolve_llm_config(None, None).is_err());
87 let api_err = resolve_llm_config(None, None).unwrap_err();
89 assert!(matches!(api_err, agent_base::AgentError::ConfigError(_)));
90
91 set("LLM_API_KEY", "sk-llm");
93 let cfg = resolve_llm_config(None, None).unwrap();
94 assert_eq!(cfg.api_key, "sk-llm");
95 assert_eq!(cfg.model, DEFAULT_MODEL);
96 rm("LLM_API_KEY");
97
98 set("OPENAI_API_KEY", "sk-openai");
100 let cfg = resolve_llm_config(None, None).unwrap();
101 assert_eq!(cfg.api_key, "sk-openai");
102 rm("OPENAI_API_KEY");
103
104 set("LLM_API_KEY", "sk-llm");
106 set("OPENAI_API_KEY", "sk-openai");
107 let cfg = resolve_llm_config(None, None).unwrap();
108 assert_eq!(cfg.api_key, "sk-llm");
109 rm("LLM_API_KEY");
110 rm("OPENAI_API_KEY");
111
112 set("LLM_API_KEY", "sk-test");
114 set("LLM_MODEL", "env-model");
115 let cfg = resolve_llm_config(Some("cli-model"), None).unwrap();
116 assert_eq!(cfg.model, "cli-model");
117 rm("LLM_MODEL");
118
119 set("LLM_MODEL", "gpt-4");
121 let cfg = resolve_llm_config(None, None).unwrap();
122 assert_eq!(cfg.model, "gpt-4");
123 rm("LLM_MODEL");
124
125 set("OPENAI_MODEL", "gpt-3.5");
127 let cfg = resolve_llm_config(None, None).unwrap();
128 assert_eq!(cfg.model, "gpt-3.5");
129 rm("OPENAI_MODEL");
130
131 let cfg = resolve_llm_config(None, None).unwrap();
133 assert_eq!(cfg.model, DEFAULT_MODEL);
134
135 set("LLM_BASE_URL", "https://env.example.com/v1");
137 let cfg = resolve_llm_config(None, Some("https://cli.example.com/v1")).unwrap();
138 assert_eq!(cfg.base_url, "https://cli.example.com/v1");
139 rm("LLM_BASE_URL");
140
141 set("LLM_BASE_URL", "https://llm.example.com/v1");
143 let cfg = resolve_llm_config(None, None).unwrap();
144 assert_eq!(cfg.base_url, "https://llm.example.com/v1");
145 rm("LLM_BASE_URL");
146
147 let cfg = resolve_llm_config(None, None).unwrap();
149 assert_eq!(cfg.base_url, DEFAULT_BASE_URL);
150
151 set("LLM_MODEL", "");
153 let cfg = resolve_llm_config(None, None).unwrap();
154 assert_eq!(cfg.model, DEFAULT_MODEL);
155 rm("LLM_MODEL");
156
157 set("OPENAI_BASE_URL", "https://openai.example.com/v1");
159 let cfg = resolve_llm_config(None, None).unwrap();
160 assert_eq!(cfg.base_url, "https://openai.example.com/v1");
161 rm("OPENAI_BASE_URL");
162
163 set("LLM_API_KEY", "sk-test");
164 }
166
167 #[test]
168 fn test_llm_config_debug_clone() {
169 let cfg = LlmConfig {
170 api_key: "sk-test".into(),
171 model: "gpt-4".into(),
172 base_url: "https://api.openai.com/v1".into(),
173 };
174 let cloned = cfg.clone();
175 assert_eq!(cloned.api_key, "sk-test");
176 assert_eq!(cloned.model, "gpt-4");
177 let _ = format!("{:?}", cfg);
179 }
180}