1use crate::pricing::calculator::ModelPrice;
2use serde::Deserialize;
3use std::collections::HashMap;
4use std::path::Path;
5
6#[derive(Debug, Deserialize, Default)]
7pub struct Config {
8 #[serde(default)]
9 pub subscription: Vec<SubscriptionPeriod>,
10 #[serde(default)]
11 pub pricing_override: HashMap<String, PriceOverride>,
12}
13
14#[derive(Debug, Deserialize)]
15pub struct SubscriptionPeriod {
16 pub start_date: String,
17 pub monthly_price_usd: f64,
18 pub plan: Option<String>,
19}
20
21#[derive(Debug, Deserialize)]
22pub struct PriceOverride {
23 pub base_input: f64,
24 pub cache_write_5m: f64,
25 pub cache_write_1h: f64,
26 pub cache_read: f64,
27 pub output: f64,
28}
29
30impl Config {
31 pub fn load(path: &Path) -> anyhow::Result<Self> {
32 let content = std::fs::read_to_string(path)?;
33 Ok(toml::from_str(&content)?)
34 }
35
36 pub fn to_model_prices(&self) -> HashMap<String, ModelPrice> {
37 self.pricing_override
38 .iter()
39 .map(|(k, v)| {
40 (k.clone(), ModelPrice {
41 base_input: v.base_input,
42 cache_write_5m: v.cache_write_5m,
43 cache_write_1h: v.cache_write_1h,
44 cache_read: v.cache_read,
45 output: v.output,
46 })
47 })
48 .collect()
49 }
50}
51
52#[cfg(test)]
53mod tests {
54 use super::*;
55
56 #[test]
57 fn parse_config() {
58 let toml_str = r#"
59[[subscription]]
60start_date = "2026-01-01"
61monthly_price_usd = 100.0
62plan = "max_5x"
63
64[[subscription]]
65start_date = "2026-03-01"
66monthly_price_usd = 200.0
67plan = "max_20x"
68
69[pricing_override.claude-opus-4-6]
70base_input = 5.0
71cache_write_5m = 6.25
72cache_write_1h = 10.0
73cache_read = 0.50
74output = 25.0
75"#;
76
77 let config: Config = toml::from_str(toml_str).unwrap();
78
79 assert_eq!(config.subscription.len(), 2);
80 assert_eq!(config.subscription[0].start_date, "2026-01-01");
81 assert!((config.subscription[0].monthly_price_usd - 100.0).abs() < f64::EPSILON);
82 assert_eq!(config.subscription[0].plan.as_deref(), Some("max_5x"));
83 assert_eq!(config.subscription[1].start_date, "2026-03-01");
84 assert!((config.subscription[1].monthly_price_usd - 200.0).abs() < f64::EPSILON);
85 assert_eq!(config.subscription[1].plan.as_deref(), Some("max_20x"));
86
87 assert_eq!(config.pricing_override.len(), 1);
88 let opus = &config.pricing_override["claude-opus-4-6"];
89 assert!((opus.base_input - 5.0).abs() < f64::EPSILON);
90 assert!((opus.cache_write_5m - 6.25).abs() < f64::EPSILON);
91 assert!((opus.cache_write_1h - 10.0).abs() < f64::EPSILON);
92 assert!((opus.cache_read - 0.50).abs() < f64::EPSILON);
93 assert!((opus.output - 25.0).abs() < f64::EPSILON);
94
95 let model_prices = config.to_model_prices();
96 assert_eq!(model_prices.len(), 1);
97 let mp = &model_prices["claude-opus-4-6"];
98 assert!((mp.base_input - 5.0).abs() < f64::EPSILON);
99 assert!((mp.output - 25.0).abs() < f64::EPSILON);
100 }
101
102 #[test]
103 fn empty_config() {
104 let config: Config = toml::from_str("").unwrap();
105 assert!(config.subscription.is_empty());
106 assert!(config.pricing_override.is_empty());
107 assert!(config.to_model_prices().is_empty());
108 }
109}