lean_ctx/core/gain/
mod.rs

1pub mod gain_score;
2pub mod model_pricing;
3pub mod task_classifier;
4
5use serde::{Deserialize, Serialize};
6
7use crate::core::a2a::cost_attribution::CostStore;
8use crate::core::gain::gain_score::GainScore;
9use crate::core::gain::model_pricing::{ModelPricing, ModelQuote};
10use crate::core::gain::task_classifier::{TaskCategory, TaskClassifier};
11use crate::core::heatmap::HeatMap;
12use crate::core::stats::StatsStore;
13
14#[derive(Clone)]
15pub struct GainEngine {
16    pub stats: StatsStore,
17    pub costs: CostStore,
18    pub heatmap: HeatMap,
19    pub pricing: ModelPricing,
20    pub events: Vec<crate::core::events::LeanCtxEvent>,
21    pub session: Option<crate::core::session::SessionState>,
22}
23
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct GainSummary {
26    pub model: ModelQuote,
27    pub total_commands: u64,
28    pub input_tokens: u64,
29    pub output_tokens: u64,
30    pub tokens_saved: u64,
31    pub gain_rate_pct: f64,
32    pub avoided_usd: f64,
33    pub tool_spend_usd: f64,
34    pub roi: Option<f64>,
35    pub score: GainScore,
36}
37
38#[derive(Debug, Clone, Serialize, Deserialize)]
39pub struct TaskGainRow {
40    pub category: TaskCategory,
41    pub commands: u64,
42    pub tokens_saved: u64,
43    pub tool_calls: u64,
44    pub tool_spend_usd: f64,
45}
46
47#[derive(Debug, Clone, Serialize, Deserialize)]
48pub struct FileGainRow {
49    pub path: String,
50    pub access_count: u32,
51    pub tokens_saved: u64,
52    pub compression_pct: f32,
53}
54
55impl GainEngine {
56    pub fn load() -> Self {
57        Self {
58            stats: crate::core::stats::load(),
59            costs: crate::core::a2a::cost_attribution::CostStore::load(),
60            heatmap: crate::core::heatmap::HeatMap::load(),
61            pricing: ModelPricing::load(),
62            events: crate::core::events::load_events_from_file(500),
63            session: crate::core::session::SessionState::load_latest(),
64        }
65    }
66
67    pub fn summary(&self, model: Option<&str>) -> GainSummary {
68        let quote = self.pricing.quote(model);
69        let tokens_saved = self
70            .stats
71            .total_input_tokens
72            .saturating_sub(self.stats.total_output_tokens);
73        let gain_rate_pct = if self.stats.total_input_tokens > 0 {
74            tokens_saved as f64 / self.stats.total_input_tokens as f64 * 100.0
75        } else {
76            0.0
77        };
78        let avoided_usd = quote.cost.estimate_usd(tokens_saved, 0, 0, 0);
79        let tool_spend_usd = self.costs.total_cost().max(0.0);
80        let roi = if tool_spend_usd > 0.0 {
81            Some(avoided_usd / tool_spend_usd)
82        } else {
83            None
84        };
85        let score = GainScore::compute(&self.stats, &self.costs, &self.pricing, model);
86        GainSummary {
87            model: quote,
88            total_commands: self.stats.total_commands,
89            input_tokens: self.stats.total_input_tokens,
90            output_tokens: self.stats.total_output_tokens,
91            tokens_saved,
92            gain_rate_pct,
93            avoided_usd,
94            tool_spend_usd,
95            roi,
96            score,
97        }
98    }
99
100    pub fn gain_score(&self, model: Option<&str>) -> GainScore {
101        GainScore::compute(&self.stats, &self.costs, &self.pricing, model)
102    }
103
104    pub fn task_breakdown(&self) -> Vec<TaskGainRow> {
105        use std::collections::HashMap;
106
107        let mut by_cat: HashMap<TaskCategory, TaskGainRow> = HashMap::new();
108
109        for (cmd_key, st) in &self.stats.commands {
110            let cat = TaskClassifier::classify_command_key(cmd_key);
111            let row = by_cat.entry(cat).or_insert(TaskGainRow {
112                category: cat,
113                commands: 0,
114                tokens_saved: 0,
115                tool_calls: 0,
116                tool_spend_usd: 0.0,
117            });
118            row.commands += st.count;
119            row.tokens_saved += st.input_tokens.saturating_sub(st.output_tokens);
120        }
121
122        for (tool, tc) in &self.costs.tools {
123            let cat = TaskClassifier::classify_tool(tool);
124            let row = by_cat.entry(cat).or_insert(TaskGainRow {
125                category: cat,
126                commands: 0,
127                tokens_saved: 0,
128                tool_calls: 0,
129                tool_spend_usd: 0.0,
130            });
131            row.tool_calls += tc.total_calls;
132            row.tool_spend_usd += tc.cost_usd;
133        }
134
135        let mut out: Vec<TaskGainRow> = by_cat.into_values().collect();
136        out.sort_by_key(|x| std::cmp::Reverse(x.tokens_saved));
137        out
138    }
139
140    pub fn heatmap_gains(&self, limit: usize) -> Vec<FileGainRow> {
141        let mut items: Vec<_> = self.heatmap.entries.values().collect();
142        items.sort_by_key(|x| std::cmp::Reverse(x.total_tokens_saved));
143        items.truncate(limit);
144        items
145            .into_iter()
146            .map(|e| FileGainRow {
147                path: e.path.clone(),
148                access_count: e.access_count,
149                tokens_saved: e.total_tokens_saved,
150                compression_pct: e.avg_compression_ratio * 100.0,
151            })
152            .collect()
153    }
154}
lean_ctx/core/gain/mod.rs

lean_ctx/core/gain/
mod.rs
