1use serde::{Deserialize, Serialize};
13
14#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
19pub struct MessageMetadata {
20 pub timestamp: u64,
22
23 pub importance_score: f64,
29
30 pub compression_status: CompressionStatus,
32
33 pub estimated_tokens: usize,
36
37 pub source: Option<String>,
40}
41
42impl MessageMetadata {
43 pub fn user_input(timestamp: u64, estimated_tokens: usize) -> Self {
45 Self {
46 timestamp,
47 importance_score: 0.5,
48 compression_status: CompressionStatus::Uncompressed,
49 estimated_tokens,
50 source: Some("user_input".into()),
51 }
52 }
53
54 pub fn llm_response(timestamp: u64, estimated_tokens: usize) -> Self {
56 Self {
57 timestamp,
58 importance_score: 0.6,
59 compression_status: CompressionStatus::Uncompressed,
60 estimated_tokens,
61 source: Some("llm_response".into()),
62 }
63 }
64
65 pub fn tool_result(timestamp: u64, estimated_tokens: usize) -> Self {
67 Self {
68 timestamp,
69 importance_score: 0.4,
70 compression_status: CompressionStatus::Uncompressed,
71 estimated_tokens,
72 source: Some("tool_result".into()),
73 }
74 }
75
76 pub fn system(timestamp: u64, estimated_tokens: usize) -> Self {
78 Self {
79 timestamp,
80 importance_score: 1.0,
81 compression_status: CompressionStatus::Uncompressed,
82 estimated_tokens,
83 source: Some("system".into()),
84 }
85 }
86
87 pub fn synthetic(timestamp: u64, estimated_tokens: usize) -> Self {
89 Self {
90 timestamp,
91 importance_score: 0.3,
92 compression_status: CompressionStatus::Uncompressed,
93 estimated_tokens,
94 source: Some("synthetic".into()),
95 }
96 }
97
98 pub fn mark_compressed(&mut self, original_tokens: usize, compressed_tokens: usize) {
100 self.compression_status = CompressionStatus::Compressed {
101 original_token_count: original_tokens,
102 summary_token_count: compressed_tokens,
103 };
104 self.estimated_tokens = compressed_tokens;
105 }
106
107 pub fn mark_summarized(&mut self, original_tokens: usize, summary_tokens: usize) {
109 self.compression_status = CompressionStatus::Summarized {
110 original_token_count: original_tokens,
111 summary_token_count: summary_tokens,
112 };
113 self.estimated_tokens = summary_tokens;
114 }
115
116 pub fn set_importance(&mut self, score: f64) {
118 self.importance_score = score.clamp(0.0, 1.0);
119 }
120
121 pub fn original_token_count(&self) -> usize {
124 match self.compression_status {
125 CompressionStatus::Uncompressed => self.estimated_tokens,
126 CompressionStatus::Compressed {
127 original_token_count,
128 ..
129 }
130 | CompressionStatus::Summarized {
131 original_token_count,
132 ..
133 } => original_token_count,
134 CompressionStatus::Dropped => 0,
135 }
136 }
137
138 pub fn effective_token_count(&self) -> usize {
140 match self.compression_status {
141 CompressionStatus::Uncompressed => self.estimated_tokens,
142 CompressionStatus::Compressed {
143 summary_token_count,
144 ..
145 }
146 | CompressionStatus::Summarized {
147 summary_token_count,
148 ..
149 } => summary_token_count,
150 CompressionStatus::Dropped => 0,
151 }
152 }
153}
154
155#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
157#[serde(rename_all = "snake_case")]
158pub enum CompressionStatus {
159 Uncompressed,
161 Compressed {
163 original_token_count: usize,
164 summary_token_count: usize,
165 },
166 Summarized {
168 original_token_count: usize,
169 summary_token_count: usize,
170 },
171 Dropped,
174}
175
176impl Default for CompressionStatus {
177 fn default() -> Self {
178 Self::Uncompressed
179 }
180}
181
182#[cfg(test)]
183mod tests {
184 use super::*;
185
186 #[test]
187 fn test_create_user_metadata() {
188 let meta = MessageMetadata::user_input(1000, 50);
189 assert_eq!(meta.timestamp, 1000);
190 assert!((meta.importance_score - 0.5).abs() < f64::EPSILON);
191 assert_eq!(meta.compression_status, CompressionStatus::Uncompressed);
192 assert_eq!(meta.estimated_tokens, 50);
193 assert_eq!(meta.source.as_deref(), Some("user_input"));
194 }
195
196 #[test]
197 fn test_create_llm_response_metadata() {
198 let meta = MessageMetadata::llm_response(2000, 150);
199 assert!((meta.importance_score - 0.6).abs() < f64::EPSILON);
200 }
201
202 #[test]
203 fn test_mark_compressed() {
204 let mut meta = MessageMetadata::user_input(1000, 200);
205 meta.mark_compressed(200, 50);
206 assert_eq!(meta.estimated_tokens, 50);
207 assert_eq!(meta.effective_token_count(), 50);
208 assert_eq!(meta.original_token_count(), 200);
209 }
210
211 #[test]
212 fn test_mark_summarized() {
213 let mut meta = MessageMetadata::user_input(1000, 300);
214 meta.mark_summarized(300, 30);
215 assert_eq!(meta.effective_token_count(), 30);
216 assert_eq!(meta.original_token_count(), 300);
217 }
218
219 #[test]
220 fn test_set_importance_clamps() {
221 let mut meta = MessageMetadata::user_input(1000, 50);
222 meta.set_importance(1.5);
223 assert!((meta.importance_score - 1.0).abs() < f64::EPSILON);
224 meta.set_importance(-0.5);
225 assert!((meta.importance_score - 0.0).abs() < f64::EPSILON);
226 }
227
228 #[test]
229 fn test_compression_status_serde_roundtrip() {
230 let status = CompressionStatus::Compressed {
231 original_token_count: 200,
232 summary_token_count: 50,
233 };
234 let json = serde_json::to_string(&status).unwrap();
235 let deserialized: CompressionStatus = serde_json::from_str(&json).unwrap();
236 assert_eq!(status, deserialized);
237 }
238}