Skip to main content

lean_ctx/core/
litm.rs

1use crate::core::attention_model;
2use crate::core::session::SessionState;
3
4#[derive(Debug, Clone, Copy)]
5pub struct LitmProfile {
6    pub alpha: f64,
7    pub beta: f64,
8    pub gamma: f64,
9    pub name: &'static str,
10}
11
12impl LitmProfile {
13    pub const CLAUDE: Self = Self {
14        alpha: 0.92,
15        beta: 0.50,
16        gamma: 0.88,
17        name: "claude",
18    };
19    pub const GPT: Self = Self {
20        alpha: 0.90,
21        beta: 0.55,
22        gamma: 0.85,
23        name: "gpt",
24    };
25    pub const GEMINI: Self = Self {
26        alpha: 0.88,
27        beta: 0.60,
28        gamma: 0.82,
29        name: "gemini",
30    };
31    pub const DEFAULT: Self = Self::GPT;
32
33    pub fn from_client_name(client: &str) -> Self {
34        if let Ok(override_val) = std::env::var("LEAN_CTX_LITM_PROFILE") {
35            return Self::from_name(&override_val);
36        }
37        let lower = client.to_lowercase();
38        if lower.contains("claude") || lower.contains("cursor") {
39            Self::CLAUDE
40        } else if lower.contains("gemini") {
41            Self::GEMINI
42        } else {
43            Self::GPT
44        }
45    }
46
47    pub fn from_name(name: &str) -> Self {
48        match name.to_lowercase().as_str() {
49            "claude" | "cursor" => Self::CLAUDE,
50            "gemini" => Self::GEMINI,
51            "gpt" | "openai" | "codex" => Self::GPT,
52            _ => Self::DEFAULT,
53        }
54    }
55}
56
57const _ALPHA: f64 = 0.9;
58const _BETA: f64 = 0.55;
59const _GAMMA: f64 = 0.85;
60
61#[allow(dead_code)]
62pub struct PositionedOutput {
63    pub begin_block: String,
64    pub end_block: String,
65}
66
67/// Sorts session state fields by attention priority:
68///   P1 (begin): task, decisions, project topology, file refs
69///   P2 (end): recent findings, test results, next steps
70///   P3 (dropped): old completed tasks, historical reads beyond limit
71pub fn position_optimize(session: &SessionState) -> PositionedOutput {
72    let mut begin_lines = Vec::new();
73    let mut end_lines = Vec::new();
74
75    begin_lines.push(format!(
76        "ACTIVE SESSION v{} | {} calls | {} tok saved",
77        session.version, session.stats.total_tool_calls, session.stats.total_tokens_saved
78    ));
79
80    if let Some(ref task) = session.task {
81        let pct = task
82            .progress_pct
83            .map_or(String::new(), |p| format!(" [{p}%]"));
84        begin_lines.push(format!("Task: {}{pct}", task.description));
85    }
86
87    if let Some(ref root) = session.project_root {
88        begin_lines.push(format!("Root: {root}"));
89    }
90
91    if !session.decisions.is_empty() {
92        let items: Vec<&str> = session
93            .decisions
94            .iter()
95            .rev()
96            .take(5)
97            .map(|d| d.summary.as_str())
98            .collect();
99        begin_lines.push(format!("Decisions: {}", items.join(" | ")));
100    }
101
102    if !session.files_touched.is_empty() {
103        let items: Vec<String> = session
104            .files_touched
105            .iter()
106            .rev()
107            .take(15)
108            .map(|f| {
109                let r = f.file_ref.as_deref().unwrap_or("?");
110                let status = if f.modified { "mod" } else { &f.last_mode };
111                format!("{r}={} [{status}]", short_path(&f.path))
112            })
113            .collect();
114        begin_lines.push(format!("Files: {}", items.join(" ")));
115    }
116
117    if !session.findings.is_empty() {
118        let items: Vec<String> = session
119            .findings
120            .iter()
121            .rev()
122            .take(5)
123            .map(|f| match (&f.file, f.line) {
124                (Some(file), Some(line)) => format!("{}:{line} — {}", short_path(file), f.summary),
125                (Some(file), None) => format!("{} — {}", short_path(file), f.summary),
126                _ => f.summary.clone(),
127            })
128            .collect();
129        end_lines.push(format!("Findings: {}", items.join(" | ")));
130    }
131
132    if let Some(ref tests) = session.test_results {
133        let status = if tests.failed > 0 { "FAIL" } else { "PASS" };
134        end_lines.push(format!(
135            "Tests [{status}]: {}/{} ({})",
136            tests.passed, tests.total, tests.command
137        ));
138    }
139
140    if !session.next_steps.is_empty() {
141        end_lines.push(format!("Next: {}", session.next_steps.join(" → ")));
142    }
143
144    PositionedOutput {
145        begin_block: begin_lines.join("\n"),
146        end_block: end_lines.join("\n"),
147    }
148}
149
150#[allow(dead_code)]
151/// Compute the theoretical LITM efficiency for a given context layout.
152/// Returns (efficiency_without_ccp, efficiency_with_ccp) as percentages.
153pub fn compute_litm_efficiency(
154    begin_tokens: usize,
155    middle_tokens: usize,
156    end_tokens: usize,
157    ccp_begin_tokens: usize,
158    ccp_end_tokens: usize,
159) -> (f64, f64) {
160    let total_without = (begin_tokens + middle_tokens + end_tokens) as f64;
161    let effective_without =
162        _ALPHA * begin_tokens as f64 + _BETA * middle_tokens as f64 + _GAMMA * end_tokens as f64;
163
164    let total_with = (ccp_begin_tokens + ccp_end_tokens) as f64;
165    let effective_with = _ALPHA * ccp_begin_tokens as f64 + _GAMMA * ccp_end_tokens as f64;
166
167    let eff_without = if total_without > 0.0 {
168        effective_without / total_without * 100.0
169    } else {
170        0.0
171    };
172    let eff_with = if total_with > 0.0 {
173        effective_with / total_with * 100.0
174    } else {
175        0.0
176    };
177
178    (eff_without, eff_with)
179}
180
181#[allow(dead_code)]
182/// Profile-aware LITM efficiency using model-specific attention weights.
183pub fn compute_litm_efficiency_for_profile(
184    begin_tokens: usize,
185    middle_tokens: usize,
186    end_tokens: usize,
187    ccp_begin_tokens: usize,
188    ccp_end_tokens: usize,
189    profile: &LitmProfile,
190) -> (f64, f64) {
191    let total_without = (begin_tokens + middle_tokens + end_tokens) as f64;
192    let effective_without = profile.alpha * begin_tokens as f64
193        + profile.beta * middle_tokens as f64
194        + profile.gamma * end_tokens as f64;
195
196    let total_with = (ccp_begin_tokens + ccp_end_tokens) as f64;
197    let effective_with =
198        profile.alpha * ccp_begin_tokens as f64 + profile.gamma * ccp_end_tokens as f64;
199
200    let eff_without = if total_without > 0.0 {
201        effective_without / total_without * 100.0
202    } else {
203        0.0
204    };
205    let eff_with = if total_with > 0.0 {
206        effective_with / total_with * 100.0
207    } else {
208        0.0
209    };
210
211    (eff_without, eff_with)
212}
213
214/// Compute content-aware attention efficiency using the heuristic attention model.
215/// Combines positional U-curve with structural importance for each line.
216pub fn content_attention_efficiency(content: &str, profile: &LitmProfile) -> f64 {
217    let lines: Vec<&str> = content.lines().collect();
218    if lines.is_empty() {
219        return 0.0;
220    }
221
222    let importances: Vec<f64> = lines
223        .iter()
224        .enumerate()
225        .map(|(i, line)| {
226            let pos = i as f64 / (lines.len() - 1).max(1) as f64;
227            attention_model::combined_attention(
228                line,
229                pos,
230                profile.alpha,
231                profile.beta,
232                profile.gamma,
233            )
234        })
235        .collect();
236
237    attention_model::attention_efficiency(&importances, profile.alpha, profile.beta, profile.gamma)
238}
239
240fn short_path(path: &str) -> String {
241    let parts: Vec<&str> = path.split('/').collect();
242    if parts.len() <= 2 {
243        return path.to_string();
244    }
245    parts.last().copied().unwrap_or(path).to_string()
246}
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251
252    #[test]
253    fn litm_efficiency_without_ccp_lower() {
254        let (eff_without, eff_with) = compute_litm_efficiency(100, 500, 100, 300, 200);
255        assert!(
256            eff_with > eff_without,
257            "CCP should improve LITM efficiency: without={eff_without:.1}%, with={eff_with:.1}%"
258        );
259    }
260
261    #[test]
262    fn litm_efficiency_zero_tokens() {
263        let (eff_without, eff_with) = compute_litm_efficiency(0, 0, 0, 0, 0);
264        assert_eq!(eff_without, 0.0);
265        assert_eq!(eff_with, 0.0);
266    }
267
268    #[test]
269    fn litm_all_at_begin_is_alpha() {
270        let (_, eff_with) = compute_litm_efficiency(0, 0, 0, 100, 0);
271        assert!((eff_with - 90.0).abs() < 0.1, "all begin should be ~90%");
272    }
273
274    #[test]
275    fn litm_all_at_end_is_gamma() {
276        let (_, eff_with) = compute_litm_efficiency(0, 0, 0, 0, 100);
277        assert!((eff_with - 85.0).abs() < 0.1, "all end should be ~85%");
278    }
279
280    #[test]
281    fn litm_middle_heavy_is_worst() {
282        let (eff_middle, _) = compute_litm_efficiency(10, 1000, 10, 0, 0);
283        let (eff_balanced, _) = compute_litm_efficiency(500, 20, 500, 0, 0);
284        assert!(
285            eff_balanced > eff_middle,
286            "middle-heavy should be less efficient"
287        );
288    }
289
290    #[test]
291    fn short_path_simple() {
292        assert_eq!(short_path("file.rs"), "file.rs");
293        assert_eq!(short_path("src/file.rs"), "src/file.rs");
294        assert_eq!(short_path("a/b/c/file.rs"), "file.rs");
295    }
296
297    #[test]
298    fn litm_profile_from_client_claude() {
299        let p = LitmProfile::from_client_name("Claude Desktop");
300        assert_eq!(p.name, "claude");
301        assert!((p.alpha - 0.92).abs() < f64::EPSILON);
302    }
303
304    #[test]
305    fn litm_profile_from_client_cursor() {
306        let p = LitmProfile::from_client_name("Cursor");
307        assert_eq!(p.name, "claude");
308    }
309
310    #[test]
311    fn litm_profile_from_client_gemini() {
312        let p = LitmProfile::from_client_name("Gemini CLI");
313        assert_eq!(p.name, "gemini");
314        assert!((p.beta - 0.60).abs() < f64::EPSILON);
315    }
316
317    #[test]
318    fn litm_profile_unknown_defaults_to_gpt() {
319        let p = LitmProfile::from_client_name("unknown-tool");
320        assert_eq!(p.name, "gpt");
321    }
322
323    #[test]
324    fn litm_profile_efficiency_differs_by_model() {
325        let (_, claude_eff) =
326            compute_litm_efficiency_for_profile(200, 0, 100, 200, 100, &LitmProfile::CLAUDE);
327        let (_, gemini_eff) =
328            compute_litm_efficiency_for_profile(200, 0, 100, 200, 100, &LitmProfile::GEMINI);
329        assert!(
330            (claude_eff - gemini_eff).abs() > 0.1,
331            "different profiles should yield different efficiencies"
332        );
333    }
334}