langchain_rust/memory/
window_buffer.rs

1use std::collections::VecDeque;
2use std::sync::Arc;
3
4use tokio::sync::Mutex;
5
6use crate::schemas::{memory::BaseMemory, messages::Message};
7
8/// A sliding-window conversation memory that retains the last `window_size` messages.
9///
10/// Internally uses a [`VecDeque`] so that evicting the oldest message on overflow
11/// is O(1) rather than the O(n) `Vec::remove(0)` it replaces.
12pub struct WindowBufferMemory {
13    window_size: usize,
14    messages: VecDeque<Message>,
15}
16
17impl Default for WindowBufferMemory {
18    fn default() -> Self {
19        Self::new(10)
20    }
21}
22
23impl WindowBufferMemory {
24    pub fn new(window_size: usize) -> Self {
25        Self {
26            messages: VecDeque::new(),
27            window_size,
28        }
29    }
30}
31
32impl Into<Arc<dyn BaseMemory>> for WindowBufferMemory {
33    fn into(self) -> Arc<dyn BaseMemory> {
34        Arc::new(self)
35    }
36}
37
38impl Into<Arc<Mutex<dyn BaseMemory>>> for WindowBufferMemory {
39    fn into(self) -> Arc<Mutex<dyn BaseMemory>> {
40        Arc::new(Mutex::new(self))
41    }
42}
43
44impl BaseMemory for WindowBufferMemory {
45    fn messages(&self) -> Vec<Message> {
46        self.messages.iter().cloned().collect()
47    }
48    fn add_message(&mut self, message: Message) {
49        if self.window_size == 0 {
50            return;
51        }
52        if self.messages.len() >= self.window_size {
53            self.messages.pop_front();
54        }
55        self.messages.push_back(message);
56    }
57    fn clear(&mut self) {
58        self.messages.clear();
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65    use crate::schemas::MessageType;
66
67    #[test]
68    fn keeps_only_last_window_size_messages() {
69        let mut memory = WindowBufferMemory::new(2);
70        memory.add_user_message(&"one");
71        memory.add_ai_message(&"two");
72        memory.add_user_message(&"three");
73
74        let messages = memory.messages();
75        assert_eq!(messages.len(), 2);
76        assert_eq!(messages[0].content, "two");
77        assert_eq!(messages[0].message_type, MessageType::AIMessage);
78        assert_eq!(messages[1].content, "three");
79        assert_eq!(messages[1].message_type, MessageType::HumanMessage);
80    }
81
82    #[test]
83    fn default_window_is_ten() {
84        let mut memory = WindowBufferMemory::default();
85        for i in 0..15 {
86            memory.add_user_message(&i);
87        }
88        let messages = memory.messages();
89        assert_eq!(messages.len(), 10);
90        assert_eq!(messages.first().unwrap().content, "5");
91        assert_eq!(messages.last().unwrap().content, "14");
92    }
93
94    #[test]
95    fn zero_window_stores_nothing_and_does_not_panic() {
96        let mut memory = WindowBufferMemory::new(0);
97        memory.add_user_message(&"hello");
98        assert!(memory.messages().is_empty());
99    }
100
101    #[test]
102    fn clear_and_to_string() {
103        let mut memory = WindowBufferMemory::new(5);
104        memory.add_user_message(&"hi");
105        memory.add_ai_message(&"hello");
106        assert_eq!(memory.to_string(), "human: hi\nai: hello");
107        memory.clear();
108        assert!(memory.messages().is_empty());
109        assert_eq!(memory.to_string(), "");
110    }
111}
langchain_rust/memory/window_buffer.rs

langchain_rust/memory/
window_buffer.rs
