1use std::path::{Path, PathBuf};
10
11use serde::{Deserialize, Serialize};
12
13#[derive(Debug)]
15pub enum TunerError {
16 ReadConfigDir(String),
18 ParseTOML(PathBuf, String),
20 UserConfigExists(PathBuf),
22 ModelNotFound(String),
24}
25
26impl std::fmt::Display for TunerError {
27 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28 match self {
29 Self::ReadConfigDir(path) => write!(f, "Cannot read config dir: {}", path),
30 Self::ParseTOML(path, msg) => write!(f, "Failed to parse {}: {}", path.display(), msg),
31 Self::UserConfigExists(path) => {
32 write!(
33 f,
34 "User config exists at {}, refusing to overwrite",
35 path.display()
36 )
37 }
38 Self::ModelNotFound(model) => write!(f, "Model not found: {}", model),
39 }
40 }
41}
42
43impl std::error::Error for TunerError {}
44
45#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
55pub struct ModelTuning {
56 pub model: String,
58
59 #[serde(default, skip_serializing_if = "Option::is_none")]
61 pub num_ctx: Option<u32>,
62
63 #[serde(default, skip_serializing_if = "Option::is_none")]
65 pub context_window: Option<u32>,
66
67 #[serde(default, skip_serializing_if = "Option::is_none")]
69 pub real_context_discovery: Option<bool>,
70
71 #[serde(default, skip_serializing_if = "Option::is_none")]
73 pub mid_loop_trim_threshold: Option<usize>,
74
75 #[serde(default, skip_serializing_if = "Option::is_none")]
77 pub mid_loop_trim_tokens: Option<usize>,
78
79 #[serde(default, skip_serializing_if = "Option::is_none")]
81 pub max_tool_rounds: Option<usize>,
82
83 #[serde(default, skip_serializing_if = "Option::is_none")]
85 pub workflow_grace_rounds: Option<usize>,
86
87 #[serde(default, skip_serializing_if = "Option::is_none")]
89 pub narration_nudge_cap: Option<usize>,
90}
91
92pub fn load_model_config(model_name: &str) -> Result<toml::Value, TunerError> {
96 let home = dirs::home_dir()
97 .ok_or_else(|| TunerError::ReadConfigDir("Cannot determine home directory".to_string()))?;
98 let config_dir = home.join(".newt").join("model_tuning");
99
100 let user_config_path = config_dir.join(format!("{}.toml", model_name));
102 if user_config_path.exists() {
103 return load_toml(&user_config_path)
104 .map_err(|e| TunerError::ParseTOML(user_config_path, e.to_string()));
105 }
106
107 let default_path = config_dir.join(format!("{}.default.toml", model_name));
109 if default_path.exists() {
110 return load_toml(&default_path)
111 .map_err(|e| TunerError::ParseTOML(default_path, e.to_string()));
112 }
113
114 Err(TunerError::ModelNotFound(model_name.to_string()))
115}
116
117pub fn init_model_tuning() -> Result<(), TunerError> {
119 let home = dirs::home_dir()
120 .ok_or_else(|| TunerError::ReadConfigDir("Cannot determine home directory".to_string()))?;
121 let config_dir = home.join(".newt").join("model_tuning");
122
123 if config_dir.exists() {
124 return Ok(());
125 }
126
127 std::fs::create_dir_all(&config_dir).map_err(|e| {
128 TunerError::ReadConfigDir(format!("Cannot create {}: {}", config_dir.display(), e))
129 })?;
130
131 Ok(())
133}
134
135pub fn model_tuning_dir() -> Result<PathBuf, TunerError> {
137 let home = dirs::home_dir()
138 .ok_or_else(|| TunerError::ReadConfigDir("Cannot determine home directory".to_string()))?;
139 Ok(home.join(".newt").join("model_tuning"))
140}
141
142fn load_toml(path: &Path) -> Result<toml::Value, String> {
143 let content = std::fs::read_to_string(path).map_err(|e| e.to_string())?;
144 toml::from_str(&content).map_err(|e| e.to_string())
145}
146
147#[cfg(test)]
148mod tests {
149 use super::*;
150
151 #[test]
152 fn test_model_tuning_dir() {
153 let dir = model_tuning_dir().expect("Should return config dir");
154 assert!(dir.ends_with(".newt/model_tuning"));
155 }
156}