language_barrier_core/
compactor.rs

1use crate::token::TokenCounter;
2
3/// Trait for compacting chat history
4pub trait ChatHistoryCompactor: Send + Sync + Clone {
5    /// Compacts the chat history to fit within a token budget
6    ///
7    /// This method should modify the history in place, removing
8    /// messages as needed, and updating the token counter.
9    fn compact(
10        &self,
11        history: &mut Vec<crate::message::Message>,
12        counter: &mut TokenCounter,
13        max_tokens: usize,
14    );
15}
16
17/// Compactor that drops oldest messages first
18#[derive(Debug, Default, Clone)]
19pub struct DropOldestCompactor {}
20
21impl ChatHistoryCompactor for DropOldestCompactor {
22    fn compact(
23        &self,
24        history: &mut Vec<crate::message::Message>,
25        counter: &mut TokenCounter,
26        max_tokens: usize,
27    ) {
28        // If history is empty or we're already under budget, nothing to do
29        if history.is_empty() || counter.under_budget(max_tokens) {
30            return;
31        }
32
33        // While we're over budget, keep removing oldest messages
34        while !counter.under_budget(max_tokens) && history.len() > 1 {
35            // Remove the oldest message
36            let removed_msg = history.remove(0);
37
38            // Update the token count based on message type
39            match &removed_msg {
40                crate::message::Message::User { content, .. } => {
41                    match content {
42                        crate::message::Content::Text(text) => {
43                            counter.subtract(text);
44                        }
45                        crate::message::Content::Parts(parts) => {
46                            // For multimodal content, count all text parts
47                            for part in parts {
48                                if let crate::message::ContentPart::Text { text } = part {
49                                    counter.subtract(text);
50                                }
51                            }
52                        }
53                    }
54                }
55                crate::message::Message::Assistant { content, .. } => {
56                    if let Some(content_data) = content {
57                        match content_data {
58                            crate::message::Content::Text(text) => {
59                                counter.subtract(text);
60                            }
61                            crate::message::Content::Parts(parts) => {
62                                // For multimodal content, count all text parts
63                                for part in parts {
64                                    if let crate::message::ContentPart::Text { text } = part {
65                                        counter.subtract(text);
66                                    }
67                                }
68                            }
69                        }
70                    }
71                }
72                crate::message::Message::System { content, .. }
73                | crate::message::Message::Tool { content, .. } => {
74                    counter.subtract(content);
75                }
76            }
77        }
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84
85    use crate::message::Message;
86
87    #[test]
88    fn test_drop_oldest_compactor() {
89        let compactor = DropOldestCompactor::default();
90        let mut history = vec![
91            Message::system("System message"),
92            Message::user("First user message"),
93            Message::assistant("First assistant message"),
94        ];
95        let mut counter = TokenCounter::default();
96        counter.observe("System message");
97        counter.observe("First user message");
98        counter.observe("First assistant message");
99
100        // Initial state
101        assert_eq!(history.len(), 3);
102        assert_eq!(counter.total(), 8); // "System" "message" "First" "user" "message" "First" "assistant" "message"
103
104        // Compact to fit within 5 tokens
105        compactor.compact(&mut history, &mut counter, 5);
106
107        // Should have removed at least the system message
108        assert!(history.len() < 3);
109        assert!(counter.total() <= 5);
110    }
111}
language_barrier_core/compactor.rs

language_barrier_core/
compactor.rs
