Skip to main content

apple_code_assistant/config/
mod.rs

1//! Configuration loading (env + config file)
2
3use std::collections::HashMap;
4use std::path::{Path, PathBuf};
5
6use serde::{Deserialize, Serialize};
7
8use crate::error::ConfigError;
9
10const CONFIG_DIR_NAME: &str = "apple-code-assistant";
11const CONFIG_FILE_NAME: &str = "config.toml";
12
13/// Prompt template configuration (inspired by smartcat's prompts.toml).
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct PromptConfig {
16    /// Optional override for the model.
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub model: Option<String>,
19    /// Optional override for temperature.
20    #[serde(skip_serializing_if = "Option::is_none")]
21    pub temperature: Option<f32>,
22    /// Optional system prompt snippet to prepend to the default system behavior.
23    #[serde(skip_serializing_if = "Option::is_none")]
24    pub system: Option<String>,
25}
26
27/// Application configuration.
28#[derive(Debug, Clone, Serialize, Deserialize)]
29pub struct Config {
30    #[serde(skip_serializing_if = "Option::is_none")]
31    pub model: Option<String>,
32    #[serde(skip_serializing_if = "Option::is_none")]
33    pub default_language: Option<String>,
34    #[serde(skip_serializing_if = "Option::is_none")]
35    pub theme: Option<String>,
36    #[serde(skip_serializing_if = "Option::is_none")]
37    pub max_tokens: Option<u32>,
38    #[serde(skip_serializing_if = "Option::is_none")]
39    pub temperature: Option<f32>,
40    /// Name of the default prompt template (if any).
41    #[serde(skip_serializing_if = "Option::is_none")]
42    pub default_prompt: Option<String>,
43    /// Prompt templates loaded from config file.
44    #[serde(skip_serializing_if = "Option::is_none")]
45    pub prompts: Option<HashMap<String, PromptConfig>>,
46    #[serde(skip)]
47    pub config_file: Option<PathBuf>,
48}
49
50impl Default for Config {
51    fn default() -> Self {
52        Self {
53            model: None,
54            default_language: Some("typescript".to_string()),
55            theme: Some("dark".to_string()),
56            max_tokens: Some(4000),
57            temperature: Some(0.7),
58            default_prompt: None,
59            prompts: None,
60            config_file: None,
61        }
62    }
63}
64
65impl Config {
66    /// Default config file path under user config dir.
67    pub fn default_config_path() -> Option<PathBuf> {
68        dirs::config_dir().map(|d| d.join(CONFIG_DIR_NAME).join(CONFIG_FILE_NAME))
69    }
70
71    /// Load config: first .env (if present), then config file (if path given or default exists), then defaults.
72    pub fn load(config_file_override: Option<&str>) -> Result<Self, ConfigError> {
73        let _ = dotenvy::dotenv();
74        let mut config = Self::default();
75
76        let config_path: Option<PathBuf> = config_file_override
77            .map(PathBuf::from)
78            .or_else(Self::default_config_path);
79
80        if let Some(ref path) = config_path {
81            if path.exists() {
82                let content = std::fs::read_to_string(path).map_err(|e| ConfigError::Io(e))?;
83                let file_config: ConfigFile = toml::from_str(&content).map_err(|e| ConfigError::Parse(e.to_string()))?;
84                config.merge_file(file_config);
85                config.config_file = Some(path.clone());
86            }
87        }
88
89        config.merge_env();
90        Ok(config)
91    }
92
93    fn merge_file(&mut self, f: ConfigFile) {
94        if f.model.is_some() {
95            self.model = f.model;
96        }
97        if f.default_language.is_some() {
98            self.default_language = f.default_language;
99        }
100        if f.theme.is_some() {
101            self.theme = f.theme;
102        }
103        if f.max_tokens.is_some() {
104            self.max_tokens = f.max_tokens;
105        }
106        if f.temperature.is_some() {
107            self.temperature = f.temperature;
108        }
109        if f.default_prompt.is_some() {
110            self.default_prompt = f.default_prompt;
111        }
112        if let Some(prompts) = f.prompts {
113            // Simple overwrite for now; could be merged per key later if needed.
114            self.prompts = Some(prompts);
115        }
116    }
117
118    fn merge_env(&mut self) {
119        if let Ok(v) = std::env::var("APPLE_FOUNDATION_MODEL") {
120            if !v.is_empty() {
121                self.model = Some(v);
122            }
123        }
124        if let Ok(v) = std::env::var("APPLE_CODE_DEFAULT_LANGUAGE") {
125            if !v.is_empty() {
126                self.default_language = Some(v);
127            }
128        }
129        if let Ok(v) = std::env::var("APPLE_CODE_THEME") {
130            if !v.is_empty() {
131                self.theme = Some(v);
132            }
133        }
134        if let Ok(v) = std::env::var("APPLE_CODE_MAX_TOKENS") {
135            if let Ok(n) = v.parse::<u32>() {
136                self.max_tokens = Some(n);
137            }
138        }
139        if let Ok(v) = std::env::var("APPLE_CODE_TEMPERATURE") {
140            if let Ok(n) = v.parse::<f32>() {
141                self.temperature = Some(n);
142            }
143        }
144    }
145
146    /// Save current config to the given path (or default path).
147    pub fn save(&self, path: Option<&Path>) -> Result<(), ConfigError> {
148        let path = path
149            .map(PathBuf::from)
150            .or_else(Self::default_config_path)
151            .ok_or_else(|| ConfigError::Invalid("no config path".to_string()))?;
152        if let Some(parent) = path.parent() {
153            std::fs::create_dir_all(parent).map_err(ConfigError::Io)?;
154        }
155        let file = ConfigFile {
156            model: self.model.clone(),
157            default_language: self.default_language.clone(),
158            theme: self.theme.clone(),
159            max_tokens: self.max_tokens,
160            temperature: self.temperature,
161            default_prompt: self.default_prompt.clone(),
162            prompts: self.prompts.clone(),
163        };
164        let toml = toml::to_string_pretty(&file).map_err(|e| ConfigError::Invalid(e.to_string()))?;
165        std::fs::write(&path, toml).map_err(ConfigError::Io)?;
166        Ok(())
167    }
168
169    /// Get a value by key name (for --get).
170    pub fn get(&self, key: &str) -> Option<String> {
171        match key {
172            "model" | "APPLE_FOUNDATION_MODEL" => self.model.clone(),
173            "default_language" | "APPLE_CODE_DEFAULT_LANGUAGE" => self.default_language.clone(),
174            "theme" | "APPLE_CODE_THEME" => self.theme.clone(),
175            "max_tokens" | "APPLE_CODE_MAX_TOKENS" => self.max_tokens.map(|n| n.to_string()),
176            "temperature" | "APPLE_CODE_TEMPERATURE" => self.temperature.map(|n| n.to_string()),
177            _ => None,
178        }
179    }
180
181    /// Set a value by key (for --set key=value).
182    pub fn set(&mut self, key: &str, value: &str) -> Result<(), ConfigError> {
183        match key {
184            "model" | "APPLE_FOUNDATION_MODEL" => self.model = Some(value.to_string()),
185            "default_language" | "APPLE_CODE_DEFAULT_LANGUAGE" => self.default_language = Some(value.to_string()),
186            "theme" | "APPLE_CODE_THEME" => self.theme = Some(value.to_string()),
187            "max_tokens" | "APPLE_CODE_MAX_TOKENS" => {
188                self.max_tokens = Some(value.parse().map_err(|_| ConfigError::Invalid(format!("invalid number: {}", value)))?);
189            }
190            "temperature" | "APPLE_CODE_TEMPERATURE" => {
191                self.temperature = Some(value.parse().map_err(|_| ConfigError::Invalid(format!("invalid number: {}", value)))?);
192            }
193            _ => return Err(ConfigError::Invalid(format!("unknown key: {}", key))),
194        }
195        Ok(())
196    }
197
198    /// All keys for --list.
199    pub fn keys() -> &'static [&'static str] {
200        &["model", "default_language", "theme", "max_tokens", "temperature"]
201    }
202}
203
204#[derive(Debug, Deserialize, Serialize)]
205struct ConfigFile {
206    model: Option<String>,
207    default_language: Option<String>,
208    theme: Option<String>,
209    max_tokens: Option<u32>,
210    temperature: Option<f32>,
211     /// Optional default prompt template name.
212    default_prompt: Option<String>,
213    /// Optional map of prompt templates.
214    prompts: Option<HashMap<String, PromptConfig>>,
215}
216
217impl Config {
218    /// Resolve a prompt template either by explicit name or by using default_prompt.
219    pub fn resolve_prompt<'a>(&'a self, name: Option<&str>) -> Option<(&'a str, &'a PromptConfig)> {
220        let prompts = self.prompts.as_ref()?;
221        let key = if let Some(name) = name {
222            name
223        } else if let Some(default) = self.default_prompt.as_deref() {
224            default
225        } else {
226            return None;
227        };
228        prompts
229            .get_key_value(key)
230            .map(|(k, v)| (k.as_str(), v))
231    }
232}