langchain_rust/memory/
window_buffer.rs

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

langchain_rust/memory/
window_buffer.rs
