Skip to main content

ai_agents_memory/
events.rs

1//! Memory hook event types
2
3use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, Serialize, Deserialize)]
6pub struct MemoryCompressEvent {
7    pub messages_compressed: usize,
8    pub summary_tokens_before: u32,
9    pub summary_tokens_after: u32,
10    pub compression_ratio: f64,
11}
12
13impl MemoryCompressEvent {
14    pub fn new(
15        messages_compressed: usize,
16        summary_tokens_before: u32,
17        summary_tokens_after: u32,
18    ) -> Self {
19        let ratio = if summary_tokens_before > 0 {
20            summary_tokens_after as f64 / summary_tokens_before as f64
21        } else {
22            0.0
23        };
24        Self {
25            messages_compressed,
26            summary_tokens_before,
27            summary_tokens_after,
28            compression_ratio: ratio,
29        }
30    }
31}
32
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct MemoryEvictEvent {
35    pub reason: EvictionReason,
36    pub messages_evicted: usize,
37    pub importance_scores: Vec<f64>,
38}
39
40#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
41#[serde(rename_all = "snake_case")]
42pub enum EvictionReason {
43    TokenBudgetExceeded,
44    MessageCountExceeded,
45    StateTransition,
46    Manual,
47}
48
49#[derive(Debug, Clone, Serialize, Deserialize)]
50pub struct MemoryBudgetEvent {
51    pub component: String,
52    pub used_tokens: u32,
53    pub budget_tokens: u32,
54    pub usage_percent: f64,
55}
56
57impl MemoryBudgetEvent {
58    pub fn new(component: impl Into<String>, used_tokens: u32, budget_tokens: u32) -> Self {
59        let usage_percent = if budget_tokens > 0 {
60            (used_tokens as f64 / budget_tokens as f64) * 100.0
61        } else {
62            0.0
63        };
64        Self {
65            component: component.into(),
66            used_tokens,
67            budget_tokens,
68            usage_percent,
69        }
70    }
71}
72
73#[derive(Debug, Clone, Serialize, Deserialize)]
74pub struct FactExtractedEvent {
75    pub category: String,
76    pub content: String,
77    pub confidence: f64,
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    #[test]
85    fn test_memory_compress_event() {
86        let event = MemoryCompressEvent::new(10, 1000, 200);
87        assert_eq!(event.messages_compressed, 10);
88        assert_eq!(event.summary_tokens_before, 1000);
89        assert_eq!(event.summary_tokens_after, 200);
90        assert!((event.compression_ratio - 0.2).abs() < 0.01);
91    }
92
93    #[test]
94    fn test_memory_compress_event_zero_before() {
95        let event = MemoryCompressEvent::new(5, 0, 100);
96        assert_eq!(event.compression_ratio, 0.0);
97    }
98
99    #[test]
100    fn test_memory_budget_event() {
101        let event = MemoryBudgetEvent::new("summary", 800, 1000);
102        assert_eq!(event.component, "summary");
103        assert_eq!(event.used_tokens, 800);
104        assert_eq!(event.budget_tokens, 1000);
105        assert!((event.usage_percent - 80.0).abs() < 0.01);
106    }
107
108    #[test]
109    fn test_memory_budget_event_zero_budget() {
110        let event = MemoryBudgetEvent::new("test", 100, 0);
111        assert_eq!(event.usage_percent, 0.0);
112    }
113
114    #[test]
115    fn test_eviction_reason_serialize() {
116        let reason = EvictionReason::TokenBudgetExceeded;
117        let json = serde_json::to_string(&reason).unwrap();
118        assert_eq!(json, "\"token_budget_exceeded\"");
119
120        let reason = EvictionReason::StateTransition;
121        let json = serde_json::to_string(&reason).unwrap();
122        assert_eq!(json, "\"state_transition\"");
123    }
124
125    #[test]
126    fn test_fact_extracted_event() {
127        let event = FactExtractedEvent {
128            category: "preference".to_string(),
129            content: "User prefers dark mode".to_string(),
130            confidence: 0.95,
131        };
132        assert_eq!(event.category, "preference");
133        assert!(event.confidence > 0.9);
134    }
135}