ai_agent/services/compact/
compact_warning_hook.rs1use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
7pub enum WarningLevel {
8 None,
10 Info,
12 Warning,
14 Critical,
16}
17
18#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct CompactWarningInfo {
21 pub current_tokens: u64,
23 pub max_tokens: u64,
25 pub utilization: f64,
27 pub level: WarningLevel,
29 pub suggestion: Option<String>,
31}
32
33pub fn get_warning_level(utilization: f64) -> WarningLevel {
35 if utilization >= 0.95 {
36 WarningLevel::Critical
37 } else if utilization >= 0.85 {
38 WarningLevel::Warning
39 } else if utilization >= 0.75 {
40 WarningLevel::Info
41 } else {
42 WarningLevel::None
43 }
44}
45
46pub fn create_compact_warning_info(current_tokens: u64, max_tokens: u64) -> CompactWarningInfo {
48 let utilization = if max_tokens > 0 {
49 (current_tokens as f64 / max_tokens as f64) * 100.0
50 } else {
51 0.0
52 };
53
54 let level = get_warning_level(utilization / 100.0);
55
56 let suggestion = match level {
57 WarningLevel::Critical => Some("Run /compact now to avoid losing context".to_string()),
58 WarningLevel::Warning => Some("Consider running /compact soon".to_string()),
59 WarningLevel::Info => Some("You may want to run /compact later".to_string()),
60 WarningLevel::None => None,
61 };
62
63 CompactWarningInfo {
64 current_tokens,
65 max_tokens,
66 utilization,
67 level,
68 suggestion,
69 }
70}
71
72#[cfg(test)]
73mod tests {
74 use super::*;
75
76 #[test]
77 fn test_get_warning_level_critical() {
78 assert_eq!(get_warning_level(0.96), WarningLevel::Critical);
79 }
80
81 #[test]
82 fn test_get_warning_level_warning() {
83 assert_eq!(get_warning_level(0.86), WarningLevel::Warning);
84 }
85
86 #[test]
87 fn test_get_warning_level_info() {
88 assert_eq!(get_warning_level(0.76), WarningLevel::Info);
89 }
90
91 #[test]
92 fn test_get_warning_level_none() {
93 assert_eq!(get_warning_level(0.5), WarningLevel::None);
94 }
95
96 #[test]
97 fn test_create_compact_warning_info() {
98 let info = create_compact_warning_info(80000, 100000);
99 assert_eq!(info.current_tokens, 80000);
100 assert_eq!(info.max_tokens, 100000);
101 assert!((info.utilization - 80.0).abs() < 0.01);
102 }
103}