Skip to main content

sqz_engine/
regret_tracker.rs

1//! Compression Regret Tracker — learns from compression mistakes to improve
2//! future decisions per-file and per-content-type.
3//!
4//! When the LLM re-reads a file that was served from dedup cache, or when
5//! the verifier triggers a safe-mode fallback, that's a "regret event."
6//! The tracker records these and adjusts compression aggressiveness per file.
7//!
8//! NOTE: this in-memory tracker is not wired into the hook path (each hook
9//! invocation is a fresh process, so nothing here persists). The shipped
10//! regret signals live in the session store's `regret_log` table — written
11//! by `CliProxy` (quick reruns) and `CacheManager::expand_prefix` (ref
12//! expands), surfaced by `sqz stats`.
13
14use std::collections::HashMap;
15
16/// A single regret event.
17#[derive(Debug, Clone)]
18pub struct RegretEvent {
19    /// The file path or content identifier.
20    pub content_id: String,
21    /// What happened.
22    pub kind: RegretKind,
23    /// Turn number when the regret occurred.
24    pub turn: u64,
25}
26
27/// Types of compression regret.
28#[derive(Debug, Clone, PartialEq, Eq)]
29pub enum RegretKind {
30    /// LLM re-read a file that was served from dedup cache.
31    DedupReRead,
32    /// Verifier triggered safe-mode fallback (compression was too aggressive).
33    VerifierFallback,
34    /// LLM asked about content that was compressed away.
35    InformationLoss,
36}
37
38/// Per-file compression profile learned from regret events.
39#[derive(Debug, Clone)]
40pub struct FileProfile {
41    /// Number of times this file triggered a regret event.
42    pub regret_count: u32,
43    /// Recommended compression aggressiveness (0.0 = safe, 1.0 = aggressive).
44    pub aggressiveness: f64,
45    /// Last turn this file was accessed.
46    pub last_access_turn: u64,
47}
48
49impl Default for FileProfile {
50    fn default() -> Self {
51        Self {
52            regret_count: 0,
53            aggressiveness: 0.5, // default: balanced
54            last_access_turn: 0,
55        }
56    }
57}
58
59/// Tracks compression regret events and learns per-file compression profiles.
60pub struct RegretTracker {
61    /// Per-file profiles.
62    profiles: HashMap<String, FileProfile>,
63    /// All regret events (for analysis).
64    events: Vec<RegretEvent>,
65    /// How much to reduce aggressiveness per regret event.
66    decay_rate: f64,
67    /// Minimum aggressiveness (never go below this).
68    min_aggressiveness: f64,
69}
70
71impl Default for RegretTracker {
72    fn default() -> Self {
73        Self::new()
74    }
75}
76
77impl RegretTracker {
78    pub fn new() -> Self {
79        Self {
80            profiles: HashMap::new(),
81            events: Vec::new(),
82            decay_rate: 0.15,
83            min_aggressiveness: 0.1,
84        }
85    }
86
87    /// Record a regret event and update the file's compression profile.
88    pub fn record_regret(&mut self, event: RegretEvent) {
89        let profile = self.profiles
90            .entry(event.content_id.clone())
91            .or_default();
92
93        profile.regret_count += 1;
94        profile.last_access_turn = event.turn;
95
96        // Reduce aggressiveness based on regret type
97        let penalty = match event.kind {
98            RegretKind::DedupReRead => self.decay_rate * 0.5,
99            RegretKind::VerifierFallback => self.decay_rate,
100            RegretKind::InformationLoss => self.decay_rate * 2.0,
101        };
102
103        profile.aggressiveness = (profile.aggressiveness - penalty)
104            .max(self.min_aggressiveness);
105
106        self.events.push(event);
107    }
108
109    /// Record a successful compression (no regret) — slowly increase aggressiveness.
110    pub fn record_success(&mut self, content_id: &str, turn: u64) {
111        let profile = self.profiles
112            .entry(content_id.to_string())
113            .or_default();
114
115        profile.last_access_turn = turn;
116
117        // Slowly recover aggressiveness after successful compressions
118        profile.aggressiveness = (profile.aggressiveness + 0.02).min(1.0);
119    }
120
121    /// Get the recommended compression aggressiveness for a file.
122    /// Returns 0.0-1.0 where 0.0 = safe mode, 1.0 = maximum compression.
123    pub fn recommended_aggressiveness(&self, content_id: &str) -> f64 {
124        self.profiles
125            .get(content_id)
126            .map(|p| p.aggressiveness)
127            .unwrap_or(0.5) // default: balanced
128    }
129
130    /// Get the profile for a specific file.
131    pub fn get_profile(&self, content_id: &str) -> Option<&FileProfile> {
132        self.profiles.get(content_id)
133    }
134
135    /// Get total regret event count.
136    pub fn total_regrets(&self) -> usize {
137        self.events.len()
138    }
139
140    /// Get a summary of the most regretted files.
141    pub fn most_regretted(&self, top_n: usize) -> Vec<(&str, &FileProfile)> {
142        let mut sorted: Vec<_> = self.profiles.iter()
143            .map(|(k, v)| (k.as_str(), v))
144            .collect();
145        sorted.sort_by(|a, b| b.1.regret_count.cmp(&a.1.regret_count));
146        sorted.truncate(top_n);
147        sorted
148    }
149
150    /// Format a human-readable regret report.
151    pub fn format_report(&self) -> String {
152        let mut out = format!("sqz regret tracker: {} events\n", self.events.len());
153
154        let top = self.most_regretted(5);
155        if top.is_empty() {
156            out.push_str("  No regret events recorded.\n");
157            return out;
158        }
159
160        out.push_str("  Most regretted files:\n");
161        for (path, profile) in &top {
162            out.push_str(&format!(
163                "    {} — {} regrets, aggressiveness: {:.2}\n",
164                path, profile.regret_count, profile.aggressiveness
165            ));
166        }
167        out
168    }
169
170    /// Reset all profiles (e.g., on project change).
171    pub fn reset(&mut self) {
172        self.profiles.clear();
173        self.events.clear();
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180
181    #[test]
182    fn test_new_tracker_empty() {
183        let tracker = RegretTracker::new();
184        assert_eq!(tracker.total_regrets(), 0);
185        assert_eq!(tracker.recommended_aggressiveness("any.rs"), 0.5);
186    }
187
188    #[test]
189    fn test_record_regret_reduces_aggressiveness() {
190        let mut tracker = RegretTracker::new();
191        let initial = tracker.recommended_aggressiveness("auth.rs");
192
193        tracker.record_regret(RegretEvent {
194            content_id: "auth.rs".into(),
195            kind: RegretKind::VerifierFallback,
196            turn: 1,
197        });
198
199        let after = tracker.recommended_aggressiveness("auth.rs");
200        assert!(after < initial, "aggressiveness should decrease after regret");
201    }
202
203    #[test]
204    fn test_multiple_regrets_compound() {
205        let mut tracker = RegretTracker::new();
206
207        for i in 0..5 {
208            tracker.record_regret(RegretEvent {
209                content_id: "config.yaml".into(),
210                kind: RegretKind::InformationLoss,
211                turn: i,
212            });
213        }
214
215        let agg = tracker.recommended_aggressiveness("config.yaml");
216        assert_eq!(agg, 0.1, "should hit minimum aggressiveness after many regrets");
217    }
218
219    #[test]
220    fn test_success_recovers_aggressiveness() {
221        let mut tracker = RegretTracker::new();
222
223        tracker.record_regret(RegretEvent {
224            content_id: "lib.rs".into(),
225            kind: RegretKind::DedupReRead,
226            turn: 1,
227        });
228
229        let after_regret = tracker.recommended_aggressiveness("lib.rs");
230
231        for i in 2..20 {
232            tracker.record_success("lib.rs", i);
233        }
234
235        let after_recovery = tracker.recommended_aggressiveness("lib.rs");
236        assert!(after_recovery > after_regret, "aggressiveness should recover after successes");
237    }
238
239    #[test]
240    fn test_most_regretted() {
241        let mut tracker = RegretTracker::new();
242
243        for _ in 0..3 {
244            tracker.record_regret(RegretEvent {
245                content_id: "a.rs".into(),
246                kind: RegretKind::DedupReRead,
247                turn: 1,
248            });
249        }
250        tracker.record_regret(RegretEvent {
251            content_id: "b.rs".into(),
252            kind: RegretKind::DedupReRead,
253            turn: 2,
254        });
255
256        let top = tracker.most_regretted(2);
257        assert_eq!(top[0].0, "a.rs");
258        assert_eq!(top[0].1.regret_count, 3);
259    }
260
261    #[test]
262    fn test_format_report() {
263        let mut tracker = RegretTracker::new();
264        tracker.record_regret(RegretEvent {
265            content_id: "test.rs".into(),
266            kind: RegretKind::VerifierFallback,
267            turn: 1,
268        });
269        let report = tracker.format_report();
270        assert!(report.contains("test.rs"));
271        assert!(report.contains("1 regrets"));
272    }
273
274    #[test]
275    fn test_reset_clears_all() {
276        let mut tracker = RegretTracker::new();
277        tracker.record_regret(RegretEvent {
278            content_id: "x.rs".into(),
279            kind: RegretKind::DedupReRead,
280            turn: 1,
281        });
282        tracker.reset();
283        assert_eq!(tracker.total_regrets(), 0);
284        assert_eq!(tracker.recommended_aggressiveness("x.rs"), 0.5);
285    }
286}