Skip to main content

newt_tuner/
lib.rs

1// Copyright (c) 2025-2026 newt contributors. Licensed under Apache-2.0.
2
3//! # newt-tuner
4//!
5//! Model tuning configuration and management for Newt-Agent.
6//! Provides per-model TOML configurations that ship defaults
7//! but never overwrite user-customized settings.
8
9use std::path::{Path, PathBuf};
10
11use serde::{Deserialize, Serialize};
12
13/// Error type for tuning operations.
14#[derive(Debug)]
15pub enum TunerError {
16    /// Failed to read the configuration directory.
17    ReadConfigDir(String),
18    /// Failed to parse a TOML file.
19    ParseTOML(PathBuf, String),
20    /// User config already exists — refusing to overwrite.
21    UserConfigExists(PathBuf),
22    /// Model not found in available configurations.
23    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/// Inference-parameter overrides for a specific model name.
46///
47/// Matched against the active model by exact string equality.  Add entries
48/// under `[[model_tuning]]` in `~/.newt/config.toml` to pin parameters
49/// for models whose defaults cause problems (e.g. context overflow).
50///
51/// Human-authored entries are never touched by the auto-tuner.  Auto-tuned
52/// entries are appended (not modified) when the harness gains high confidence
53/// in its empirical measurements.
54#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
55pub struct ModelTuning {
56    /// Model name as it appears in Ollama (e.g. `"nemotron3:33b"`).
57    pub model: String,
58
59    /// Ollama `options.num_ctx` — hard cap on KV-cache allocation.
60    #[serde(default, skip_serializing_if = "Option::is_none")]
61    pub num_ctx: Option<u32>,
62
63    /// Per-model context window (total tokens the backend accepts).
64    #[serde(default, skip_serializing_if = "Option::is_none")]
65    pub context_window: Option<u32>,
66
67    /// Per-model override of `[tui].real_context_discovery`.
68    #[serde(default, skip_serializing_if = "Option::is_none")]
69    pub real_context_discovery: Option<bool>,
70
71    /// Per-model `mid_loop_trim_threshold` override.
72    #[serde(default, skip_serializing_if = "Option::is_none")]
73    pub mid_loop_trim_threshold: Option<usize>,
74
75    /// Per-model `mid_loop_trim_tokens` override.
76    #[serde(default, skip_serializing_if = "Option::is_none")]
77    pub mid_loop_trim_tokens: Option<usize>,
78
79    /// Per-model `max_tool_rounds` override.
80    #[serde(default, skip_serializing_if = "Option::is_none")]
81    pub max_tool_rounds: Option<usize>,
82
83    /// Per-model `[tui].workflow_grace_rounds` override.
84    #[serde(default, skip_serializing_if = "Option::is_none")]
85    pub workflow_grace_rounds: Option<usize>,
86
87    /// Per-model `[tui].narration_nudge_cap` override.
88    #[serde(default, skip_serializing_if = "Option::is_none")]
89    pub narration_nudge_cap: Option<usize>,
90}
91
92/// Load tuning configuration for a specific model.
93///
94/// Returns the user's config if it exists, otherwise returns the bundled default.
95pub 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    // Check for user config first
101    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    // Fall back to bundled default
108    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
117/// Initialize the model tuning directory with bundled defaults.
118pub 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    // Copy bundled defaults (this would be implemented with embedded resources)
132    Ok(())
133}
134
135/// Get the path to the model tuning configuration directory.
136pub 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}