1use std::collections::BTreeMap;
7use std::fmt;
8use std::io;
9use std::path::{Path, PathBuf};
10
11use serde::Deserialize;
12
13#[derive(Debug, Clone, Default, Deserialize)]
15#[serde(default)]
16pub struct Config {
17 pub general: General,
18 pub sources: BTreeMap<String, Source>,
20 pub pricing: BTreeMap<String, BTreeMap<String, ModelPrice>>,
22}
23
24#[derive(Debug, Clone, Deserialize)]
25#[serde(default)]
26pub struct General {
27 pub data_dir: Option<PathBuf>,
29 pub index_prompt_text: bool,
31}
32
33impl Default for General {
34 fn default() -> Self {
35 Self {
36 data_dir: None,
37 index_prompt_text: true,
38 }
39 }
40}
41
42#[derive(Debug, Clone, Deserialize)]
43#[serde(default)]
44pub struct Source {
45 pub enabled: bool,
46 pub path: Option<PathBuf>,
48}
49
50impl Default for Source {
51 fn default() -> Self {
52 Self {
53 enabled: true,
54 path: None,
55 }
56 }
57}
58
59#[derive(Debug, Clone, Copy, Deserialize)]
62pub struct ModelPrice {
63 #[serde(default)]
64 pub input: f64,
65 #[serde(default)]
66 pub output: f64,
67 #[serde(default)]
68 pub cache_read: f64,
69 #[serde(default)]
70 pub cache_write: Option<f64>,
71}
72
73#[derive(Debug, Clone, Copy, Default)]
75pub struct TokenCounts {
76 pub input: Option<u64>,
77 pub output: Option<u64>,
78 pub cache_read: Option<u64>,
79 pub cache_write: Option<u64>,
80}
81
82impl Config {
83 pub fn load_from_dir(dir: &Path) -> Result<Self, ConfigError> {
86 Self::load_file(&dir.join("config.toml"))
87 }
88
89 pub fn load_file(path: &Path) -> Result<Self, ConfigError> {
91 let text = match std::fs::read_to_string(path) {
92 Ok(text) => text,
93 Err(err) if err.kind() == io::ErrorKind::NotFound => return Ok(Self::default()),
94 Err(err) => return Err(ConfigError::Read(path.to_path_buf(), err)),
95 };
96 toml::from_str(&text).map_err(|err| ConfigError::Parse(path.to_path_buf(), err))
97 }
98
99 pub fn source(&self, adapter: &str) -> Source {
101 self.sources.get(adapter).cloned().unwrap_or_default()
102 }
103
104 pub fn price(&self, provider: &str, model: &str) -> Option<ModelPrice> {
106 self.pricing.get(provider)?.get(model).copied()
107 }
108
109 pub fn estimate_cost(&self, provider: &str, model: &str, tokens: TokenCounts) -> Option<f64> {
114 self.pricing().estimate_cost(provider, model, tokens)
115 }
116
117 pub fn pricing(&self) -> Pricing {
123 Pricing {
124 table: self.pricing.clone(),
125 }
126 }
127}
128
129#[derive(Debug, Clone, Default)]
131pub struct Pricing {
132 table: BTreeMap<String, BTreeMap<String, ModelPrice>>,
133}
134
135impl Pricing {
136 pub fn price(&self, provider: &str, model: &str) -> Option<ModelPrice> {
138 self.table.get(provider)?.get(model).copied()
139 }
140
141 pub fn estimate_cost(&self, provider: &str, model: &str, tokens: TokenCounts) -> Option<f64> {
143 let price = self.price(provider, model)?;
144 let per_million = |count: Option<u64>, rate: f64| count.unwrap_or(0) as f64 * rate / 1e6;
145 let cache_write_rate = price.cache_write.unwrap_or(price.input);
148 Some(
149 per_million(tokens.input, price.input)
150 + per_million(tokens.output, price.output)
151 + per_million(tokens.cache_read, price.cache_read)
152 + per_million(tokens.cache_write, cache_write_rate),
153 )
154 }
155}
156
157#[derive(Debug)]
158pub enum ConfigError {
159 Read(PathBuf, io::Error),
160 Parse(PathBuf, toml::de::Error),
161}
162
163impl fmt::Display for ConfigError {
164 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
165 match self {
166 ConfigError::Read(path, err) => write!(f, "reading {}: {err}", path.display()),
167 ConfigError::Parse(path, err) => write!(f, "parsing {}: {err}", path.display()),
168 }
169 }
170}
171
172impl std::error::Error for ConfigError {}
173
174#[cfg(test)]
175mod tests {
176 use super::*;
177
178 const SAMPLE: &str = r#"
179[general]
180index_prompt_text = false
181
182[sources.claude-code]
183enabled = true
184path = "/logs/claude"
185
186[pricing.anthropic]
187"claude-sonnet-4-6" = { input = 3.0, output = 15.0, cache_read = 0.3 }
188"#;
189
190 fn write(dir: &Path, text: &str) {
191 std::fs::write(dir.join("config.toml"), text).unwrap();
192 }
193
194 #[test]
195 fn missing_file_yields_defaults() {
196 let dir = tempfile::tempdir().unwrap();
197 let cfg = Config::load_from_dir(dir.path()).unwrap();
198 assert!(cfg.general.index_prompt_text);
199 assert!(cfg.general.data_dir.is_none());
200 assert!(cfg.pricing.is_empty());
201 let source = cfg.source("claude-code");
203 assert!(source.enabled);
204 assert!(source.path.is_none());
205 }
206
207 #[test]
208 fn parses_all_sections() {
209 let dir = tempfile::tempdir().unwrap();
210 write(dir.path(), SAMPLE);
211 let cfg = Config::load_from_dir(dir.path()).unwrap();
212 assert!(!cfg.general.index_prompt_text);
213 assert_eq!(
214 cfg.source("claude-code").path.as_deref(),
215 Some(Path::new("/logs/claude"))
216 );
217 let price = cfg.price("anthropic", "claude-sonnet-4-6").unwrap();
218 assert_eq!(price.input, 3.0);
219 assert_eq!(price.cache_write, None);
220 }
221
222 #[test]
223 fn estimates_cost_from_config() {
224 let dir = tempfile::tempdir().unwrap();
225 write(dir.path(), SAMPLE);
226 let cfg = Config::load_from_dir(dir.path()).unwrap();
227 let tokens = TokenCounts {
228 input: Some(1_000_000),
229 output: Some(1_000_000),
230 cache_read: Some(1_000_000),
231 cache_write: None,
232 };
233 let cost = cfg
234 .estimate_cost("anthropic", "claude-sonnet-4-6", tokens)
235 .unwrap();
236 assert!((cost - 18.3).abs() < 1e-9, "got {cost}");
237 }
238
239 #[test]
240 fn unpriced_model_returns_none_not_zero() {
241 let dir = tempfile::tempdir().unwrap();
242 write(dir.path(), SAMPLE);
243 let cfg = Config::load_from_dir(dir.path()).unwrap();
244 let tokens = TokenCounts {
245 input: Some(1_000),
246 ..TokenCounts::default()
247 };
248 assert!(cfg
249 .estimate_cost("anthropic", "some-unlisted-model", tokens)
250 .is_none());
251 assert!(cfg.estimate_cost("openai", "gpt-x", tokens).is_none());
252 }
253
254 #[test]
255 fn empty_config_prices_nothing() {
256 let cfg = Config::default();
257 assert!(cfg
258 .estimate_cost("anthropic", "claude-sonnet-4-6", TokenCounts::default())
259 .is_none());
260 }
261}