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 { original_token_count, .. }
127 | CompressionStatus::Summarized { original_token_count, .. } => original_token_count,
128 CompressionStatus::Dropped => 0,
129 }
130 }
131
132 pub fn effective_token_count(&self) -> usize {
134 match self.compression_status {
135 CompressionStatus::Uncompressed => self.estimated_tokens,
136 CompressionStatus::Compressed { summary_token_count, .. }
137 | CompressionStatus::Summarized { summary_token_count, .. } => summary_token_count,
138 CompressionStatus::Dropped => 0,
139 }
140 }
141}
142
143#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
145#[serde(rename_all = "snake_case")]
146pub enum CompressionStatus {
147 Uncompressed,
149 Compressed {
151 original_token_count: usize,
152 summary_token_count: usize,
153 },
154 Summarized {
156 original_token_count: usize,
157 summary_token_count: usize,
158 },
159 Dropped,
162}
163
164#[allow(clippy::derivable_impls)]
165impl Default for CompressionStatus {
166 fn default() -> Self {
167 Self::Uncompressed
168 }
169}
170
171#[cfg(test)]
172mod tests {
173 use super::*;
174
175 #[test]
176 fn test_create_user_metadata() {
177 let meta = MessageMetadata::user_input(1000, 50);
178 assert_eq!(meta.timestamp, 1000);
179 assert!((meta.importance_score - 0.5).abs() < f64::EPSILON);
180 assert_eq!(meta.compression_status, CompressionStatus::Uncompressed);
181 assert_eq!(meta.estimated_tokens, 50);
182 assert_eq!(meta.source.as_deref(), Some("user_input"));
183 }
184
185 #[test]
186 fn test_create_llm_response_metadata() {
187 let meta = MessageMetadata::llm_response(2000, 150);
188 assert!((meta.importance_score - 0.6).abs() < f64::EPSILON);
189 }
190
191 #[test]
192 fn test_mark_compressed() {
193 let mut meta = MessageMetadata::user_input(1000, 200);
194 meta.mark_compressed(200, 50);
195 assert_eq!(meta.estimated_tokens, 50);
196 assert_eq!(meta.effective_token_count(), 50);
197 assert_eq!(meta.original_token_count(), 200);
198 }
199
200 #[test]
201 fn test_mark_summarized() {
202 let mut meta = MessageMetadata::user_input(1000, 300);
203 meta.mark_summarized(300, 30);
204 assert_eq!(meta.effective_token_count(), 30);
205 assert_eq!(meta.original_token_count(), 300);
206 }
207
208 #[test]
209 fn test_set_importance_clamps() {
210 let mut meta = MessageMetadata::user_input(1000, 50);
211 meta.set_importance(1.5);
212 assert!((meta.importance_score - 1.0).abs() < f64::EPSILON);
213 meta.set_importance(-0.5);
214 assert!((meta.importance_score - 0.0).abs() < f64::EPSILON);
215 }
216
217 #[test]
218 fn test_compression_status_serde_roundtrip() {
219 let status =
220 CompressionStatus::Compressed { original_token_count: 200, summary_token_count: 50 };
221 let json = serde_json::to_string(&status).unwrap();
222 let deserialized: CompressionStatus = serde_json::from_str(&json).unwrap();
223 assert_eq!(status, deserialized);
224 }
225}